113 lines
2.6 KiB
Go
113 lines
2.6 KiB
Go
package util
|
|
|
|
import "sync"
|
|
|
|
import (
|
|
"errors"
|
|
"sync/atomic"
|
|
)
|
|
|
|
// ErrorNoSuchLock is returned when the requested lock does not exist
|
|
var ErrorNoSuchLock = errors.New("no such lock")
|
|
|
|
// InMemoryLockRegistry provides a locking mechanism based on the passed in reference name
|
|
type InMemoryLockRegistry struct {
|
|
mu sync.Mutex
|
|
locks map[string]*lockCtr
|
|
}
|
|
|
|
// lockCtr is used by InMemoryLockRegistry to represent a lock with a given name.
|
|
type lockCtr struct {
|
|
mu sync.Mutex
|
|
// waiters is the number of waiters waiting to acquire the lock
|
|
// this is int32 instead of uint32, so we can add `-1` in `dec()`
|
|
waiters int32
|
|
}
|
|
|
|
// inc increments the number of waiters waiting for the lock
|
|
func (l *lockCtr) inc() {
|
|
atomic.AddInt32(&l.waiters, 1)
|
|
}
|
|
|
|
// dec decrements the number of waiters waiting on the lock
|
|
func (l *lockCtr) dec() {
|
|
atomic.AddInt32(&l.waiters, -1)
|
|
}
|
|
|
|
// count gets the current number of waiters
|
|
func (l *lockCtr) count() int32 {
|
|
return atomic.LoadInt32(&l.waiters)
|
|
}
|
|
|
|
// Lock locks the mutex
|
|
func (l *lockCtr) Lock() {
|
|
l.mu.Lock()
|
|
}
|
|
|
|
// Unlock unlocks the mutex
|
|
func (l *lockCtr) Unlock() {
|
|
l.mu.Unlock()
|
|
}
|
|
|
|
// NewInMemoryLockRegistry creates a new InMemoryLockRegistry
|
|
func NewInMemoryLockRegistry() *InMemoryLockRegistry {
|
|
return &InMemoryLockRegistry{
|
|
locks: make(map[string]*lockCtr),
|
|
}
|
|
}
|
|
|
|
// Clear clears all locks by initializing a new map
|
|
func (l *InMemoryLockRegistry) Clear() {
|
|
l.locks = make(map[string]*lockCtr)
|
|
}
|
|
|
|
// Exists exists a lock by name
|
|
func (l *InMemoryLockRegistry) Exists(name string) bool {
|
|
l.mu.Lock()
|
|
_, exists := l.locks[name]
|
|
l.mu.Unlock()
|
|
return exists
|
|
}
|
|
|
|
// Lock locks a mutex with the given name. If it doesn't exist, one is created
|
|
func (l *InMemoryLockRegistry) Lock(name string) {
|
|
l.mu.Lock()
|
|
if l.locks == nil {
|
|
l.locks = make(map[string]*lockCtr)
|
|
}
|
|
|
|
nameLock, exists := l.locks[name]
|
|
if !exists {
|
|
nameLock = &lockCtr{}
|
|
l.locks[name] = nameLock
|
|
}
|
|
|
|
// increment the nameLock waiters while inside the main mutex
|
|
// this makes sure that the lock isn't deleted if `Lock` and `Unlock` are called concurrently
|
|
nameLock.inc()
|
|
l.mu.Unlock()
|
|
|
|
// Lock the nameLock outside the main mutex, so we don't block other operations
|
|
// once locked then we can decrement the number of waiters for this lock
|
|
nameLock.Lock()
|
|
nameLock.dec()
|
|
}
|
|
|
|
// Unlock unlocks the mutex with the given name
|
|
// If the given lock is not being waited on by any other callers, it is deleted
|
|
func (l *InMemoryLockRegistry) Unlock(name string) error {
|
|
l.mu.Lock()
|
|
nameLock, exists := l.locks[name]
|
|
if !exists {
|
|
l.mu.Unlock()
|
|
return ErrorNoSuchLock
|
|
}
|
|
|
|
if nameLock.count() == 0 {
|
|
delete(l.locks, name)
|
|
}
|
|
nameLock.Unlock()
|
|
|
|
l.mu.Unlock()
|
|
return nil
|
|
}
|