diff --git a/lru/lru.go b/lru/lru.go index 47f13528..f0b67462 100644 --- a/lru/lru.go +++ b/lru/lru.go @@ -16,6 +16,7 @@ limitations under the License. package lru import ( + "fmt" "sync" groupcache "k8s.io/utils/internal/third_party/forked/golang/golang-lru" @@ -44,6 +45,15 @@ func NewWithEvictionFunc(size int, f EvictionFunc) *Cache { return c } +// SetEvictionFunc updates the eviction func +func (c *Cache) SetEvictionFunc(f EvictionFunc) error { + if c.cache.OnEvicted != nil { + return fmt.Errorf("lru cache eviction function is already set") + } + c.cache.OnEvicted = f + return nil +} + // Add adds a value to the cache. func (c *Cache) Add(key Key, value interface{}) { c.lock.Lock() diff --git a/lru/lru_test.go b/lru/lru_test.go index 2fcaf794..df8b0d5a 100644 --- a/lru/lru_test.go +++ b/lru/lru_test.go @@ -130,3 +130,31 @@ func TestEviction(t *testing.T) { t.Errorf("unexpected eviction data: key=%v val=%v", seenKey, seenVal) } } + +func TestSetEviction(t *testing.T) { + var seenKey Key + var seenVal interface{} + + lru := New(1) + + err := lru.SetEvictionFunc(func(key Key, value interface{}) { + seenKey = key + seenVal = value + }) + + if err != nil { + t.Errorf("unexpected error setting eviction function: %v", err) + } + + lru.Add(1, 2) + lru.Add(3, 4) + + if seenKey != 1 || seenVal != 2 { + t.Errorf("unexpected eviction data: key=%v val=%v", seenKey, seenVal) + } + + err = lru.SetEvictionFunc(func(key Key, value interface{}) {}) + if err == nil { + t.Errorf("expected error but got none") + } +}