upda/util/locker_memory.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
}