diff --git a/bigcache.go b/bigcache.go index 5620c0ef..7c18db39 100644 --- a/bigcache.go +++ b/bigcache.go @@ -153,6 +153,13 @@ func (c *BigCache) Set(key string, entry []byte) error { return shard.set(key, hashedKey, entry) } +// SetOrGet saves entry under the key unless already exist in which case return current value under the key +func (c *BigCache) SetOrGet(key string, entry []byte) (actual []byte, loaded bool, err error) { + hashedKey := c.hash.Sum64(key) + shard := c.getShard(hashedKey) + return shard.setOrGet(key, hashedKey, entry) +} + // Append appends entry under the key if key exists, otherwise // it will set the key (same behaviour as Set()). With Append() you can // concatenate multiple entries under the same key in an lock optimized way. diff --git a/bigcache_bench_test.go b/bigcache_bench_test.go index b6d044e5..d4070da1 100644 --- a/bigcache_bench_test.go +++ b/bigcache_bench_test.go @@ -14,7 +14,9 @@ var message = blob('a', 256) func BenchmarkWriteToCacheWith1Shard(b *testing.B) { writeToCache(b, 1, 100*time.Second, b.N) } - +func BenchmarkWriteToCacheUsingSetOrGetWith1Shard(b *testing.B) { + writeToCacheWithSetOrGet(b, 1, 100*time.Second, b.N) +} func BenchmarkWriteToLimitedCacheWithSmallInitSizeAnd1Shard(b *testing.B) { m := blob('a', 1024) cache, _ := New(context.Background(), Config{ @@ -53,6 +55,13 @@ func BenchmarkWriteToCache(b *testing.B) { }) } } +func BenchmarkWriteToCacheUsingSetOrGet(b *testing.B) { + for _, shards := range []int{1, 512, 1024, 8192} { + b.Run(fmt.Sprintf("%d-shards", shards), func(b *testing.B) { + writeToCacheWithSetOrGet(b, shards, 100*time.Second, b.N) + }) + } +} func BenchmarkAppendToCache(b *testing.B) { for _, shards := range []int{1, 512, 1024, 8192} { b.Run(fmt.Sprintf("%d-shards", shards), func(b *testing.B) { @@ -112,7 +121,9 @@ func BenchmarkIterateOverCache(b *testing.B) { func BenchmarkWriteToCacheWith1024ShardsAndSmallShardInitSize(b *testing.B) { writeToCache(b, 1024, 100*time.Second, 100) } - +func BenchmarkWriteUsingSetOrGetToCacheWith1024ShardsAndSmallShardInitSize(b *testing.B) { + writeToCacheWithSetOrGet(b, 1024, 100*time.Second, 100) +} func BenchmarkReadFromCacheNonExistentKeys(b *testing.B) { for _, shards := range []int{1, 512, 1024, 8192} { b.Run(fmt.Sprintf("%d-shards", shards), func(b *testing.B) { @@ -142,6 +153,25 @@ func writeToCache(b *testing.B, shards int, lifeWindow time.Duration, requestsIn }) } +func writeToCacheWithSetOrGet(b *testing.B, shards int, lifeWindow time.Duration, requestsInLifeWindow int) { + cache, _ := New(context.Background(), Config{ + Shards: shards, + LifeWindow: lifeWindow, + MaxEntriesInWindow: max(requestsInLifeWindow, 100), + MaxEntrySize: 500, + }) + rand.Seed(time.Now().Unix()) + + b.RunParallel(func(pb *testing.PB) { + id := rand.Int() + + b.ReportAllocs() + for pb.Next() { + _, _, _ = cache.SetOrGet(fmt.Sprintf("key-%d", id), message) + } + }) +} + func appendToCache(b *testing.B, shards int, lifeWindow time.Duration, requestsInLifeWindow int) { cache, _ := New(context.Background(), Config{ Shards: shards, diff --git a/bigcache_test.go b/bigcache_test.go index 43ec3f57..287582a1 100644 --- a/bigcache_test.go +++ b/bigcache_test.go @@ -945,6 +945,83 @@ func TestEntryUpdate(t *testing.T) { assertEqual(t, []byte("value2"), cachedValue) } +func TestSetOrGet(t *testing.T) { + t.Parallel() + + // given + clock := mockedClock{value: 0} + cache, _ := newBigCache(context.Background(), Config{ + Shards: 1, + LifeWindow: 6 * time.Second, + MaxEntriesInWindow: 1, + MaxEntrySize: 256, + }, &clock) + + // when + entry1, loaded1, _ := cache.SetOrGet("key1", []byte("value1")) + entry2, loaded2, _ := cache.SetOrGet("key1", []byte("value2")) + entry3, loaded3, _ := cache.SetOrGet("key2", []byte("value3")) + + cachedValue, _ := cache.Get("key1") + + // then + assertEqual(t, []byte("value1"), entry1) + assertEqual(t, []byte("value1"), entry2) + assertEqual(t, []byte("value3"), entry3) + assertEqual(t, []byte("value1"), cachedValue) + assertEqual(t, false, loaded1) + assertEqual(t, true, loaded2) + assertEqual(t, false, loaded3) +} + +func TestSetOrGetCollision(t *testing.T) { + t.Parallel() + + // given + cache, _ := New(context.Background(), Config{ + Shards: 1, + LifeWindow: 5 * time.Second, + MaxEntriesInWindow: 10, + MaxEntrySize: 256, + Verbose: true, + Hasher: hashStub(5), + }) + + //when + entry1, loaded1, _ := cache.SetOrGet("a", []byte("value1")) + entry2, loaded2, _ := cache.SetOrGet("b", []byte("value2")) + + // then + assertEqual(t, []byte("value1"), entry1) + assertEqual(t, []byte("value2"), entry2) + assertEqual(t, false, loaded1) + assertEqual(t, false, loaded2) + assertEqual(t, cache.Stats().Collisions, int64(1)) + +} + +func TestSetOrGetErrorBiggerThanShardSize(t *testing.T) { + t.Parallel() + + // given + cache, _ := New(context.Background(), Config{ + Shards: 1, + LifeWindow: 5 * time.Second, + MaxEntriesInWindow: 1, + MaxEntrySize: 1, + HardMaxCacheSize: 1, + }) + + // when + entry, loaded, err := cache.SetOrGet("key1", blob('a', 1024*1025)) + + // then + assertEqual(t, blob('a', 1024*1025), entry) + assertEqual(t, false, loaded) + assertEqual(t, "entry is bigger than max shard size", err.Error()) + +} + func TestOldestEntryDeletionWhenMaxCacheSizeIsReached(t *testing.T) { t.Parallel() diff --git a/shard.go b/shard.go index 4f03b53e..000a5724 100644 --- a/shard.go +++ b/shard.go @@ -151,6 +151,37 @@ func (s *cacheShard) set(key string, hashedKey uint64, entry []byte) error { } } +func (s *cacheShard) setOrGet(key string, hashedKey uint64, entry []byte) (actual []byte, loaded bool, err error) { + s.lock.Lock() + defer s.lock.Unlock() + + wrappedEntry, err := s.getWrappedEntry(hashedKey) + if err == nil { + if entryKey := readKeyFromEntry(wrappedEntry); key == entryKey { + actual = readEntry(wrappedEntry) + s.hit(hashedKey) + return actual, true, nil + } else { + + s.collision() + if s.isVerbose { + s.logger.Printf("Collision detected. Both %q and %q have the same hash %x", key, entryKey, hashedKey) + } + + delete(s.hashmap, hashedKey) + s.onRemove(wrappedEntry, Deleted) + if s.statsEnabled { + delete(s.hashmapStats, hashedKey) + } + resetHashFromEntry(wrappedEntry) + } + } else if !errors.Is(err, ErrEntryNotFound) { + return entry, false, err + } + + return entry, false, s.addNewWithoutLock(key, hashedKey, entry) +} + func (s *cacheShard) addNewWithoutLock(key string, hashedKey uint64, entry []byte) error { currentTimestamp := uint64(s.clock.Epoch())