diff --git a/pkg/generic/mutexmap.go b/pkg/generic/mutexmap.go new file mode 100644 index 0000000..54002cd --- /dev/null +++ b/pkg/generic/mutexmap.go @@ -0,0 +1,76 @@ +package generic + +import ( + "maps" + "sync" +) + +type MutexMap[K comparable, V any] struct { + mu sync.RWMutex + contents map[K]V +} + +func NewMutexMap[K comparable, V any]() *MutexMap[K, V] { + return &MutexMap[K, V]{contents: make(map[K]V)} +} + +func (m *MutexMap[K, V]) Set(key K, value V) { + m.mu.Lock() + defer m.mu.Unlock() + m.contents[key] = value +} + +func (m *MutexMap[K, V]) Get(key K) (V, bool) { + m.mu.RLock() + defer m.mu.RUnlock() + v, ok := m.contents[key] + return v, ok +} + +func (m *MutexMap[K, V]) Len() int { + m.mu.RLock() + defer m.mu.RUnlock() + return len(m.contents) +} + +// Upd atomically updates the value for key. `fn` is used to determine how to update an existing +// value, if there is one. It will be passed the existing value, and should return `(newValue V, shouldReplace bool)`. +// +// Returns the value that ends up associated with key after the call and whether an update occurred. +func (m *MutexMap[K, V]) Upd(key K, fn func(prev V) (next V, shouldReplace bool)) (V, bool) { + m.mu.Lock() + defer m.mu.Unlock() + prev := m.contents[key] + if next, shouldReplace := fn(prev); shouldReplace { + m.contents[key] = next + return next, true + } + return prev, false +} + +// Del removes key from the map. Returns true if the key was present. +func (m *MutexMap[K, V]) Del(key K) bool { //nolint: boolnames // Return value is incidental, not the main purpose of the function + m.mu.Lock() + defer m.mu.Unlock() + if _, ok := m.contents[key]; !ok { + return false + } + delete(m.contents, key) + return true +} + +func (m *MutexMap[K, V]) Cpy() map[K]V { + m.mu.RLock() + defer m.mu.RUnlock() + cp := make(map[K]V, len(m.contents)) + maps.Copy(cp, m.contents) + return cp +} + +func (m *MutexMap[K, V]) Pop(key K) (ret V, isOk bool) { + m.mu.Lock() + defer m.mu.Unlock() + ret, isOk = m.contents[key] + delete(m.contents, key) + return +}