diff --git a/concurrent_map.go b/concurrent_map.go index baccab0..44faf1a 100644 --- a/concurrent_map.go +++ b/concurrent_map.go @@ -75,6 +75,30 @@ func (m ConcurrentMap[K, V]) Set(key K, value V) { shard.Unlock() } +// Callback to update an element in the map. +// If the element doesn't exist in the map, the parameter will receive the zero value for the value type. +// The returned value will be stored in the map replacing the existing value. +// Returning false for the second return value aborts the update. +// It is called while lock is held, therefore it MUST NOT +// try to access other keys in same map, as it can lead to deadlock since +// Go sync.RWLock is not reentrant +type UpdateCb[V any] func(exist bool, valueInMap V) (V, bool) + +// Update an existing element using UpdateCb, assuming the key exists. +// If it does not, the zero value for the value type is passed to the callback. +// if the callback return false for the second return value, the map will not be updated. +func (m ConcurrentMap[K, V]) Update(key K, cb UpdateCb[V]) (res V) { + shard := m.GetShard(key) + shard.Lock() + v, ok := shard.items[key] + res, update := cb(ok, v) + if update { + shard.items[key] = res + } + shard.Unlock() + return res +} + // Callback to return new element to be inserted into the map // It is called while lock is held, therefore it MUST NOT // try to access other keys in same map, as it can lead to deadlock since diff --git a/concurrent_map_test.go b/concurrent_map_test.go index 4e5d52e..a62c52b 100644 --- a/concurrent_map_test.go +++ b/concurrent_map_test.go @@ -479,6 +479,68 @@ func TestFnv32(t *testing.T) { } +func TestUpdate(t *testing.T) { + m := New[Animal]() + lion := Animal{"lion"} + + m.Set("safari", lion) + m.Update("safari", func(exists bool, valueInMap Animal) (Animal, bool) { + if !exists { + t.Error("Update recieved false exists flag for existing key") + } + valueInMap.name = "tiger" + return valueInMap, true + }) + safari, ok := m.Get("safari") + if safari.name != "tiger" || !ok { + t.Error("Set, then Update failed") + } + + m.Update("marine", func(exists bool, valueInMap Animal) (Animal, bool) { + if exists { + t.Error("Update recieved exists flag for empty key") + } + if valueInMap.name != "" { + t.Error("Update did not receive zero value for non existing key") + } + valueInMap.name = "whale" + return valueInMap, true + }) + marineAnimals, ok := m.Get("marine") + if marineAnimals.name != "whale" || !ok { + t.Error("Update on non-existing key failed") + } + + // return false to prevent updateing map + m.Set("safari", lion) + m.Update("safari", func(exists bool, valueInMap Animal) (Animal, bool) { + if !exists { + t.Error("Update recieved false exists flag for existing key") + } + valueInMap.name = "tiger" + return valueInMap, false + }) + safari, ok = m.Get("safari") + if safari.name != "lion" || !ok { + t.Error("Set, then aborting Update failed") + } + + m.Update("tundra", func(exists bool, valueInMap Animal) (Animal, bool) { + if exists { + t.Error("Update recieved exists flag for empty key") + } + if valueInMap.name != "" { + t.Error("Update did not receive zero value for non existing key") + } + valueInMap.name = "moose" + return valueInMap, false + }) + _, ok = m.Get("tundra") + if ok { + t.Error("Update aborting on non-existing key failed") + } +} + func TestUpsert(t *testing.T) { dolphin := Animal{"dolphin"} whale := Animal{"whale"}