From 2bbc7d88bdfc31bbd580125124f1edfdb6371041 Mon Sep 17 00:00:00 2001 From: Draco Date: Tue, 15 Jul 2025 11:03:30 -0400 Subject: [PATCH 1/3] feat: add eviction callback in LRU cache --- cache/lru/cache.go | 28 ++++++++++++++++++++++++++-- cache/lru/cache_test.go | 29 +++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/cache/lru/cache.go b/cache/lru/cache.go index dd83ee0cca65..94adba0ba69d 100644 --- a/cache/lru/cache.go +++ b/cache/lru/cache.go @@ -20,6 +20,16 @@ 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 evicted entry, if set. + onEvict func(K, V) +} + +// SetOnEvict sets a callback to be called with the key and value of an evicted entry. +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,8 +44,11 @@ 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, _ := c.elements.Oldest() c.elements.Delete(oldestKey) + if c.onEvict != nil { + c.onEvict(oldestKey, oldestValue) + } } c.elements.Put(key, value) } @@ -55,14 +68,25 @@ func (c *Cache[K, V]) Get(key K) (V, bool) { func (c *Cache[K, _]) Evict(key K) { c.lock.Lock() defer c.lock.Unlock() - c.elements.Delete(key) + if c.onEvict != nil { + value, _ := c.elements.Get(key) + c.onEvict(key, value) + } } 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..907c1c24c48d 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,30 @@ func TestCacheEviction(t *testing.T) { c := NewCache[ids.ID, int64](2) cachetest.Eviction(t, c) } + +func TestCacheFlushWithOnEvict(t *testing.T) { + c := NewCache[ids.ID, int64](2) + + // Track which elements were evicted + evicted := make(map[ids.ID]int64) + c.SetOnEvict(func(key ids.ID, value int64) { + evicted[key] = value + }) + + cachetest.Eviction(t, c) + require.Zero(t, c.Len()) + require.Len(t, evicted, 3) +} + +func TestCachePutWithOnEvict(t *testing.T) { + c := NewCache[ids.ID, int64](1) + + evicted := make(map[ids.ID]int64) + c.SetOnEvict(func(key ids.ID, value int64) { + evicted[key] = value + }) + + cachetest.Basic(t, c) + require.Len(t, evicted, 1) + require.Equal(t, evicted[ids.ID{1}], int64(1)) +} From 5991a645b10e03b940f642d3489c955581ea2bf9 Mon Sep 17 00:00:00 2001 From: Draco Date: Tue, 15 Jul 2025 15:39:24 -0400 Subject: [PATCH 2/3] add note on deadlock --- cache/lru/cache.go | 21 +++++++++++++-------- cache/lru/cache_test.go | 3 +-- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/cache/lru/cache.go b/cache/lru/cache.go index 94adba0ba69d..2d9f66f677f0 100644 --- a/cache/lru/cache.go +++ b/cache/lru/cache.go @@ -21,11 +21,15 @@ type Cache[K comparable, V any] struct { elements *linked.Hashmap[K, V] size int - // onEvict is called with the key and value of an evicted entry, if set. + // 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 evicted entry. +// 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() @@ -44,11 +48,11 @@ func (c *Cache[K, V]) Put(key K, value V) { defer c.lock.Unlock() if c.elements.Len() == c.size { - oldestKey, oldestValue, _ := c.elements.Oldest() - c.elements.Delete(oldestKey) - if c.onEvict != nil { + oldestKey, oldestValue, found := c.elements.Oldest() + if c.onEvict != nil && found { c.onEvict(oldestKey, oldestValue) } + c.elements.Delete(oldestKey) } c.elements.Put(key, value) } @@ -68,11 +72,12 @@ func (c *Cache[K, V]) Get(key K) (V, bool) { func (c *Cache[K, _]) Evict(key K) { c.lock.Lock() defer c.lock.Unlock() - c.elements.Delete(key) if c.onEvict != nil { - value, _ := c.elements.Get(key) - c.onEvict(key, value) + if value, found := c.elements.Get(key); found { + c.onEvict(key, value) + } } + c.elements.Delete(key) } func (c *Cache[_, _]) Flush() { diff --git a/cache/lru/cache_test.go b/cache/lru/cache_test.go index 907c1c24c48d..62e3ff1c3a7e 100644 --- a/cache/lru/cache_test.go +++ b/cache/lru/cache_test.go @@ -25,7 +25,6 @@ func TestCacheEviction(t *testing.T) { func TestCacheFlushWithOnEvict(t *testing.T) { c := NewCache[ids.ID, int64](2) - // Track which elements were evicted evicted := make(map[ids.ID]int64) c.SetOnEvict(func(key ids.ID, value int64) { evicted[key] = value @@ -46,5 +45,5 @@ func TestCachePutWithOnEvict(t *testing.T) { cachetest.Basic(t, c) require.Len(t, evicted, 1) - require.Equal(t, evicted[ids.ID{1}], int64(1)) + require.Equal(t, int64(1), evicted[ids.ID{1}]) } From 236e6fd2e31c82b2572246c6f78a9895554f7971 Mon Sep 17 00:00:00 2001 From: Draco Date: Thu, 17 Jul 2025 12:01:51 -0400 Subject: [PATCH 3/3] improved OnEvict tests --- cache/lru/cache_test.go | 100 ++++++++++++++++++++++++++++++++-------- 1 file changed, 80 insertions(+), 20 deletions(-) diff --git a/cache/lru/cache_test.go b/cache/lru/cache_test.go index 62e3ff1c3a7e..38c56635b9b4 100644 --- a/cache/lru/cache_test.go +++ b/cache/lru/cache_test.go @@ -22,28 +22,88 @@ func TestCacheEviction(t *testing.T) { cachetest.Eviction(t, c) } -func TestCacheFlushWithOnEvict(t *testing.T) { - c := NewCache[ids.ID, int64](2) - - evicted := make(map[ids.ID]int64) - c.SetOnEvict(func(key ids.ID, value int64) { - evicted[key] = value - }) +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 + }, + }, + } - cachetest.Eviction(t, c) - require.Zero(t, c.Len()) - require.Len(t, evicted, 3) -} + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + c := NewCache[ids.ID, int64](tt.cacheSize) -func TestCachePutWithOnEvict(t *testing.T) { - c := NewCache[ids.ID, int64](1) + evicted := make(map[ids.ID]int64) + c.SetOnEvict(func(key ids.ID, value int64) { + evicted[key] = value + }) - evicted := make(map[ids.ID]int64) - c.SetOnEvict(func(key ids.ID, value int64) { - evicted[key] = value - }) + tt.operations(c) - cachetest.Basic(t, c) - require.Len(t, evicted, 1) - require.Equal(t, int64(1), evicted[ids.ID{1}]) + require.Equal(t, tt.expectedEvicted, evicted) + }) + } }