GoToSocial/vendor/codeberg.org/gruf/go-mutexes/map.go
2021-11-13 12:29:08 +01:00

106 lines
2.2 KiB
Go

package mutexes
import (
"sync"
)
// MutexMap is a structure that allows having a map of self-evicting mutexes
// by key. You do not need to worry about managing the contents of the map,
// only requesting RLock/Lock for keys, and ensuring to call the returned
// unlock functions.
type MutexMap struct {
// NOTE:
// Individual keyed mutexes should ONLY ever
// be locked within the protection of the outer
// mapMu lock. If you lock these outside the
// protection of this, there is a chance for
// deadlocks
mus map[string]RWMutex
mapMu sync.Mutex
pool sync.Pool
}
// NewMap returns a new MutexMap instance based on supplied
// RWMutex allocator function, nil implies use default
func NewMap(newFn func() RWMutex) MutexMap {
if newFn == nil {
newFn = NewRW
}
return MutexMap{
mus: make(map[string]RWMutex),
mapMu: sync.Mutex{},
pool: sync.Pool{
New: func() interface{} {
return newFn()
},
},
}
}
func (mm *MutexMap) evict(key string, mu RWMutex) {
// Acquire map lock
mm.mapMu.Lock()
// Toggle mutex lock to
// ensure it is unused
unlock := mu.Lock()
unlock()
// Delete mutex key
delete(mm.mus, key)
mm.mapMu.Unlock()
// Release to pool
mm.pool.Put(mu)
}
// RLock acquires a mutex read lock for supplied key, returning an RUnlock function
func (mm *MutexMap) RLock(key string) func() {
return mm.getLock(key, func(mu RWMutex) func() {
return mu.RLock()
})
}
// Lock acquires a mutex lock for supplied key, returning an Unlock function
func (mm *MutexMap) Lock(key string) func() {
return mm.getLock(key, func(mu RWMutex) func() {
return mu.Lock()
})
}
func (mm *MutexMap) getLock(key string, doLock func(RWMutex) func()) func() {
// Get map lock
mm.mapMu.Lock()
// Look for mutex
mu, ok := mm.mus[key]
if ok {
// Lock and return
// its unlocker func
unlock := doLock(mu)
mm.mapMu.Unlock()
return unlock
}
// Note: even though the mutex data structure is
// small, benchmarking does actually show that pooled
// alloc of mutexes here is faster
// Acquire mu + add
mu = mm.pool.Get().(RWMutex)
mm.mus[key] = mu
// Lock mutex + unlock map
unlockFn := doLock(mu)
mm.mapMu.Unlock()
return func() {
// Unlock mutex
unlockFn()
// Release function
go mm.evict(key, mu)
}
}