[refactor] - Decouple Metrics From Cache Implementation (#3355)
* decouple metrics from cache logic * delete * address comments * update
This commit is contained in:
Vendored
+19
@@ -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() }
|
||||
|
||||
Vendored
+379
@@ -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)
|
||||
}
|
||||
Vendored
+37
-35
@@ -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())
|
||||
}
|
||||
|
||||
Vendored
+104
-48
@@ -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)
|
||||
}
|
||||
|
||||
Vendored
-41
@@ -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()
|
||||
}
|
||||
Vendored
+74
-56
@@ -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()
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user