From c66c0d2c9633b3cc24fef2fe2ba0735e75140b31 Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 23 Dec 2025 17:43:52 +0800 Subject: [PATCH 1/3] feat: implement SetIfAbsent --- common_test.go | 38 ++++++++++++++++++++++++++++++++++++++ doc.go | 2 ++ kv.go | 13 +++++++++++++ kv_cache.go | 13 +++++++++++++ kv_cache_test.go | 32 ++++++++++++++++++++++++++++++++ kv_test.go | 28 ++++++++++++++++++++++++++++ locker.go | 10 ++++++++++ locker_test.go | 8 ++++++++ map.go | 13 +++++++++++++ map_ttl.go | 13 +++++++++++++ ring.go | 17 +++++++++++++++++ shard.go | 4 ++++ updater.go | 4 ++++ 13 files changed, 195 insertions(+) diff --git a/common_test.go b/common_test.go index 4355a63..874d053 100644 --- a/common_test.go +++ b/common_test.go @@ -53,6 +53,42 @@ func testSetIfPresentThenGet(t *testing.T, imp Geche[string, string]) { } } +func testSetThenSetIfAbsentThenGet(t *testing.T, imp Geche[string, string]) { + imp.Set("key", "value") + existing, inserted := imp.SetIfAbsent("key", "value2") + if inserted { + t.Errorf("expected SetIfAbsent to not insert a new value for existing key") + } + + if existing != "value" { + t.Errorf("expected existing value %q to be returned from SetIfPresent, got %q", "value", existing) + } + + val, err := imp.Get("key") + if err != nil { + t.Errorf("unexpected error in Get: %v", err) + } + + if val != "value" { + t.Errorf("expected value %q, got %q", "value", val) + } +} + +func testSetIfAbsentThenGet(t *testing.T, imp Geche[string, string]) { + if _, inserted := imp.SetIfAbsent("key", "value"); !inserted { + t.Errorf("expected SetIfAbsent to insert a new value for non-existing key") + } + + val, err := imp.Get("key") + if err != nil { + t.Errorf("expected no error in Get: %v", err) + } + + if val != "value" { + t.Errorf("expected inserted value %q to be found, but got %q", "value", val) + } +} + func testGetNonExist(t *testing.T, imp Geche[string, string]) { _, err := imp.Get("key") if err != ErrNotFound { @@ -236,6 +272,8 @@ func TestCommon(t *testing.T) { {"SnapshotLen", testSnapshotLen}, {"SetSetIfPresentGet", testSetThenSetIfPresentThenGet}, {"SetIfPresentGet", testSetIfPresentThenGet}, + {"SetSetIfAbsentGet", testSetThenSetIfAbsentThenGet}, + {"SetIfAbsentGet", testSetIfAbsentThenGet}, } for _, ci := range caches { for _, tc := range tab { diff --git a/doc.go b/doc.go index a764b8c..f9207f6 100644 --- a/doc.go +++ b/doc.go @@ -11,6 +11,8 @@ type Geche[K comparable, V any] interface { Set(K, V) // SetIfPresent sets the kv only if the key was already present, and returns the previous value (if any) and whether the insertion was performed SetIfPresent(K, V) (V, bool) + // SetIfAbsent sets the kv only if the key didn't exist yet, and returns the existing value (if any) and whether the insertion was performed + SetIfAbsent(K, V) (V, bool) Get(K) (V, error) Del(K) error Snapshot() map[K]V diff --git a/kv.go b/kv.go index 72b8237..edba2ef 100644 --- a/kv.go +++ b/kv.go @@ -126,6 +126,19 @@ func (kv *KV[V]) SetIfPresent(key string, value V) (V, bool) { return previousVal, false } +func (kv *KV[V]) SetIfAbsent(key string, value V) (V, bool) { + kv.mux.Lock() + defer kv.mux.Unlock() + + previousVal, err := kv.data.Get(key) + if err == nil { + return previousVal, false + } + + kv.set(key, value) + return previousVal, true +} + // Set key-value pair while updating the trie. // Panics if key is empty. func (kv *KV[V]) Set(key string, value V) { diff --git a/kv_cache.go b/kv_cache.go index 6e5f26c..1a87955 100644 --- a/kv_cache.go +++ b/kv_cache.go @@ -63,6 +63,19 @@ func (kv *KVCache[K, V]) SetIfPresent(key K, value V) (V, bool) { return kv.zero, false } +func (kv *KVCache[K, V]) SetIfAbsent(key K, value V) (V, bool) { + kv.mux.Lock() + defer kv.mux.Unlock() + + old, found := kv.get(key) + if found { + return old, false + } + + kv.insert(key, value) + return kv.zero, true +} + // Get retrieves a value by key. func (kv *KVCache[K, V]) Get(key K) (V, error) { kv.mux.RLock() diff --git a/kv_cache_test.go b/kv_cache_test.go index 0bee3de..2eee759 100644 --- a/kv_cache_test.go +++ b/kv_cache_test.go @@ -616,6 +616,38 @@ func TestKVCacheSetIfPresent(t *testing.T) { } } +func TestKVCacheSetIfAbsent(t *testing.T) { + cache := NewKVCache[string, string]() + cache.Set("a", "test2") + + old, inserted := cache.SetIfAbsent("a", "test5") + if inserted { + t.Errorf("key \"a\" is present in cache, SetIfAbsent should return false") + } + + if old != "test2" { + t.Errorf("expected %q, got %q", "test2", old) + } + + old, inserted = cache.SetIfAbsent("b", "test6") + if !inserted { + t.Errorf("key \"b\" is absent from cache, SetIfAbsent should return true") + } + + if old != "" { + t.Errorf("expected %q, got %q", "", old) + } + + old, inserted = cache.SetIfAbsent("b", "test7") + if inserted { + t.Errorf("key \"b\" is present in cache, SetIfAbsent should return false") + } + + if old != "test6" { + t.Errorf("previous value should be %q, but got %q", "test6", old) + } +} + func TestKVCacheSetIfPresentConcurrent(t *testing.T) { cache := NewKVCache[string, string]() cache.Set("a", "startA") diff --git a/kv_test.go b/kv_test.go index c22f39a..7195b5a 100644 --- a/kv_test.go +++ b/kv_test.go @@ -336,6 +336,7 @@ type MockErrCache struct{} func (m *MockErrCache) Set(key string, value string) {} func (m *MockErrCache) SetIfPresent(key string, value string) (string, bool) { return "", false } +func (m *MockErrCache) SetIfAbsent(key string, value string) (string, bool) { return "", true } func (m *MockErrCache) Del(key string) error { return nil } func (m *MockErrCache) Snapshot() map[string]string { return nil } func (m *MockErrCache) Len() int { return 0 } @@ -762,6 +763,33 @@ func TestSetIfPresent(t *testing.T) { } } +func TestSetIfAbsent(t *testing.T) { + kv := NewKV[string](NewMapCache[string, string]()) + kv.Set("a", "test2") + + old, inserted := kv.SetIfAbsent("a", "test5") + if inserted { + t.Errorf("key \"a\" is present in kv, SetIfAbsent should return false") + } + + if old != "test2" { + t.Errorf("old value is %q, SetIfAbsent should return true", old) + } + + old, inserted = kv.SetIfAbsent("d", "test3") + if !inserted { + t.Errorf("key \"bbb\" is not present in kv, SetIfAbsent should return true") + } + + if old != "" { + t.Errorf("there was no old value for \"bbb\", SetIfAbsent should have returned zero-value for old value") + } + + if _, inserted := kv.SetIfAbsent("d", "test3"); inserted { + t.Errorf("key \"d\" is present in kv, SetIfPresent should return false") + } +} + func TestSetIfPresentConcurrent(t *testing.T) { kv := NewKV[string](NewMapCache[string, string]()) kv.Set("a", "startA") diff --git a/locker.go b/locker.go index ac24371..4c04bd2 100644 --- a/locker.go +++ b/locker.go @@ -92,6 +92,16 @@ func (tx *Tx[K, V]) SetIfPresent(key K, value V) (V, bool) { return tx.cache.SetIfPresent(key, value) } +func (tx *Tx[K, V]) SetIfAbsent(key K, value V) (V, bool) { + if atomic.LoadInt32(&tx.unlocked) == 1 { + panic("cannot use unlocked transaction") + } + if !tx.writable { + panic("cannot set in read-only transaction") + } + return tx.cache.SetIfAbsent(key, value) +} + // Get value by key from the underlying sharded cache. func (tx *Tx[K, V]) Get(key K) (V, error) { if atomic.LoadInt32(&tx.unlocked) == 1 { diff --git a/locker_test.go b/locker_test.go index 29ef89b..7e54e6b 100644 --- a/locker_test.go +++ b/locker_test.go @@ -97,6 +97,10 @@ func TestLockerRPanics(t *testing.T) { t.Errorf("expected panic (SetIfPresent on RLocked)") } + if !panics(func() { tx.SetIfAbsent(1, 1) }) { + t.Errorf("expected panic (SetIfAbsent on RLocked)") + } + tx.Unlock() if !panics(func() { tx.Unlock() }) { t.Errorf("expected panic (Unlock on already unlocked)") @@ -125,6 +129,10 @@ func TestLockerRPanics(t *testing.T) { if !panics(func() { tx.SetIfPresent(1, 1) }) { t.Errorf("expected panic (SetIfPresent on already unlocked)") } + + if !panics(func() { tx.SetIfAbsent(1, 1) }) { + t.Errorf("expected panic (SetIfAbsent on already unlocked)") + } } func TestLockerRWPanics(t *testing.T) { diff --git a/map.go b/map.go index 5ecab0a..e40b587 100644 --- a/map.go +++ b/map.go @@ -38,6 +38,19 @@ func (c *MapCache[K, V]) SetIfPresent(key K, value V) (V, bool) { return old, false } +func (c *MapCache[K, V]) SetIfAbsent(key K, value V) (V, bool) { + c.mux.Lock() + defer c.mux.Unlock() + + old, ok := c.data[key] + if ok { + return old, false + } + + c.data[key] = value + return old, true +} + // Get returns ErrNotFound if key does not exist in the cache. func (c *MapCache[K, V]) Get(key K) (V, error) { c.mux.RLock() diff --git a/map_ttl.go b/map_ttl.go index 9a61d2b..c81480b 100644 --- a/map_ttl.go +++ b/map_ttl.go @@ -90,6 +90,19 @@ func (c *MapTTLCache[K, V]) SetIfPresent(key K, value V) (V, bool) { return old, false } +func (c *MapTTLCache[K, V]) SetIfAbsent(key K, value V) (V, bool) { + c.mux.Lock() + defer c.mux.Unlock() + + old, err := c.get(key) + if err == nil { + return old, false + } + + c.set(key, value) + return old, true +} + // Get returns ErrNotFound if key is not found in the cache or record is outdated. func (c *MapTTLCache[K, V]) Get(key K) (V, error) { c.mux.RLock() diff --git a/ring.go b/ring.go index e7d0a9e..d35ed1f 100644 --- a/ring.go +++ b/ring.go @@ -44,6 +44,10 @@ func NewRingBuffer[K comparable, V any](size int) *RingBuffer[K, V] { func (c *RingBuffer[K, V]) Set(key K, value V) { c.mux.Lock() defer c.mux.Unlock() + c.set(key, value) +} + +func (c *RingBuffer[K, V]) set(key K, value V) { // Remove the key which value we are overwriting // from the map. GC does not cleanup preallocated map, // so no pressure here. @@ -72,6 +76,19 @@ func (c *RingBuffer[K, V]) SetIfPresent(key K, value V) (V, bool) { return c.zeroV, false } +func (c *RingBuffer[K, V]) SetIfAbsent(key K, value V) (V, bool) { + c.mux.Lock() + defer c.mux.Unlock() + + i, present := c.index[key] + if present { + return c.data[i].V, false + } + + c.set(key, value) + return c.zeroV, true +} + // Get returns cached value for the key, or ErrNotFound if the key does not exist. func (c *RingBuffer[K, V]) Get(key K) (V, error) { c.mux.RLock() diff --git a/shard.go b/shard.go index faed636..d639757 100644 --- a/shard.go +++ b/shard.go @@ -79,6 +79,10 @@ func (s *Sharded[K, V]) SetIfPresent(key K, value V) (V, bool) { return s.shards[s.mapper.Map(key, s.N)].SetIfPresent(key, value) } +func (s *Sharded[K, V]) SetIfAbsent(key K, value V) (V, bool) { + return s.shards[s.mapper.Map(key, s.N)].SetIfAbsent(key, value) +} + // Get value by key from the underlying sharded cache. func (s *Sharded[K, V]) Get(key K) (V, error) { return s.shards[s.mapper.Map(key, s.N)].Get(key) diff --git a/updater.go b/updater.go index abce0dd..01d9235 100644 --- a/updater.go +++ b/updater.go @@ -64,6 +64,10 @@ func (u *Updater[K, V]) SetIfPresent(key K, value V) (V, bool) { return u.cache.SetIfPresent(key, value) } +func (u *Updater[K, V]) SetIfAbsent(key K, value V) (V, bool) { + return u.cache.SetIfAbsent(key, value) +} + // Get returns value from the cache. If the value is not in the cache, // it calls updateFn to get the value and update the cache first. // Since updateFn can return error, Get is not guaranteed to always return the value. From a35fcd378fced21d2f26bb96cb2a88e021064086 Mon Sep 17 00:00:00 2001 From: Sergey Melekhin Date: Tue, 23 Dec 2025 17:44:46 +0700 Subject: [PATCH 2/3] Update kv_test.go --- kv_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kv_test.go b/kv_test.go index 7195b5a..3311810 100644 --- a/kv_test.go +++ b/kv_test.go @@ -778,7 +778,7 @@ func TestSetIfAbsent(t *testing.T) { old, inserted = kv.SetIfAbsent("d", "test3") if !inserted { - t.Errorf("key \"bbb\" is not present in kv, SetIfAbsent should return true") + t.Errorf("key \"d\" is not present in kv, SetIfAbsent should return true") } if old != "" { From 194143cbdf637453c9878a6fc6c39601024e129b Mon Sep 17 00:00:00 2001 From: Sergey Melekhin Date: Tue, 23 Dec 2025 17:44:53 +0700 Subject: [PATCH 3/3] Update kv_test.go --- kv_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kv_test.go b/kv_test.go index 3311810..0127b75 100644 --- a/kv_test.go +++ b/kv_test.go @@ -782,7 +782,7 @@ func TestSetIfAbsent(t *testing.T) { } if old != "" { - t.Errorf("there was no old value for \"bbb\", SetIfAbsent should have returned zero-value for old value") + t.Errorf("there was no old value for \"d\", SetIfAbsent should have returned zero-value for old value") } if _, inserted := kv.SetIfAbsent("d", "test3"); inserted {