diff --git a/concurrent_map.go b/concurrent_map.go index baccab0..16436b3 100644 --- a/concurrent_map.go +++ b/concurrent_map.go @@ -75,6 +75,25 @@ 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. +// 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(valueInMap V) V + +// 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. +func (m ConcurrentMap[K, V]) Update(key K, cb UpdateCb[V]) (res V) { + shard := m.GetShard(key) + shard.Lock() + v := shard.items[key] + shard.items[key] = cb(v) + 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..0de5ca2 100644 --- a/concurrent_map_test.go +++ b/concurrent_map_test.go @@ -479,6 +479,34 @@ func TestFnv32(t *testing.T) { } +func TestUpdate(t *testing.T) { + m := New[Animal]() + lion := Animal{"lion"} + + m.Set("safari", lion) + m.Update("safari", func(valueInMap Animal) Animal { + valueInMap.name = "tiger" + return valueInMap + }) + safari, ok := m.Get("safari") + if safari.name != "tiger" || !ok { + t.Error("Set, then Update failed") + } + + m.Update("marine", func(valueInMap Animal) Animal { + if valueInMap.name != "" { + t.Error("Update did not receive zero value for non existing key") + } + valueInMap.name = "whale" + return valueInMap + }) + marineAnimals, ok := m.Get("marine") + if marineAnimals.name != "whale" || !ok { + t.Error("Update on non-existing key failed") + } + +} + func TestUpsert(t *testing.T) { dolphin := Animal{"dolphin"} whale := Animal{"whale"}