chore: add metrics and tests to all cache implementations (#2874)

Author: miparnisari

pkg/cache/cache.go

@@ -68,7 +68,13 @@ type Cache[K KeyString, V any] interface {
 	// Get returns the value for the given key in the cache, if it exists.
 	Get(key K) (V, bool)
 
-	// Set sets a value for the key in the cache, with the given cost.
+	// GetTTL returns the TTL of entries in the cache.
+	// If zero is used, entries are not deleted.
+	GetTTL() time.Duration
+
+	// Set is a best-effort attempt to set a value for the key in the cache, with the given cost.
+	// If GetTTL returns zero, the entry never expires.
+	// Returns true if the value could be set, false if the cost was too high.
 	Set(key K, entry V, cost int64) bool
 
 	// Wait waits for the cache to process and apply updates.
@@ -78,6 +84,7 @@ type Cache[K KeyString, V any] interface {
 	Close()
 
 	// GetMetrics returns the metrics block for the cache.
+	// Some implementations may chose to not return some of these metrics.
 	GetMetrics() Metrics
 
 	zerolog.LogObjectMarshaler
@@ -106,6 +113,7 @@ type noopCache[K KeyString, V any] struct{}
 var _ Cache[StringKey, any] = (*noopCache[StringKey, any])(nil)
 
 func (no *noopCache[K, V]) Get(_ K) (V, bool)          { return *new(V), false }
+func (no *noopCache[K, V]) GetTTL() time.Duration      { return time.Duration(0) }
 func (no *noopCache[K, V]) Set(_ K, _ V, _ int64) bool { return false }
 func (no *noopCache[K, V]) Wait()                      {}
 func (no *noopCache[K, V]) Close()                     {}

pkg/cache/cache_otter.go

@@ -5,6 +5,7 @@ package cache
 import (
 	"math"
 	"sync/atomic"
+	"time"
 
 	"github.com/ccoveille/go-safecast/v2"
 	"github.com/maypok86/otter/v2"
@@ -13,15 +14,20 @@ import (
 )
 
 func NewOtterCacheWithMetrics[K KeyString, V any](name string, config *Config) (Cache[K, V], error) {
-	return NewOtterCache[K, V](config)
+	cache, err := NewOtterCache[K, V](name, config)
+	if err != nil {
+		return nil, err
+	}
+	mustRegisterCache(name, cache)
+	return cache, nil
 }
 
 type valueAndCost[V any] struct {
 	value V
 	cost  uint32
 }
 
-func NewOtterCache[K KeyString, V any](config *Config) (Cache[K, V], error) {
+func NewOtterCache[K KeyString, V any](name string, config *Config) (Cache[K, V], error) {
 	uintCost, err := safecast.Convert[uint64](config.MaxCost)
 	if err != nil {
 		return nil, err
@@ -41,14 +47,22 @@ func NewOtterCache[K KeyString, V any](config *Config) (Cache[K, V], error) {
 
 	cache, err := otter.New(opts)
 	return &otterCache[K, V]{
+		name,
 		cache,
 		otterMetrics{atomic.Uint64{}, counter},
+		config.DefaultTTL,
 	}, err
 }
 
 type otterCache[K KeyString, V any] struct {
+	name    string
 	cache   *otter.Cache[string, valueAndCost[V]]
 	metrics otterMetrics
+	ttl     time.Duration
+}
+
+func (wtc *otterCache[K, V]) GetTTL() time.Duration {
+	return wtc.ttl
 }
 
 func (wtc *otterCache[K, V]) Get(key K) (V, bool) {
@@ -67,13 +81,23 @@ func (wtc *otterCache[K, V]) Set(key K, value V, cost int64) bool {
 		// was too big, so we set to maxint in that case.
 		uintCost = math.MaxUint32
 	}
+
 	wtc.metrics.costAdded.Add(uint64(uintCost))
+	_, ok := wtc.Get(key)
+	if ok {
+		wtc.cache.Invalidate(key.KeyString())
+	}
 	wtc.cache.Set(key.KeyString(), valueAndCost[V]{value, uintCost})
-	return true // Otter doesn't drop insertions for performance
+	if wtc.ttl > 0 {
+		wtc.cache.SetExpiresAfter(key.KeyString(), wtc.ttl)
+	}
+	return true
 }
 
-func (wtc *otterCache[K, V]) Wait()  {}
-func (wtc *otterCache[K, V]) Close() {}
+func (wtc *otterCache[K, V]) Wait() {}
+func (wtc *otterCache[K, V]) Close() {
+	unregisterCache(wtc.name)
+}
 
 type otterMetrics struct {
 	costAdded atomic.Uint64

pkg/cache/cache_ristretto.go

@@ -43,32 +43,32 @@ func NewRistrettoCacheWithMetrics[K KeyString, V any](name string, config *Confi
 		return nil, err
 	}
 
-	cache := wrapped[K, V]{name, config, config.DefaultTTL, rcache}
+	cache := ristretoCache[K, V]{name, config, config.DefaultTTL, rcache}
 	mustRegisterCache(name, cache)
 	return &cache, nil
 }
 
 // NewRistrettoCache creates a new ristretto cache from the given config.
 func NewRistrettoCache[K KeyString, V any](config *Config) (Cache[K, V], error) {
 	rcache, err := ristretto.NewCache(ristrettoConfig(config))
-	return &wrapped[K, V]{"", config, config.DefaultTTL, rcache}, err
+	return &ristretoCache[K, V]{"", config, config.DefaultTTL, rcache}, err
 }
 
-type wrapped[K any, V any] struct {
+type ristretoCache[K any, V any] struct {
 	name       string
 	config     *Config
 	defaultTTL time.Duration
 	ristretto  *ristretto.Cache
 }
 
-func (w wrapped[K, V]) Set(key K, entry V, cost int64) bool {
+func (w ristretoCache[K, V]) Set(key K, entry V, cost int64) bool {
 	if w.defaultTTL <= 0 {
 		return w.ristretto.Set(key, entry, cost)
 	}
 	return w.ristretto.SetWithTTL(key, entry, cost, w.defaultTTL)
 }
 
-func (w wrapped[K, V]) Get(key K) (V, bool) {
+func (w ristretoCache[K, V]) Get(key K) (V, bool) {
 	found, ok := w.ristretto.Get(key)
 	if !ok {
 		return *new(V), false
@@ -77,16 +77,20 @@ func (w wrapped[K, V]) Get(key K) (V, bool) {
 	return found.(V), true
 }
 
-func (w wrapped[K, V]) Wait() {
+func (w ristretoCache[K, V]) Wait() {
 	w.ristretto.Wait()
 }
 
-var _ Cache[StringKey, any] = (*wrapped[StringKey, any])(nil)
+func (w ristretoCache[K, V]) GetTTL() time.Duration {
+	return w.defaultTTL
+}
+
+var _ Cache[StringKey, any] = (*ristretoCache[StringKey, any])(nil)
 
-func (w wrapped[K, V]) GetMetrics() Metrics                   { return w.ristretto.Metrics }
-func (w wrapped[K, V]) MarshalZerologObject(e *zerolog.Event) { e.EmbedObject(w.config) }
+func (w ristretoCache[K, V]) GetMetrics() Metrics                   { return w.ristretto.Metrics }
+func (w ristretoCache[K, V]) MarshalZerologObject(e *zerolog.Event) { e.EmbedObject(w.config) }
 
-func (w wrapped[K, V]) Close() {
+func (w ristretoCache[K, V]) Close() {
 	w.ristretto.Close()
 	unregisterCache(w.name)
 }

pkg/cache/cache_test.go

@@ -0,0 +1,148 @@
+//go:build !wasm
+
+package cache
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestCacheWithMetrics(t *testing.T) {
+	t.Parallel()
+
+	config := &Config{
+		NumCounters: 10000,
+		MaxCost:     1000,
+		DefaultTTL:  10 * time.Hour,
+	}
+
+	t.Run("otter", func(t *testing.T) {
+		t.Parallel()
+		testCacheImplementation(t, func() (Cache[StringKey, string], error) {
+			return NewOtterCacheWithMetrics[StringKey, string]("test-otter", config)
+		})
+	})
+
+	t.Run("ristretto", func(t *testing.T) {
+		t.Parallel()
+		testCacheImplementation(t, func() (Cache[StringKey, string], error) {
+			// Use the metrics version for proper metrics tracking
+			return NewRistrettoCacheWithMetrics[StringKey, string]("test-ristretto", config)
+		})
+	})
+
+	t.Run("theine", func(t *testing.T) {
+		t.Parallel()
+		testCacheImplementation(t, func() (Cache[StringKey, string], error) {
+			return NewTheineCacheWithMetrics[StringKey, string]("test-theine", config)
+		})
+	})
+}
+
+func testCacheImplementation(t *testing.T, factory func() (Cache[StringKey, string], error)) {
+	t.Run("Set and Get", func(t *testing.T) {
+		cache, err := factory()
+		require.NoError(t, err)
+		defer cache.Close()
+
+		// Set multiple entries
+		entries := []struct {
+			key   StringKey
+			value string
+		}{
+			{"key1", "value1"},
+			{"key2", "value2"},
+			{"key3", "value3"},
+		}
+
+		for _, entry := range entries {
+			ok := cache.Set(entry.key, entry.value, 10)
+			require.True(t, ok)
+		}
+
+		// Wait for all sets to be processed
+		cache.Wait()
+
+		// Verify all entries
+		for _, entry := range entries {
+			retrieved, found := cache.Get(entry.key)
+			require.True(t, found, "expected key %s to be found", entry.key)
+			require.Equal(t, entry.value, retrieved, "expected value for key %s to match", entry.key)
+		}
+	})
+
+	t.Run("Set same key with diff values", func(t *testing.T) {
+		cache, err := factory()
+		require.NoError(t, err)
+		defer cache.Close()
+
+		ok := cache.Set(StringKey("metric-key-1"), "value1", 10)
+		require.True(t, ok)
+		cache.Wait()
+		val, found := cache.Get("metric-key-1")
+		require.True(t, found)
+		require.Equal(t, "value1", val)
+
+		// same key set, diff value
+		ok = cache.Set(StringKey("metric-key-1"), "value2", 10)
+		require.True(t, ok)
+		cache.Wait()
+		val, found = cache.Get("metric-key-1")
+		require.True(t, found)
+		require.Equal(t, "value2", val)
+	})
+
+	t.Run("Close multiple times", func(t *testing.T) {
+		cache, err := factory()
+		require.NoError(t, err)
+
+		for i := 0; i < 10; i++ {
+			cache.Close()
+		}
+	})
+
+	t.Run("GetTTL", func(t *testing.T) {
+		cache, err := factory()
+		require.NoError(t, err)
+		defer cache.Close()
+
+		require.Equal(t, 10*time.Hour, cache.GetTTL())
+	})
+
+	t.Run("GetMetrics", func(t *testing.T) {
+		cache, err := factory()
+		require.NoError(t, err)
+		defer cache.Close()
+
+		// Set some values
+		ok := cache.Set(StringKey("metric-key-1"), "value1", 10)
+		require.True(t, ok)
+		ok = cache.Set(StringKey("metric-key-2"), "value2", 20)
+		require.True(t, ok)
+
+		cache.Wait()
+
+		// Perform some gets (hits and misses)
+		_, ok = cache.Get(StringKey("metric-key-1")) // hit
+		require.True(t, ok)
+		_, ok = cache.Get(StringKey("metric-key-2")) // hit
+		require.True(t, ok)
+		_, ok = cache.Get(StringKey("non-existent")) // miss
+		require.False(t, ok)
+
+		metrics := cache.GetMetrics()
+		require.NotNil(t, metrics, "expected metrics to be available")
+
+		// Verify hits and misses are tracked
+		hits := metrics.Hits()
+		misses := metrics.Misses()
+		require.GreaterOrEqual(t, hits, uint64(1), "expected at least one hit")
+		require.GreaterOrEqual(t, misses, uint64(1), "expected at least one miss")
+
+		// Verify cost tracking
+		costAdded := metrics.CostAdded()
+		require.GreaterOrEqual(t, costAdded, uint64(10), "expected cost to be tracked")
+	})
+}