wip

keyedmutex/keyedmutex.go

@@ -0,0 +1,67 @@
+package keyedmutex
+
+import (
+	"container/list"
+	"sync"
+)
+
+type KeyedMutex[K comparable] struct {
+	mu       *sync.Mutex
+	waitList map[K]*list.List
+}
+
+func New[K comparable]() KeyedMutex[K] {
+	return KeyedMutex[K]{
+		mu:       new(sync.Mutex),
+		waitList: map[K]*list.List{},
+	}
+}
+
+func (m KeyedMutex[K]) Lock(key K) {
+	if ch := m.lock(key); ch != nil {
+		<-ch
+	}
+}
+
+func (m KeyedMutex[K]) lock(key K) chan struct{} {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	if waitList, ok := m.waitList[key]; ok {
+		ch := make(chan struct{})
+		waitList.PushBack(ch)
+		return ch
+	}
+
+	m.waitList[key] = list.New()
+	return nil
+}
+
+func (m KeyedMutex[K]) TryLock(key K) bool {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	if _, ok := m.waitList[key]; ok {
+		return false
+	}
+
+	m.waitList[key] = list.New()
+	return true
+}
+
+func (m KeyedMutex[K]) Unlock(key K) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	waitList, ok := m.waitList[key]
+	if !ok {
+		panic("unlock of unlocked mutex")
+	}
+
+	if waitList.Len() == 0 {
+		delete(m.waitList, key)
+	} else {
+		ch := waitList.Remove(waitList.Front()).(chan struct{})
+		ch <- struct{}{}
+	}
+}