Initial commit

keyedmutex.go

@@ -0,0 +1,69 @@
+package keyedmutex
+
+import (
+	"sync"
+)
+
+type keyedLock struct {
+	lock  sync.Mutex
+	count int
+}
+
+type KeyedMutex[K comparable] struct {
+	mu    sync.Mutex
+	byKey map[K]*keyedLock
+}
+
+func New[K comparable]() *KeyedMutex[K] {
+	return &KeyedMutex[K]{
+		byKey: map[K]*keyedLock{},
+	}
+}
+
+func (m *KeyedMutex[K]) getLock(key K) *keyedLock {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	item, ok := m.byKey[key]
+	if !ok {
+		item = &keyedLock{}
+		m.byKey[key] = item
+	}
+	item.count++
+	return item
+}
+
+func (m *KeyedMutex[K]) release(key K, unlock bool) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	item, ok := m.byKey[key]
+	if !ok {
+		panic("unlock of unlocked mutex")
+	}
+
+	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)
+}