Initial commit.

keyedmutex.go

@@ -1,67 +1,69 @@
 package keyedmutex
 
 import (
-	"container/list"
 	"sync"
 )
 
+type keyedLock struct {
+	lock  sync.Mutex
+	count int
+}
+
 type KeyedMutex[K comparable] struct {
-	mu       *sync.Mutex
-	waitList map[K]*list.List
+	mu    sync.Mutex
+	byKey map[K]*keyedLock
 }
 
-func New[K comparable]() KeyedMutex[K] {
-	return KeyedMutex[K]{
-		mu:       new(sync.Mutex),
-		waitList: map[K]*list.List{},
+func New[K comparable]() *KeyedMutex[K] {
+	return &KeyedMutex[K]{
+		byKey: map[K]*keyedLock{},
 	}
 }
 
-func (m KeyedMutex[K]) Lock(key K) {
-	if ch := m.lock(key); ch != nil {
-		<-ch
-	}
-}
-
-func (m KeyedMutex[K]) lock(key K) chan struct{} {
+func (m *KeyedMutex[K]) getLock(key K) *keyedLock {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	if waitList, ok := m.waitList[key]; ok {
-		ch := make(chan struct{})
-		waitList.PushBack(ch)
-		return ch
+	item, ok := m.byKey[key]
+	if !ok {
+		item = &keyedLock{}
+		m.byKey[key] = item
 	}
-
-	m.waitList[key] = list.New()
-	return nil
+	item.count++
+	return item
 }
 
-func (m KeyedMutex[K]) TryLock(key K) bool {
+func (m *KeyedMutex[K]) release(key K, unlock 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]
+	item, ok := m.byKey[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{}{}
+	item.count--
+	if unlock {
+		item.lock.Unlock()
+	}
+
+	if item.count == 0 {
+		delete(m.byKey, key)
 	}
 }
+
+func (m *KeyedMutex[K]) Lock(key K) {
+	m.getLock(key).lock.Lock()
+}
+
+func (m *KeyedMutex[K]) TryLock(key K) bool {
+	if ok := m.getLock(key).lock.TryLock(); !ok {
+		m.release(key, false)
+		return false
+	}
+	return true
+}
+
+func (m *KeyedMutex[K]) Unlock(key K) {
+	m.release(key, true)
+}