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{}{} } }