toolbox/keyedmutex/keyedmutex.go

68 lines
1.0 KiB
Go

package keyedmutex
import (
"container/list"
"sync"
)
type KeyedMutex struct {
mu *sync.Mutex
waitList map[string]*list.List
}
func New() KeyedMutex {
return KeyedMutex{
mu: new(sync.Mutex),
waitList: map[string]*list.List{},
}
}
func (m KeyedMutex) Lock(key string) {
if ch := m.lock(key); ch != nil {
<-ch
}
}
func (m KeyedMutex) lock(key string) 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) TryLock(key string) 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) Unlock(key string) {
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{}{}
}
}