package mutexes import ( "sync" "unsafe" ) const ( // possible lock types. lockTypeRead = uint8(1) << 0 lockTypeWrite = uint8(1) << 1 // frequency of GC cycles // per no. unlocks. i.e. // every 'gcfreq' unlocks. gcfreq = 1024 ) // MutexMap is a structure that allows read / write locking // per key, performing as you'd expect a map[string]*RWMutex // to perform, without you needing to worry about deadlocks // between competing read / write locks and the map's own mutex. // It uses memory pooling for the internal "mutex" (ish) types // and performs self-eviction of keys. // // Under the hood this is achieved using a single mutex for the // map, state tracking for individual keys, and some simple waitgroup // type structures to park / block goroutines waiting for keys. type MutexMap struct { mapmu sync.Mutex mumap map[string]*rwmutexish mupool rwmutexPool count uint32 } // checkInit ensures MutexMap is initialized (UNSAFE). func (mm *MutexMap) checkInit() { if mm.mumap == nil { mm.mumap = make(map[string]*rwmutexish) } } // Lock acquires a write lock on key in map, returning unlock function. func (mm *MutexMap) Lock(key string) func() { return mm.lock(key, lockTypeWrite) } // RLock acquires a read lock on key in map, returning runlock function. func (mm *MutexMap) RLock(key string) func() { return mm.lock(key, lockTypeRead) } func (mm *MutexMap) lock(key string, lt uint8) func() { // Perform first map lock // and check initialization // OUTSIDE the main loop. mm.mapmu.Lock() mm.checkInit() for { // Check map for mu. mu := mm.mumap[key] if mu == nil { // Allocate new mutex. mu = mm.mupool.Acquire() mm.mumap[key] = mu } if !mu.Lock(lt) { // Wait on mutex unlock, after // immediately relocking map mu. mu.WaitRelock(&mm.mapmu) continue } // Done with map. mm.mapmu.Unlock() // Return mutex unlock function. return func() { mm.unlock(key, mu) } } } func (mm *MutexMap) unlock(key string, mu *rwmutexish) { // Get map lock. mm.mapmu.Lock() // Unlock mutex. if mu.Unlock() { // Mutex fully unlocked // with zero waiters. Self // evict and release it. delete(mm.mumap, key) mm.mupool.Release(mu) } if mm.count++; mm.count%gcfreq == 0 { // Every 'gcfreq' unlocks perform // a garbage collection to keep // us squeaky clean :] mm.mupool.GC() } // Done with map. mm.mapmu.Unlock() } // rwmutexPool is a very simply memory rwmutexPool. type rwmutexPool struct { current []*rwmutexish victim []*rwmutexish } // Acquire will returns a rwmutexState from rwmutexPool (or alloc new). func (p *rwmutexPool) Acquire() *rwmutexish { // First try the current queue if l := len(p.current) - 1; l >= 0 { mu := p.current[l] p.current = p.current[:l] return mu } // Next try the victim queue. if l := len(p.victim) - 1; l >= 0 { mu := p.victim[l] p.victim = p.victim[:l] return mu } // Lastly, alloc new. mu := new(rwmutexish) return mu } // Release places a sync.rwmutexState back in the rwmutexPool. func (p *rwmutexPool) Release(mu *rwmutexish) { p.current = append(p.current, mu) } // GC will clear out unused entries from the rwmutexPool. func (p *rwmutexPool) GC() { current := p.current p.current = nil p.victim = current } // rwmutexish is a RW mutex (ish), i.e. the representation // of one only to be accessed within type rwmutexish struct { tr trigger ln int32 // no. locks wn int32 // no. waiters lt uint8 // lock type } // Lock will lock the mutex for given lock type, in the // sense that it will update the internal state tracker // accordingly. Return value is true on successful lock. func (mu *rwmutexish) Lock(lt uint8) bool { switch mu.lt { case lockTypeRead: // already read locked, // only permit more reads. if lt != lockTypeRead { return false } case lockTypeWrite: // already write locked, // no other locks allowed. return false default: // Fully unlocked. mu.lt = lt } // Update // count. mu.ln++ return true } // Unlock will unlock the mutex, in the sense that // it will update the internal state tracker accordingly. // On any unlock it will awaken sleeping waiting threads. // Returned boolean is if unlocked=true AND waiters=0. func (mu *rwmutexish) Unlock() bool { var ok bool switch mu.ln--; { case mu.ln > 0 && mu.lt == lockTypeWrite: panic("BUG: multiple writer locks") case mu.ln < 0: panic("BUG: negative lock count") case mu.ln == 0: // Fully unlocked. mu.lt = 0 // Only return true // with no waiters. ok = (mu.wn == 0) } // Awake all waiting // goroutines for mu. mu.tr.Trigger() return ok } // WaitRelock expects a mutex to be passed in already in // the lock state. It incr the rwmutexish waiter count before // unlocking the outer mutex and blocking on internal trigger. // On awake it will relock outer mutex and decr wait count. func (mu *rwmutexish) WaitRelock(outer *sync.Mutex) { mu.wn++ outer.Unlock() mu.tr.Wait() outer.Lock() mu.wn-- } // trigger uses the internals of sync.Cond to provide // a waitgroup type structure (including goroutine parks) // without such a heavy reliance on a delta value. type trigger struct{ notifyList } func (t *trigger) Trigger() { runtime_notifyListNotifyAll(&t.notifyList) } func (t *trigger) Wait() { v := runtime_notifyListAdd(&t.notifyList) runtime_notifyListWait(&t.notifyList, v) } // Approximation of notifyList in runtime/sema.go. type notifyList struct { wait uint32 notify uint32 lock uintptr // key field of the mutex head unsafe.Pointer tail unsafe.Pointer } // See runtime/sema.go for documentation. // //go:linkname runtime_notifyListAdd sync.runtime_notifyListAdd func runtime_notifyListAdd(l *notifyList) uint32 // See runtime/sema.go for documentation. // //go:linkname runtime_notifyListWait sync.runtime_notifyListWait func runtime_notifyListWait(l *notifyList, t uint32) // See runtime/sema.go for documentation. // //go:linkname runtime_notifyListNotifyAll sync.runtime_notifyListNotifyAll func runtime_notifyListNotifyAll(l *notifyList)