84 lines
1.3 KiB
Go
84 lines
1.3 KiB
Go
package sync
|
|
|
|
import "sync"
|
|
|
|
type locker struct {
|
|
mtx sync.RWMutex
|
|
userCount int // not protected by mtx, must used outer mutex to update concurrently
|
|
}
|
|
|
|
type KeyLocker[K comparable] struct {
|
|
lockers map[K]*locker
|
|
lockersMtx sync.Mutex
|
|
}
|
|
|
|
func NewKeyLocker[K comparable]() *KeyLocker[K] {
|
|
return &KeyLocker[K]{
|
|
lockers: make(map[K]*locker),
|
|
}
|
|
}
|
|
|
|
func (l *KeyLocker[K]) Lock(key K) {
|
|
l.lock(key, false)
|
|
}
|
|
|
|
func (l *KeyLocker[K]) RLock(key K) {
|
|
l.lock(key, true)
|
|
}
|
|
|
|
func (l *KeyLocker[K]) lock(key K, read bool) {
|
|
l.lockersMtx.Lock()
|
|
|
|
if locker, found := l.lockers[key]; found {
|
|
locker.userCount++
|
|
l.lockersMtx.Unlock()
|
|
|
|
if read {
|
|
locker.mtx.RLock()
|
|
} else {
|
|
locker.mtx.Lock()
|
|
}
|
|
return
|
|
}
|
|
|
|
locker := &locker{
|
|
userCount: 1,
|
|
}
|
|
if read {
|
|
locker.mtx.RLock()
|
|
} else {
|
|
locker.mtx.Lock()
|
|
}
|
|
|
|
l.lockers[key] = locker
|
|
l.lockersMtx.Unlock()
|
|
}
|
|
|
|
func (l *KeyLocker[K]) Unlock(key K) {
|
|
l.unlock(key, false)
|
|
}
|
|
|
|
func (l *KeyLocker[K]) RUnlock(key K) {
|
|
l.unlock(key, true)
|
|
}
|
|
|
|
func (l *KeyLocker[K]) unlock(key K, read bool) {
|
|
l.lockersMtx.Lock()
|
|
defer l.lockersMtx.Unlock()
|
|
|
|
locker, found := l.lockers[key]
|
|
if !found {
|
|
return
|
|
}
|
|
|
|
if locker.userCount == 1 {
|
|
delete(l.lockers, key)
|
|
}
|
|
locker.userCount--
|
|
|
|
if read {
|
|
locker.mtx.RUnlock()
|
|
} else {
|
|
locker.mtx.Unlock()
|
|
}
|
|
}
|