diff --git a/cache/lru/cache.go b/cache/lru/cache.go index dd83ee0cca65..2d9f66f677f0 100644 --- a/cache/lru/cache.go +++ b/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,7 +72,11 @@ 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) } @@ -63,6 +84,14 @@ 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() } diff --git a/cache/lru/cache_test.go b/cache/lru/cache_test.go index 345ce42f1f76..38c56635b9b4 100644 --- a/cache/lru/cache_test.go +++ b/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) + }) + } +}