From c98c092a7107f6ec14362f8e38525c09e98d9e52 Mon Sep 17 00:00:00 2001 From: ahrav Date: Fri, 4 Oct 2024 13:25:10 -0700 Subject: [PATCH] [refactor] - Decouple Metrics From Cache Implementation (#3355) * decouple metrics from cache logic * delete * address comments * update --- pkg/cache/decorator.go | 19 ++ pkg/cache/decorator_test.go | 379 ++++++++++++++++++++++++++++++++++++ pkg/cache/lru/lru.go | 72 +++---- pkg/cache/lru/lru_test.go | 152 ++++++++++----- pkg/cache/lru/metrics.go | 41 ---- pkg/cache/metrics.go | 130 +++++++------ 6 files changed, 613 insertions(+), 180 deletions(-) create mode 100644 pkg/cache/decorator_test.go delete mode 100644 pkg/cache/lru/metrics.go diff --git a/pkg/cache/decorator.go b/pkg/cache/decorator.go index 5740dec7e..61dde6b8c 100644 --- a/pkg/cache/decorator.go +++ b/pkg/cache/decorator.go @@ -62,3 +62,22 @@ func (c *WithMetrics[T]) Clear() { c.wrapped.Clear() c.metrics.RecordClear(c.cacheName) } + +// Count returns the number of entries in the cache. It also records a count metric +// for the cache using the provided metrics collector and cache name. +func (c *WithMetrics[T]) Count() int { + count := c.wrapped.Count() + return count +} + +// Keys returns all keys in the cache. It also records a keys metric +// for the cache using the provided metrics collector and cache name. +func (c *WithMetrics[T]) Keys() []string { return c.wrapped.Keys() } + +// Values returns all values in the cache. It also records a values metric +// for the cache using the provided metrics collector and cache name. +func (c *WithMetrics[T]) Values() []T { return c.wrapped.Values() } + +// Contents returns all keys in the cache as a string. It also records a contents metric +// for the cache using the provided metrics collector and cache name. +func (c *WithMetrics[T]) Contents() string { return c.wrapped.Contents() } diff --git a/pkg/cache/decorator_test.go b/pkg/cache/decorator_test.go new file mode 100644 index 000000000..781aa6dd3 --- /dev/null +++ b/pkg/cache/decorator_test.go @@ -0,0 +1,379 @@ +package cache + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" +) + +type mockCollector struct{ mock.Mock } + +func (m *mockCollector) RecordHits(cacheName string, hits uint64) { m.Called(cacheName, hits) } +func (m *mockCollector) RecordMisses(cacheName string, misses uint64) { m.Called(cacheName, misses) } + +func (m *mockCollector) RecordSet(cacheName string) { m.Called(cacheName) } +func (m *mockCollector) RecordHit(cacheName string) { m.Called(cacheName) } +func (m *mockCollector) RecordMiss(cacheName string) { m.Called(cacheName) } +func (m *mockCollector) RecordDelete(cacheName string) { m.Called(cacheName) } +func (m *mockCollector) RecordClear(cacheName string) { m.Called(cacheName) } + +type mockCache[T any] struct{ mock.Mock } + +func (m *mockCache[T]) Set(key string, val T) { m.Called(key, val) } + +func (m *mockCache[T]) Get(key string) (T, bool) { + args := m.Called(key) + var zero T + if args.Get(0) != nil { + return args.Get(0).(T), args.Bool(1) + } + return zero, args.Bool(1) +} + +func (m *mockCache[T]) Exists(key string) bool { + args := m.Called(key) + return args.Bool(0) +} + +func (m *mockCache[T]) Delete(key string) { m.Called(key) } + +func (m *mockCache[T]) Clear() { m.Called() } + +func (m *mockCache[T]) Count() int { + args := m.Called() + return args.Int(0) +} + +func (m *mockCache[T]) Keys() []string { + args := m.Called() + return args.Get(0).([]string) +} + +func (m *mockCache[T]) Values() []T { + args := m.Called() + return args.Get(0).([]T) +} + +func (m *mockCache[T]) Contents() string { + args := m.Called() + return args.String(0) +} + +// setupCache initializes the mock cache and metrics collector, then wraps them with the WithMetrics decorator. +func setupCache[T any](t *testing.T) (*WithMetrics[T], *mockCache[T], *mockCollector) { + t.Helper() + + collector := new(mockCollector) + cache := new(mockCache[T]) + wrappedCache := NewCacheWithMetrics[T](cache, collector, "test_cache") + assert.NotNil(t, wrappedCache, "WithMetrics cache should not be nil") + + return wrappedCache, cache, collector +} + +func TestNewLRUCache(t *testing.T) { + c, _, _ := setupCache[int](t) + assert.Equal(t, "test_cache", c.cacheName) +} + +func TestCacheSet(t *testing.T) { + c, cacheMock, collectorMock := setupCache[string](t) + + collectorMock.On("RecordSet", "test_cache").Once() + cacheMock.On("Set", "key", "value").Once() + + c.Set("key", "value") + + collectorMock.AssertCalled(t, "RecordSet", "test_cache") + cacheMock.AssertCalled(t, "Set", "key", "value") +} + +func TestCacheGet(t *testing.T) { + c, cacheMock, collectorMock := setupCache[string](t) + + collectorMock.On("RecordSet", "test_cache").Once() + cacheMock.On("Set", "key", "value").Once() + + collectorMock.On("RecordHit", "test_cache").Once() + cacheMock.On("Get", "key").Return("value", true).Once() + + collectorMock.On("RecordMiss", "test_cache").Once() + cacheMock.On("Get", "non_existent").Return("", false).Once() + + c.Set("key", "value") + collectorMock.AssertCalled(t, "RecordSet", "test_cache") + cacheMock.AssertCalled(t, "Set", "key", "value") + + value, found := c.Get("key") + assert.True(t, found, "Expected to find the key") + assert.Equal(t, "value", value, "Expected value to match") + collectorMock.AssertCalled(t, "RecordHit", "test_cache") + cacheMock.AssertCalled(t, "Get", "key") + + _, found = c.Get("non_existent") + assert.False(t, found, "Expected not to find the key") + collectorMock.AssertCalled(t, "RecordMiss", "test_cache") + cacheMock.AssertCalled(t, "Get", "non_existent") + + collectorMock.AssertExpectations(t) + cacheMock.AssertExpectations(t) +} + +func TestCacheExists(t *testing.T) { + c, cacheMock, collectorMock := setupCache[string](t) + + collectorMock.On("RecordSet", "test_cache").Once() + cacheMock.On("Set", "key", "value").Once() + + collectorMock.On("RecordHit", "test_cache").Once() + cacheMock.On("Exists", "key").Return(true).Once() + + collectorMock.On("RecordMiss", "test_cache").Once() + cacheMock.On("Exists", "non_existent").Return(false).Once() + + c.Set("key", "value") + collectorMock.AssertCalled(t, "RecordSet", "test_cache") + cacheMock.AssertCalled(t, "Set", "key", "value") + + exists := c.Exists("key") + assert.True(t, exists, "Expected the key to exist") + collectorMock.AssertCalled(t, "RecordHit", "test_cache") + cacheMock.AssertCalled(t, "Exists", "key") + + exists = c.Exists("non_existent") + assert.False(t, exists, "Expected the key not to exist") + collectorMock.AssertCalled(t, "RecordMiss", "test_cache") + cacheMock.AssertCalled(t, "Exists", "non_existent") + + collectorMock.AssertExpectations(t) + cacheMock.AssertExpectations(t) +} + +func TestCacheDelete(t *testing.T) { + c, cacheMock, collectorMock := setupCache[string](t) + + collectorMock.On("RecordSet", "test_cache").Once() + cacheMock.On("Set", "key", "value").Once() + + collectorMock.On("RecordDelete", "test_cache").Once() + cacheMock.On("Delete", "key").Once() + + cacheMock.On("Get", "key").Return("", false).Once() + + collectorMock.On("RecordMiss", "test_cache").Once() + + c.Set("key", "value") + collectorMock.AssertCalled(t, "RecordSet", "test_cache") + cacheMock.AssertCalled(t, "Set", "key", "value") + + c.Delete("key") + collectorMock.AssertCalled(t, "RecordDelete", "test_cache") + cacheMock.AssertCalled(t, "Delete", "key") + + _, found := c.Get("key") + assert.False(t, found, "Expected not to find the deleted key") + collectorMock.AssertCalled(t, "RecordMiss", "test_cache") + cacheMock.AssertCalled(t, "Get", "key") + + collectorMock.AssertExpectations(t) + cacheMock.AssertExpectations(t) +} + +func TestCacheClear(t *testing.T) { + c, cacheMock, collectorMock := setupCache[string](t) + + collectorMock.On("RecordSet", "test_cache").Twice() + cacheMock.On("Set", "key1", "value1").Once() + cacheMock.On("Set", "key2", "value2").Once() + + collectorMock.On("RecordClear", "test_cache").Once() + cacheMock.On("Clear").Once() + + cacheMock.On("Get", "key1").Return("", false).Once() + cacheMock.On("Get", "key2").Return("", false).Once() + + c.Set("key1", "value1") + c.Set("key2", "value2") + collectorMock.AssertNumberOfCalls(t, "RecordSet", 2) + cacheMock.AssertCalled(t, "Set", "key1", "value1") + cacheMock.AssertCalled(t, "Set", "key2", "value2") + + c.Clear() + collectorMock.AssertCalled(t, "RecordClear", "test_cache") + cacheMock.AssertCalled(t, "Clear") + + collectorMock.On("RecordMiss", "test_cache").Twice() + _, found1 := c.Get("key1") + _, found2 := c.Get("key2") + assert.False(t, found1, "Expected not to find key1 after clear") + assert.False(t, found2, "Expected not to find key2 after clear") + collectorMock.AssertNumberOfCalls(t, "RecordMiss", 2) + cacheMock.AssertCalled(t, "Get", "key1") + cacheMock.AssertCalled(t, "Get", "key2") + + collectorMock.AssertExpectations(t) + cacheMock.AssertExpectations(t) +} + +func TestCacheCount(t *testing.T) { + c, cacheMock, collectorMock := setupCache[string](t) + + collectorMock.On("RecordSet", "test_cache").Times(3) + cacheMock.On("Set", mock.Anything, mock.Anything).Times(3) + + cacheMock.On("Count").Return(3).Once() + + collectorMock.On("RecordDelete", "test_cache").Once() + cacheMock.On("Delete", "key2").Once() + cacheMock.On("Count").Return(2).Once() + + collectorMock.On("RecordClear", "test_cache").Once() + cacheMock.On("Clear").Once() + cacheMock.On("Count").Return(0).Once() + + c.Set("key1", "value1") + c.Set("key2", "value2") + c.Set("key3", "value3") + assert.Equal(t, 3, c.Count(), "Expected count to be 3") + collectorMock.AssertNumberOfCalls(t, "RecordSet", 3) + cacheMock.AssertNumberOfCalls(t, "Set", 3) + cacheMock.AssertCalled(t, "Count") + + c.Delete("key2") + assert.Equal(t, 2, c.Count(), "Expected count to be 2 after deletion") + collectorMock.AssertCalled(t, "RecordDelete", "test_cache") + cacheMock.AssertCalled(t, "Delete", "key2") + cacheMock.AssertCalled(t, "Count") + + c.Clear() + assert.Equal(t, 0, c.Count(), "Expected count to be 0 after clear") + collectorMock.AssertCalled(t, "RecordClear", "test_cache") + cacheMock.AssertCalled(t, "Clear") + cacheMock.AssertCalled(t, "Count") + + collectorMock.AssertExpectations(t) + cacheMock.AssertExpectations(t) +} + +func TestCacheKeys(t *testing.T) { + c, cacheMock, collectorMock := setupCache[string](t) + + collectorMock.On("RecordSet", "test_cache").Times(3) + cacheMock.On("Set", mock.Anything, mock.Anything).Times(3) + + collectorMock.On("RecordDelete", "test_cache").Once() + cacheMock.On("Delete", "key2").Once() + cacheMock.On("Clear").Once() + collectorMock.On("RecordClear", "test_cache").Once() + + cacheMock.On("Keys").Return([]string{"key1", "key2", "key3"}).Once() + cacheMock.On("Keys").Return([]string{"key1", "key3"}).Once() + cacheMock.On("Keys").Return([]string{}).Once() + + c.Set("key1", "value1") + c.Set("key2", "value2") + c.Set("key3", "value3") + collectorMock.AssertNumberOfCalls(t, "RecordSet", 3) + cacheMock.AssertNumberOfCalls(t, "Set", 3) + + keys := c.Keys() + assert.Len(t, keys, 3, "Expected 3 keys") + assert.ElementsMatch(t, []string{"key1", "key2", "key3"}, keys, "Keys do not match expected values") + + c.Delete("key2") + keys = c.Keys() + assert.Len(t, keys, 2, "Expected 2 keys after deletion") + assert.ElementsMatch(t, []string{"key1", "key3"}, keys, "Keys do not match expected values after deletion") + collectorMock.AssertCalled(t, "RecordDelete", "test_cache") + + c.Clear() + keys = c.Keys() + assert.Len(t, keys, 0, "Expected no keys after clear") + collectorMock.AssertCalled(t, "RecordClear", "test_cache") + + collectorMock.AssertExpectations(t) + cacheMock.AssertExpectations(t) +} + +func TestCacheValues(t *testing.T) { + c, cacheMock, collectorMock := setupCache[string](t) + + collectorMock.On("RecordSet", "test_cache").Times(3) + cacheMock.On("Set", mock.Anything, mock.Anything).Times(3) + + collectorMock.On("RecordDelete", "test_cache").Once() + cacheMock.On("Delete", "key2").Once() + collectorMock.On("RecordClear", "test_cache").Once() + cacheMock.On("Clear").Once() + + cacheMock.On("Values").Return([]string{"value1", "value2", "value3"}).Once() + cacheMock.On("Values").Return([]string{"value1", "value3"}).Once() + cacheMock.On("Values").Return([]string{}).Once() + + c.Set("key1", "value1") + c.Set("key2", "value2") + c.Set("key3", "value3") + collectorMock.AssertNumberOfCalls(t, "RecordSet", 3) + cacheMock.AssertNumberOfCalls(t, "Set", 3) + + values := c.Values() + assert.Len(t, values, 3, "Expected 3 values") + assert.ElementsMatch(t, []string{"value1", "value2", "value3"}, values, "Values do not match expected values") + + c.Delete("key2") + values = c.Values() + assert.Len(t, values, 2, "Expected 2 values after deletion") + assert.ElementsMatch(t, []string{"value1", "value3"}, values, "Values do not match expected values after deletion") + collectorMock.AssertCalled(t, "RecordDelete", "test_cache") + + c.Clear() + values = c.Values() + assert.Len(t, values, 0, "Expected no values after clear") + collectorMock.AssertCalled(t, "RecordClear", "test_cache") + + collectorMock.AssertExpectations(t) + cacheMock.AssertExpectations(t) +} + +func TestCacheContents(t *testing.T) { + c, cacheMock, collectorMock := setupCache[string](t) + + collectorMock.On("RecordSet", "test_cache").Times(3) + cacheMock.On("Set", mock.Anything, mock.Anything).Times(3) + + collectorMock.On("RecordDelete", "test_cache").Once() + cacheMock.On("Delete", "key2").Once() + collectorMock.On("RecordClear", "test_cache").Once() + cacheMock.On("Clear").Once() + + cacheMock.On("Contents").Return("key1, key2, key3").Once() + cacheMock.On("Contents").Return("key1, key3").Once() + cacheMock.On("Contents").Return("[]").Once() + + c.Set("key1", "value1") + c.Set("key2", "value2") + c.Set("key3", "value3") + collectorMock.AssertNumberOfCalls(t, "RecordSet", 3) + cacheMock.AssertNumberOfCalls(t, "Set", 3) + + contents := c.Contents() + assert.Contains(t, contents, "key1", "Contents should contain key1") + assert.Contains(t, contents, "key2", "Contents should contain key2") + assert.Contains(t, contents, "key3", "Contents should contain key3") + + c.Delete("key2") + contents = c.Contents() + assert.Contains(t, contents, "key1", "Contents should contain key1") + assert.NotContains(t, contents, "key2", "Contents should not contain key2") + assert.Contains(t, contents, "key3", "Contents should contain key3") + collectorMock.AssertCalled(t, "RecordDelete", "test_cache") + + c.Clear() + contents = c.Contents() + assert.Equal(t, "[]", contents, "Contents should be empty after clear") + collectorMock.AssertCalled(t, "RecordClear", "test_cache") + + collectorMock.AssertExpectations(t) + cacheMock.AssertExpectations(t) +} diff --git a/pkg/cache/lru/lru.go b/pkg/cache/lru/lru.go index f8e117a8a..e82d2999a 100644 --- a/pkg/cache/lru/lru.go +++ b/pkg/cache/lru/lru.go @@ -12,42 +12,32 @@ import ( lru "github.com/hashicorp/golang-lru/v2" "github.com/trufflesecurity/trufflehog/v3/pkg/cache" - "github.com/trufflesecurity/trufflehog/v3/pkg/common" ) -// collector is an interface that extends cache.BaseMetricsCollector -// and adds methods for recording cache evictions. -type collector interface { - cache.BaseMetricsCollector - - RecordEviction(cacheName string) -} - // Cache is a generic LRU-sized cache that stores key-value pairs with a maximum size limit. // It wraps the lru.Cache library and adds support for custom metrics collection. type Cache[T any] struct { cache *lru.Cache[string, T] - cacheName string - capacity int - metrics collector + cacheName string + capacity int + evictMetrics cache.EvictionMetricsCollector } // Option defines a functional option for configuring the Cache. type Option[T any] func(*Cache[T]) -// WithMetricsCollector is a functional option to set a custom metrics collector. -// It sets the metrics field of the Cache. -func WithMetricsCollector[T any](collector collector) Option[T] { - return func(lc *Cache[T]) { lc.metrics = collector } -} - // WithCapacity is a functional option to set the maximum number of items the cache can hold. // If the capacity is not set, the default value (128_000) is used. func WithCapacity[T any](capacity int) Option[T] { return func(lc *Cache[T]) { lc.capacity = capacity } } +// WithMetricsCollector is a functional option to set a custom metrics collector. +func WithMetricsCollector[T any](collector cache.EvictionMetricsCollector) Option[T] { + return func(lc *Cache[T]) { lc.evictMetrics = collector } +} + // NewCache creates a new Cache with optional configuration parameters. // It takes a cache name and a variadic list of options. func NewCache[T any](cacheName string, opts ...Option[T]) (*Cache[T], error) { @@ -55,7 +45,6 @@ func NewCache[T any](cacheName string, opts ...Option[T]) (*Cache[T], error) { const defaultSize = 128_000 sizedLRU := &Cache[T]{ - metrics: NewSizedLRUMetricsCollector(common.MetricsNamespace, common.MetricsSubsystem), cacheName: cacheName, } @@ -63,9 +52,12 @@ func NewCache[T any](cacheName string, opts ...Option[T]) (*Cache[T], error) { opt(sizedLRU) } - // Provide a evict callback function to record evictions. - onEvicted := func(string, T) { - sizedLRU.metrics.RecordEviction(sizedLRU.cacheName) + var onEvicted func(string, T) + // Provide a evict callback function to record evictions if a custom metrics collector is provided. + if sizedLRU.evictMetrics != nil { + onEvicted = func(string, T) { + sizedLRU.evictMetrics.RecordEviction(sizedLRU.cacheName) + } } lcache, err := lru.NewWithEvict[string, T](defaultSize, onEvicted) @@ -79,19 +71,14 @@ func NewCache[T any](cacheName string, opts ...Option[T]) (*Cache[T], error) { } // Set adds a key-value pair to the cache. -func (lc *Cache[T]) Set(key string, val T) { - lc.cache.Add(key, val) - lc.metrics.RecordSet(lc.cacheName) -} +func (lc *Cache[T]) Set(key string, val T) { lc.cache.Add(key, val) } // Get retrieves a value from the cache by key. func (lc *Cache[T]) Get(key string) (T, bool) { value, found := lc.cache.Get(key) if found { - lc.metrics.RecordHit(lc.cacheName) return value, true } - lc.metrics.RecordMiss(lc.cacheName) var zero T return zero, false } @@ -99,22 +86,37 @@ func (lc *Cache[T]) Get(key string) (T, bool) { // Exists checks if a key exists in the cache. func (lc *Cache[T]) Exists(key string) bool { _, found := lc.cache.Get(key) - if found { - lc.metrics.RecordHit(lc.cacheName) - } else { - lc.metrics.RecordMiss(lc.cacheName) - } return found } // Delete removes a key from the cache. func (lc *Cache[T]) Delete(key string) { lc.cache.Remove(key) - lc.metrics.RecordDelete(lc.cacheName) } // Clear removes all keys from the cache. func (lc *Cache[T]) Clear() { lc.cache.Purge() - lc.metrics.RecordClear(lc.cacheName) +} + +// Count returns the number of key-value pairs in the cache. +func (lc *Cache[T]) Count() int { return lc.cache.Len() } + +// Keys returns all keys in the cache. +func (lc *Cache[T]) Keys() []string { return lc.cache.Keys() } + +// Values returns all values in the cache. +func (lc *Cache[T]) Values() []T { + items := lc.cache.Keys() + res := make([]T, 0, len(items)) + for _, k := range items { + v, _ := lc.cache.Get(k) + res = append(res, v) + } + return res +} + +// Contents returns all keys in the cache encoded as a string. +func (lc *Cache[T]) Contents() string { + return fmt.Sprintf("%v", lc.cache.Keys()) } diff --git a/pkg/cache/lru/lru_test.go b/pkg/cache/lru/lru_test.go index 1bb08e04f..7c3c04ba5 100644 --- a/pkg/cache/lru/lru_test.go +++ b/pkg/cache/lru/lru_test.go @@ -9,22 +9,8 @@ import ( type mockCollector struct{ mock.Mock } -func (m *mockCollector) RecordHits(cacheName string, hits uint64) { m.Called(cacheName, hits) } - -func (m *mockCollector) RecordMisses(cacheName string, misses uint64) { m.Called(cacheName, misses) } - func (m *mockCollector) RecordEviction(cacheName string) { m.Called(cacheName) } -func (m *mockCollector) RecordSet(cacheName string) { m.Called(cacheName) } - -func (m *mockCollector) RecordHit(cacheName string) { m.Called(cacheName) } - -func (m *mockCollector) RecordMiss(cacheName string) { m.Called(cacheName) } - -func (m *mockCollector) RecordDelete(cacheName string) { m.Called(cacheName) } - -func (m *mockCollector) RecordClear(cacheName string) { m.Called(cacheName) } - // setupCache initializes the metrics and cache. // If withCollector is true, it sets up a cache with a custom metrics collector. // Otherwise, it sets up a cache without a custom metrics collector. @@ -52,7 +38,6 @@ func TestNewLRUCache(t *testing.T) { t.Run("default configuration", func(t *testing.T) { c, _ := setupCache[int](t, false) assert.Equal(t, "test_cache", c.cacheName) - assert.NotNil(t, c.metrics, "Cache metrics should not be nil") }) t.Run("with custom max cost", func(t *testing.T) { @@ -64,97 +49,168 @@ func TestNewLRUCache(t *testing.T) { c, collector := setupCache[int](t, true) assert.NotNil(t, c) assert.Equal(t, "test_cache", c.cacheName) - assert.Equal(t, collector, c.metrics, "Cache metrics should match the collector") + assert.Equal(t, collector, c.evictMetrics, "Cache metrics should match the collector") }) } func TestCacheSet(t *testing.T) { - c, collector := setupCache[string](t, true) + c, _ := setupCache[string](t, true) - collector.On("RecordSet", "test_cache").Once() c.Set("key", "value") - - collector.AssertCalled(t, "RecordSet", "test_cache") + value, found := c.Get("key") + assert.True(t, found, "Expected to find the key") + assert.Equal(t, "value", value, "Expected value to match") } func TestCacheGet(t *testing.T) { - c, collector := setupCache[string](t, true) - - collector.On("RecordSet", "test_cache").Once() - collector.On("RecordHit", "test_cache").Once() - collector.On("RecordMiss", "test_cache").Once() + c, _ := setupCache[string](t, true) c.Set("key", "value") - collector.AssertCalled(t, "RecordSet", "test_cache") value, found := c.Get("key") assert.True(t, found, "Expected to find the key") assert.Equal(t, "value", value, "Expected value to match") - collector.AssertCalled(t, "RecordHit", "test_cache") _, found = c.Get("non_existent") assert.False(t, found, "Expected not to find the key") - collector.AssertCalled(t, "RecordMiss", "test_cache") } func TestCacheExists(t *testing.T) { - c, collector := setupCache[string](t, true) - - collector.On("RecordSet", "test_cache").Once() - collector.On("RecordHit", "test_cache").Twice() - collector.On("RecordMiss", "test_cache").Once() + c, _ := setupCache[string](t, true) c.Set("key", "value") - collector.AssertCalled(t, "RecordSet", "test_cache") exists := c.Exists("key") assert.True(t, exists, "Expected the key to exist") - collector.AssertCalled(t, "RecordHit", "test_cache") exists = c.Exists("non_existent") assert.False(t, exists, "Expected the key not to exist") - collector.AssertCalled(t, "RecordMiss", "test_cache") } func TestCacheDelete(t *testing.T) { c, collector := setupCache[string](t, true) - collector.On("RecordSet", "test_cache").Once() - collector.On("RecordDelete", "test_cache").Once() - collector.On("RecordMiss", "test_cache").Once() collector.On("RecordEviction", "test_cache").Once() c.Set("key", "value") - collector.AssertCalled(t, "RecordSet", "test_cache") c.Delete("key") - collector.AssertCalled(t, "RecordDelete", "test_cache") collector.AssertCalled(t, "RecordEviction", "test_cache") _, found := c.Get("key") assert.False(t, found, "Expected not to find the deleted key") - collector.AssertCalled(t, "RecordMiss", "test_cache") } func TestCacheClear(t *testing.T) { c, collector := setupCache[string](t, true) - collector.On("RecordSet", "test_cache").Twice() - collector.On("RecordClear", "test_cache").Once() - collector.On("RecordMiss", "test_cache").Twice() collector.On("RecordEviction", "test_cache").Twice() c.Set("key1", "value1") c.Set("key2", "value2") - collector.AssertNumberOfCalls(t, "RecordSet", 2) c.Clear() - collector.AssertCalled(t, "RecordClear", "test_cache") collector.AssertNumberOfCalls(t, "RecordEviction", 2) _, found1 := c.Get("key1") _, found2 := c.Get("key2") assert.False(t, found1, "Expected not to find key1 after clear") assert.False(t, found2, "Expected not to find key2 after clear") - collector.AssertNumberOfCalls(t, "RecordMiss", 2) +} + +func TestCacheCount(t *testing.T) { + c, collector := setupCache[string](t, true) + + collector.On("RecordEviction", "test_cache").Times(3) + + c.Set("key1", "value1") + c.Set("key2", "value2") + c.Set("key3", "value3") + + assert.Equal(t, 3, c.Count(), "Expected count to be 3") + + c.Delete("key2") + assert.Equal(t, 2, c.Count(), "Expected count to be 2 after deletion") + collector.AssertNumberOfCalls(t, "RecordEviction", 1) + + c.Clear() + assert.Equal(t, 0, c.Count(), "Expected count to be 0 after clear") + collector.AssertNumberOfCalls(t, "RecordEviction", 3) +} + +func TestCacheKeys(t *testing.T) { + c, collector := setupCache[string](t, true) + + collector.On("RecordEviction", "test_cache").Times(3) + + c.Set("key1", "value1") + c.Set("key2", "value2") + c.Set("key3", "value3") + + keys := c.Keys() + assert.Len(t, keys, 3, "Expected 3 keys") + assert.ElementsMatch(t, []string{"key1", "key2", "key3"}, keys, "Keys do not match expected values") + + c.Delete("key2") + keys = c.Keys() + assert.Len(t, keys, 2, "Expected 2 keys after deletion") + assert.ElementsMatch(t, []string{"key1", "key3"}, keys, "Keys do not match expected values after deletion") + collector.AssertNumberOfCalls(t, "RecordEviction", 1) + + c.Clear() + keys = c.Keys() + assert.Len(t, keys, 0, "Expected no keys after clear") + collector.AssertNumberOfCalls(t, "RecordEviction", 3) +} + +func TestCacheValues(t *testing.T) { + c, collector := setupCache[string](t, true) + + collector.On("RecordEviction", "test_cache").Times(3) + + c.Set("key1", "value1") + c.Set("key2", "value2") + c.Set("key3", "value3") + + values := c.Values() + assert.Len(t, values, 3, "Expected 3 values") + assert.ElementsMatch(t, []string{"value1", "value2", "value3"}, values, "Values do not match expected values") + + c.Delete("key2") + values = c.Values() + assert.Len(t, values, 2, "Expected 2 values after deletion") + assert.ElementsMatch(t, []string{"value1", "value3"}, values, "Values do not match expected values after deletion") + collector.AssertNumberOfCalls(t, "RecordEviction", 1) + + c.Clear() + values = c.Values() + assert.Len(t, values, 0, "Expected no values after clear") + collector.AssertNumberOfCalls(t, "RecordEviction", 3) +} + +func TestCacheContents(t *testing.T) { + c, collector := setupCache[string](t, true) + + collector.On("RecordEviction", "test_cache").Times(3) + + c.Set("key1", "value1") + c.Set("key2", "value2") + c.Set("key3", "value3") + + contents := c.Contents() + assert.Contains(t, contents, "key1", "Contents should contain key1") + assert.Contains(t, contents, "key2", "Contents should contain key2") + assert.Contains(t, contents, "key3", "Contents should contain key3") + + c.Delete("key2") + contents = c.Contents() + assert.Contains(t, contents, "key1", "Contents should contain key1") + assert.NotContains(t, contents, "key2", "Contents should not contain key2") + assert.Contains(t, contents, "key3", "Contents should contain key3") + collector.AssertNumberOfCalls(t, "RecordEviction", 1) + + c.Clear() + contents = c.Contents() + assert.Equal(t, "[]", contents, "Contents should be empty after clear") + collector.AssertNumberOfCalls(t, "RecordEviction", 3) } diff --git a/pkg/cache/lru/metrics.go b/pkg/cache/lru/metrics.go deleted file mode 100644 index 3085a4e3a..000000000 --- a/pkg/cache/lru/metrics.go +++ /dev/null @@ -1,41 +0,0 @@ -package lru - -import ( - "github.com/prometheus/client_golang/prometheus" - - "github.com/trufflesecurity/trufflehog/v3/pkg/cache" -) - -// MetricsCollector should implement the collector interface. -var _ collector = (*MetricsCollector)(nil) - -// MetricsCollector extends the BaseMetricsCollector with Sized LRU specific metrics. -// It provides methods to record cache evictions. -type MetricsCollector struct { - // BaseMetricsCollector is embedded to provide the base metrics functionality. - cache.BaseMetricsCollector - - totalEvicts *prometheus.CounterVec -} - -// NewSizedLRUMetricsCollector initializes a new MetricsCollector with the provided namespace and subsystem. -func NewSizedLRUMetricsCollector(namespace, subsystem string) *MetricsCollector { - base := cache.GetMetricsCollector() - - totalEvicts := prometheus.NewCounterVec(prometheus.CounterOpts{ - Namespace: namespace, - Subsystem: subsystem, - Name: "evictions_total", - Help: "Total number of cache evictions.", - }, []string{"cache_name"}) - - return &MetricsCollector{ - BaseMetricsCollector: base, - totalEvicts: totalEvicts, - } -} - -// RecordEviction increments the total number of cache evictions for the specified cache. -func (c *MetricsCollector) RecordEviction(cacheName string) { - c.totalEvicts.WithLabelValues(cacheName).Inc() -} diff --git a/pkg/cache/metrics.go b/pkg/cache/metrics.go index 36a6ba867..449b36c00 100644 --- a/pkg/cache/metrics.go +++ b/pkg/cache/metrics.go @@ -1,8 +1,6 @@ package cache import ( - "sync" - "github.com/prometheus/client_golang/prometheus" "github.com/prometheus/client_golang/prometheus/promauto" @@ -19,10 +17,15 @@ type BaseMetricsCollector interface { RecordClear(cacheName string) } -// MetricsCollector encapsulates all Prometheus metrics with labels. +// EvictionMetricsCollector defines the interface for recording cache-specific eviction metrics. +type EvictionMetricsCollector interface { + RecordEviction(cacheName string) +} + +// baseCollector encapsulates all Prometheus metrics with labels. // It holds Prometheus counters for cache operations, which help track // the performance and usage of the cache. -type MetricsCollector struct { +type baseCollector struct { // Base metrics. hits *prometheus.CounterVec misses *prometheus.CounterVec @@ -32,77 +35,92 @@ type MetricsCollector struct { } func init() { - // Initialize the singleton MetricsCollector. + // Initialize the singleton baseCollector. // Set up Prometheus counters for cache operations (hits, misses, sets, deletes, clears). - collectorOnce.Do(func() { - collector = &MetricsCollector{ - hits: promauto.NewCounterVec(prometheus.CounterOpts{ - Namespace: common.MetricsNamespace, - Subsystem: common.MetricsSubsystem, - Name: "hits_total", - Help: "Total number of cache hits.", - }, []string{"cache_name"}), + baseMetricsInstance = &baseCollector{ + hits: promauto.NewCounterVec(prometheus.CounterOpts{ + Namespace: common.MetricsNamespace, + Subsystem: common.MetricsSubsystem, + Name: "hits_total", + Help: "Total number of cache hits.", + }, []string{"cache_name"}), - misses: promauto.NewCounterVec(prometheus.CounterOpts{ - Namespace: common.MetricsNamespace, - Subsystem: common.MetricsSubsystem, - Name: "misses_total", - Help: "Total number of cache misses.", - }, []string{"cache_name"}), + misses: promauto.NewCounterVec(prometheus.CounterOpts{ + Namespace: common.MetricsNamespace, + Subsystem: common.MetricsSubsystem, + Name: "misses_total", + Help: "Total number of cache misses.", + }, []string{"cache_name"}), - sets: promauto.NewCounterVec(prometheus.CounterOpts{ - Namespace: common.MetricsNamespace, - Subsystem: common.MetricsSubsystem, - Name: "sets_total", - Help: "Total number of cache set operations.", - }, []string{"cache_name"}), + sets: promauto.NewCounterVec(prometheus.CounterOpts{ + Namespace: common.MetricsNamespace, + Subsystem: common.MetricsSubsystem, + Name: "sets_total", + Help: "Total number of cache set operations.", + }, []string{"cache_name"}), - deletes: promauto.NewCounterVec(prometheus.CounterOpts{ - Namespace: common.MetricsNamespace, - Subsystem: common.MetricsSubsystem, - Name: "deletes_total", - Help: "Total number of cache delete operations.", - }, []string{"cache_name"}), + deletes: promauto.NewCounterVec(prometheus.CounterOpts{ + Namespace: common.MetricsNamespace, + Subsystem: common.MetricsSubsystem, + Name: "deletes_total", + Help: "Total number of cache delete operations.", + }, []string{"cache_name"}), - clears: promauto.NewCounterVec(prometheus.CounterOpts{ - Namespace: common.MetricsNamespace, - Subsystem: common.MetricsSubsystem, - Name: "clears_total", - Help: "Total number of cache clear operations.", - }, []string{"cache_name"}), - } - }) + clears: promauto.NewCounterVec(prometheus.CounterOpts{ + Namespace: common.MetricsNamespace, + Subsystem: common.MetricsSubsystem, + Name: "clears_total", + Help: "Total number of cache clear operations.", + }, []string{"cache_name"}), + } + + // Initialize the singleton evictionMetrics. + // Set up Prometheus counters for cache evictions. + evictionMetricsInstance = &evictionMetrics{ + evictions: promauto.NewCounterVec(prometheus.CounterOpts{ + Namespace: common.MetricsNamespace, + Subsystem: common.MetricsSubsystem, + Name: "evictions_total", + Help: "Total number of cache evictions.", + }, []string{"cache_name"}), + } } var ( - collectorOnce sync.Once // Ensures that the collector is initialized only once. - collector *MetricsCollector + baseMetricsInstance *baseCollector + evictionMetricsInstance *evictionMetrics ) -// GetMetricsCollector returns the singleton MetricsCollector instance. -// It panics if InitializeMetrics has not been called to ensure metrics are properly initialized. -// Must be called after InitializeMetrics to avoid runtime issues. -// If you do it before, BAD THINGS WILL HAPPEN. -func GetMetricsCollector() *MetricsCollector { - if collector == nil { - panic("MetricsCollector not initialized. Call InitializeMetrics first.") - } - return collector -} +// GetBaseMetricsCollector returns the singleton baseCollector instance. +func GetBaseMetricsCollector() BaseMetricsCollector { return baseMetricsInstance } + +// GetEvictionMetricsCollector returns the singleton evictionMetrics instance. +func GetEvictionMetricsCollector() EvictionMetricsCollector { return evictionMetricsInstance } // Implement BaseMetricsCollector interface methods. // RecordHit increments the counter for cache hits, tracking how often cache lookups succeed. -func (m *MetricsCollector) RecordHit(cacheName string) { m.hits.WithLabelValues(cacheName).Inc() } +func (m *baseCollector) RecordHit(cacheName string) { m.hits.WithLabelValues(cacheName).Inc() } // RecordMiss increments the counter for cache misses, tracking how often cache lookups fail. -func (m *MetricsCollector) RecordMiss(cacheName string) { m.misses.WithLabelValues(cacheName).Inc() } +func (m *baseCollector) RecordMiss(cacheName string) { m.misses.WithLabelValues(cacheName).Inc() } // RecordSet increments the counter for cache set operations, tracking how often items are added/updated. -func (m *MetricsCollector) RecordSet(cacheName string) { m.sets.WithLabelValues(cacheName).Inc() } +func (m *baseCollector) RecordSet(cacheName string) { m.sets.WithLabelValues(cacheName).Inc() } // RecordDelete increments the counter for cache delete operations, tracking how often items are removed. -func (m *MetricsCollector) RecordDelete(cacheName string) { m.deletes.WithLabelValues(cacheName).Inc() } +func (m *baseCollector) RecordDelete(cacheName string) { m.deletes.WithLabelValues(cacheName).Inc() } // RecordClear increments the counter for cache clear operations, tracking how often the cache is completely cleared. -func (m *MetricsCollector) RecordClear(cacheName string) { m.clears.WithLabelValues(cacheName).Inc() } +func (m *baseCollector) RecordClear(cacheName string) { m.clears.WithLabelValues(cacheName).Inc() } + +// evictionMetrics implements EvictionMetricsCollector interface. +type evictionMetrics struct { + evictions *prometheus.CounterVec +} + +// Implement EvictionMetricsCollector interface method. + +func (em *evictionMetrics) RecordEviction(cacheName string) { + em.evictions.WithLabelValues(cacheName).Inc() +}