70 lines
1.0 KiB
Go
70 lines
1.0 KiB
Go
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)
|
|
}
|