68 lines
1.1 KiB
Go
68 lines
1.1 KiB
Go
|
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{}{}
|
||
|
|
}
|
||
|
|
}
|