Add a helper method to set eviction function after construction

Signed-off-by: Davanum Srinivas <[email protected]>

Author: dims

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()

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")
+	}
+}