-
Notifications
You must be signed in to change notification settings - Fork 1.7k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
* add impl for lru sized cache * update error message * address comments * rename * update comments
- Loading branch information
Showing
5 changed files
with
493 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,64 @@ | ||
package cache | ||
|
||
// WithMetrics is a decorator that adds metrics collection to any Cache implementation. | ||
type WithMetrics[T any] struct { | ||
wrapped Cache[T] | ||
metrics BaseMetricsCollector | ||
cacheName string | ||
} | ||
|
||
// NewCacheWithMetrics creates a new WithMetrics decorator that wraps the provided Cache | ||
// and collects metrics using the provided BaseMetricsCollector. | ||
// The cacheName parameter is used to identify the cache in the collected metrics. | ||
func NewCacheWithMetrics[T any](wrapped Cache[T], metrics BaseMetricsCollector, cacheName string) *WithMetrics[T] { | ||
return &WithMetrics[T]{ | ||
wrapped: wrapped, | ||
metrics: metrics, | ||
cacheName: cacheName, | ||
} | ||
} | ||
|
||
// Set sets the value for the given key in the cache. It also records a set metric | ||
// for the cache using the provided metrics collector and cache name. | ||
func (c *WithMetrics[T]) Set(key string, val T) { | ||
c.metrics.RecordSet(c.cacheName) | ||
c.wrapped.Set(key, val) | ||
} | ||
|
||
// Get retrieves the value for the given key from the underlying cache. It also records | ||
// a hit or miss metric for the cache using the provided metrics collector and cache name. | ||
func (c *WithMetrics[T]) Get(key string) (T, bool) { | ||
val, found := c.wrapped.Get(key) | ||
if found { | ||
c.metrics.RecordHit(c.cacheName) | ||
} else { | ||
c.metrics.RecordMiss(c.cacheName) | ||
} | ||
return val, found | ||
} | ||
|
||
// Exists checks if the given key exists in the cache. It records a hit or miss metric | ||
// for the cache using the provided metrics collector and cache name. | ||
func (c *WithMetrics[T]) Exists(key string) bool { | ||
found := c.wrapped.Exists(key) | ||
if found { | ||
c.metrics.RecordHit(c.cacheName) | ||
} else { | ||
c.metrics.RecordMiss(c.cacheName) | ||
} | ||
return found | ||
} | ||
|
||
// Delete removes the value for the given key from the cache. It also records a delete metric | ||
// for the cache using the provided metrics collector and cache name. | ||
func (c *WithMetrics[T]) Delete(key string) { | ||
c.wrapped.Delete(key) | ||
c.metrics.RecordDelete(c.cacheName) | ||
} | ||
|
||
// Clear removes all entries from the cache. It also records a clear metric | ||
// for the cache using the provided metrics collector and cache name. | ||
func (c *WithMetrics[T]) Clear() { | ||
c.wrapped.Clear() | ||
c.metrics.RecordClear(c.cacheName) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,120 @@ | ||
// Package lru provides a generic, size-limited, LRU (Least Recently Used) cache with optional | ||
// metrics collection and reporting. It wraps the golang-lru/v2 caching library, adding support for custom | ||
// metrics tracking cache hits, misses, evictions, and other cache operations. | ||
// | ||
// This package supports configuring key aspects of cache behavior, including maximum cache size, | ||
// and custom metrics collection. | ||
package lru | ||
|
||
import ( | ||
"fmt" | ||
|
||
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 | ||
} | ||
|
||
// 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 } | ||
} | ||
|
||
// 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) { | ||
// Default values for cache configuration. | ||
const defaultSize = 128_000 | ||
|
||
sizedLRU := &Cache[T]{ | ||
metrics: NewSizedLRUMetricsCollector(common.MetricsNamespace, common.MetricsSubsystem), | ||
cacheName: cacheName, | ||
} | ||
|
||
for _, opt := range opts { | ||
opt(sizedLRU) | ||
} | ||
|
||
// Provide a evict callback function to record evictions. | ||
onEvicted := func(string, T) { | ||
sizedLRU.metrics.RecordEviction(sizedLRU.cacheName) | ||
} | ||
|
||
lcache, err := lru.NewWithEvict[string, T](defaultSize, onEvicted) | ||
if err != nil { | ||
return nil, fmt.Errorf("failed to create lrusized cache: %w", err) | ||
} | ||
|
||
sizedLRU.cache = lcache | ||
|
||
return sizedLRU, nil | ||
} | ||
|
||
// 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) | ||
} | ||
|
||
// 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 | ||
} | ||
|
||
// 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) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,160 @@ | ||
package lru | ||
|
||
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) 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. | ||
func setupCache[T any](t *testing.T, withCollector bool) (*Cache[T], *mockCollector) { | ||
t.Helper() | ||
|
||
var collector *mockCollector | ||
var c *Cache[T] | ||
var err error | ||
|
||
if withCollector { | ||
collector = new(mockCollector) | ||
c, err = NewCache[T]("test_cache", WithMetricsCollector[T](collector)) | ||
} else { | ||
c, err = NewCache[T]("test_cache") | ||
} | ||
|
||
assert.NoError(t, err, "Failed to create cache") | ||
assert.NotNil(t, c, "Cache should not be nil") | ||
|
||
return c, collector | ||
} | ||
|
||
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) { | ||
c, _ := setupCache[int](t, false) | ||
assert.NotNil(t, c) | ||
}) | ||
|
||
t.Run("with metrics collector", func(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") | ||
}) | ||
} | ||
|
||
func TestCacheSet(t *testing.T) { | ||
c, collector := setupCache[string](t, true) | ||
|
||
collector.On("RecordSet", "test_cache").Once() | ||
c.Set("key", "value") | ||
|
||
collector.AssertCalled(t, "RecordSet", "test_cache") | ||
} | ||
|
||
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.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.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) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,41 @@ | ||
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() | ||
} |
Oops, something went wrong.