From 22082ee84934afb0338608b2f468ae8b49ede8de Mon Sep 17 00:00:00 2001
From: Monis Khan <i@monis.app>
Date: Thu, 9 Feb 2023 14:31:41 -0500
Subject: [PATCH] lru: add support for setting eviction func

Signed-off-by: Monis Khan <mok@microsoft.com>
---
 lru/lru.go      | 10 +++++++++-
 lru/lru_test.go | 17 +++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/lru/lru.go b/lru/lru.go
index 5d0077ab..47f13528 100644
--- a/lru/lru.go
+++ b/lru/lru.go
@@ -5,7 +5,7 @@ Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
 
-    http://www.apache.org/licenses/LICENSE-2.0
+	http://www.apache.org/licenses/LICENSE-2.0
 
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
@@ -22,6 +22,7 @@ import (
 )
 
 type Key = groupcache.Key
+type EvictionFunc = func(key Key, value interface{})
 
 // Cache is a thread-safe fixed size LRU cache.
 type Cache struct {
@@ -36,6 +37,13 @@ func New(size int) *Cache {
 	}
 }
 
+// NewWithEvictionFunc creates an LRU of the given size with the given eviction func.
+func NewWithEvictionFunc(size int, f EvictionFunc) *Cache {
+	c := New(size)
+	c.cache.OnEvicted = f
+	return c
+}
+
 // 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 25b824bc..2fcaf794 100644
--- a/lru/lru_test.go
+++ b/lru/lru_test.go
@@ -113,3 +113,20 @@ func TestGetRace(t *testing.T) {
 	// let them run
 	time.Sleep(5 * time.Second)
 }
+
+func TestEviction(t *testing.T) {
+	var seenKey Key
+	var seenVal interface{}
+
+	lru := NewWithEvictionFunc(1, func(key Key, value interface{}) {
+		seenKey = key
+		seenVal = value
+	})
+
+	lru.Add(1, 2)
+	lru.Add(3, 4)
+
+	if seenKey != 1 || seenVal != 2 {
+		t.Errorf("unexpected eviction data: key=%v val=%v", seenKey, seenVal)
+	}
+}