feat: add eviction callback in LRU cache

cache/lru/cache.go

@@ -20,6 +20,20 @@ type Cache[K comparable, V any] struct {
 	lock     sync.Mutex
 	elements *linked.Hashmap[K, V]
 	size     int
+
+	// onEvict is called with the key and value of an entry before eviction, if set.
+	onEvict func(K, V)
+}
+
+// SetOnEvict sets a callback to be called with the key and value of an entry before eviction.
+// The onEvict callback is called while holding the cache lock.
+// Do not call any cache methods (Get, Put, Evict, Flush) from within the callback
+// as this will cause a deadlock. The callback should only be used for cleanup
+// operations like closing files or releasing resources.
+func (c *Cache[K, V]) SetOnEvict(cb func(K, V)) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	c.onEvict = cb
 }
 
 func NewCache[K comparable, V any](size int) *Cache[K, V] {
@@ -34,7 +48,10 @@ func (c *Cache[K, V]) Put(key K, value V) {
 	defer c.lock.Unlock()
 
 	if c.elements.Len() == c.size {
-		oldestKey, _, _ := c.elements.Oldest()
+		oldestKey, oldestValue, found := c.elements.Oldest()
+		if c.onEvict != nil && found {
+			c.onEvict(oldestKey, oldestValue)
+		}
 		c.elements.Delete(oldestKey)
 	}
 	c.elements.Put(key, value)
@@ -55,14 +72,26 @@ func (c *Cache[K, V]) Get(key K) (V, bool) {
 func (c *Cache[K, _]) Evict(key K) {
 	c.lock.Lock()
 	defer c.lock.Unlock()
-
+	if c.onEvict != nil {
+		if value, found := c.elements.Get(key); found {
+			c.onEvict(key, value)
+		}
+	}
 	c.elements.Delete(key)
 }
 
 func (c *Cache[_, _]) Flush() {
 	c.lock.Lock()
 	defer c.lock.Unlock()
 
+	// Call onEvict for each element before clearing
+	if c.onEvict != nil {
+		iter := c.elements.NewIterator()
+		for iter.Next() {
+			c.onEvict(iter.Key(), iter.Value())
+		}
+	}
+
 	c.elements.Clear()
 }
 

cache/lru/cache_test.go

@@ -6,6 +6,8 @@ package lru
 import (
 	"testing"
 
+	"github.com/stretchr/testify/require"
+
 	"github.com/ava-labs/avalanchego/cache/cachetest"
 	"github.com/ava-labs/avalanchego/ids"
 )
@@ -19,3 +21,89 @@ func TestCacheEviction(t *testing.T) {
 	c := NewCache[ids.ID, int64](2)
 	cachetest.Eviction(t, c)
 }
+
+func TestCacheOnEvict(t *testing.T) {
+	tests := []struct {
+		name            string
+		cacheSize       int
+		operations      func(*Cache[ids.ID, int64])
+		expectedEvicted map[ids.ID]int64
+	}{
+		{
+			name:      "OnEvict on Put with size limit",
+			cacheSize: 1,
+			operations: func(c *Cache[ids.ID, int64]) {
+				// Put first item
+				c.Put(ids.ID{1}, int64(1))
+				// Put second item, should evict first
+				c.Put(ids.ID{2}, int64(2))
+			},
+			expectedEvicted: map[ids.ID]int64{
+				{1}: int64(1),
+			},
+		},
+		{
+			name:      "OnEvict on explicit Evict",
+			cacheSize: 2,
+			operations: func(c *Cache[ids.ID, int64]) {
+				// Put two items
+				c.Put(ids.ID{1}, int64(1))
+				c.Put(ids.ID{2}, int64(2))
+				// Explicitly evict one
+				c.Evict(ids.ID{1})
+			},
+			expectedEvicted: map[ids.ID]int64{
+				{1}: int64(1),
+			},
+		},
+		{
+			name:      "OnEvict on Flush",
+			cacheSize: 2,
+			operations: func(c *Cache[ids.ID, int64]) {
+				// Put two items
+				c.Put(ids.ID{1}, int64(1))
+				c.Put(ids.ID{2}, int64(2))
+				// Flush should evict both
+				c.Flush()
+			},
+			expectedEvicted: map[ids.ID]int64{
+				{1}: int64(1),
+				{2}: int64(2),
+			},
+		},
+		{
+			name:      "OnEvict on multiple operations",
+			cacheSize: 2,
+			operations: func(c *Cache[ids.ID, int64]) {
+				// Put three items, should evict first
+				c.Put(ids.ID{1}, int64(1))
+				c.Put(ids.ID{2}, int64(2))
+				c.Put(ids.ID{3}, int64(3))
+				// Evict one more
+				c.Evict(ids.ID{2})
+				// Flush remaining
+				c.Flush()
+			},
+			expectedEvicted: map[ids.ID]int64{
+				{1}: int64(1), // evicted by Put
+				{2}: int64(2), // evicted by Evict
+				{3}: int64(3), // evicted by Flush
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c := NewCache[ids.ID, int64](tt.cacheSize)
+
+			evicted := make(map[ids.ID]int64)
+			c.SetOnEvict(func(key ids.ID, value int64) {
+				evicted[key] = value
+			})
+
+			tt.operations(c)
+
+			require.Equal(t, tt.expectedEvicted, evicted)
+		})
+	}
+}