Skip to content

Commit

Permalink
Implement Update method
Browse files Browse the repository at this point in the history
The update method allows atomically mutating a value in the map.
If the key does not exist, the zero value for the value type is mutated
and saved.
  • Loading branch information
tbrisker committed Apr 2, 2023
1 parent 85296bc commit ba8662a
Show file tree
Hide file tree
Showing 2 changed files with 47 additions and 0 deletions.
19 changes: 19 additions & 0 deletions concurrent_map.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
28 changes: 28 additions & 0 deletions concurrent_map_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"}
Expand Down

0 comments on commit ba8662a

Please sign in to comment.