[refactor] - Decouple Metrics From Cache Implementation (#3355)

* decouple metrics from cache logic

* delete

* address comments

* update
This commit is contained in:
ahrav
2024-10-04 13:25:10 -07:00
committed by GitHub
parent 1e5b831d16
commit c98c092a71
6 changed files with 613 additions and 180 deletions
+19
View File
@@ -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() }
+379
View File
@@ -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)
}
+37 -35
View File
@@ -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())
}
+104 -48
View File
@@ -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)
}
-41
View File
@@ -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()
}
+74 -56
View File
@@ -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()
}