Skip to content

Commit

Permalink
[feat] - Add SizedLRU Cache (#3344)
Browse files Browse the repository at this point in the history
* add impl for lru sized cache

* update error message

* address comments

* rename

* update comments
  • Loading branch information
ahrav authored Sep 30, 2024
1 parent 350db3a commit a5b0995
Show file tree
Hide file tree
Showing 5 changed files with 493 additions and 0 deletions.
64 changes: 64 additions & 0 deletions pkg/cache/decorator.go
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)
}
120 changes: 120 additions & 0 deletions pkg/cache/lru/lru.go
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)
}
160 changes: 160 additions & 0 deletions pkg/cache/lru/lru_test.go
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)
}
41 changes: 41 additions & 0 deletions pkg/cache/lru/metrics.go
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()
}
Loading

0 comments on commit a5b0995

Please sign in to comment.