synapsecns/sanguine

package mapmutex

import (
    "fmt"
    "sync"
)

// untypedMapMutex wraps a map of mutexes.  Each key locks separately.
// this implementation is a modified version of https://stackoverflow.com/a/62562831/1011803
// this is unexported on purpose to prevent possible issues with pointers not being equivalent.
// TODO: different implementations need to be moved to ginny: https://flaviocopes.com/golang-generic-generate/
type untypedMapMutex interface {
    Lock(key interface{}) Unlocker
    TryLock(key interface{}) (Unlocker, bool)
    Keys() []interface{}
}

type untypedMapMutexImpl struct {
    ml sync.Mutex              // lock for entry map
    ma map[interface{}]*mentry // entry map
}

type mentry struct {
    // m point back to untypedMapMutexImpl, so we can synchronize removing this mentry when cnt==0
    m *untypedMapMutexImpl
    // el is an entry-specific lock
    el sync.Mutex
    // cnt is the reference count
    cnt int
    // key is the key of the memory entry
    key interface{}
}

// Unlocker provides an Unlock method to release the lock.
type Unlocker interface {
    Unlock()
}

// newMapMutex returns an initialized untypedMapMutexImpl.
func newMapMutex() untypedMapMutex {
    return &untypedMapMutexImpl{ma: make(map[interface{}]*mentry)}
}

// Lock acquires a lock corresponding to this key.
// This method will never return nil and Unlock() must be called
// to release the lock when done.
func (m *untypedMapMutexImpl) Lock(key interface{}) Unlocker {
    // read or create entry for this key atomically
    m.ml.Lock()
    e, ok := m.ma[key]
    if !ok {
        e = &mentry{m: m, key: key}
        m.ma[key] = e
    }
    e.cnt++ // ref count
    m.ml.Unlock()

    // acquire lock, will block here until e.cnt==1
    e.el.Lock()

    return e
}

// TryLock tries to acquire the lock, if this can't be done instantly false is returned. Otherwise
// true and the unlocker are returned.
func (m *untypedMapMutexImpl) TryLock(key interface{}) (Unlocker, bool) {
    // read or create entry for this key atomically
    m.ml.Lock()
    defer m.ml.Unlock()

    e, ok := m.ma[key]
    if !ok {
        e = &mentry{m: m, key: key}
        m.ma[key] = e
    }

    if e.el.TryLock() {
        e.cnt++
        return e, true
    }

    return nil, false
}

// Keys returns all keys in the map.
func (m *untypedMapMutexImpl) Keys() []interface{} {
    m.ml.Lock()
    defer m.ml.Unlock()

    keys := make([]interface{}, 0, len(m.ma))
    for k := range m.ma {
        keys = append(keys, k)
    }
    return keys
}

// Unlock releases the lock for this entry.
func (me *mentry) Unlock() {
    m := me.m

    // decrement and if needed remove entry atomically
    m.ml.Lock()
    e, ok := m.ma[me.key]
    if !ok { // entry must exist
        m.ml.Unlock()
        panic(fmt.Errorf("unlock requested for key=%v but no entry found", me.key))
    }
    e.cnt--        // ref count
    if e.cnt < 1 { // if it hits zero then we own it and remove from map
        delete(m.ma, me.key)
    }
    m.ml.Unlock()

    // now that map stuff is handled, we unlock and let
    // anything else waiting on this key through
    e.el.Unlock()
}