From 6b3a92e7f232bdf1a76f1bace000ffe6981ea3f6 Mon Sep 17 00:00:00 2001
From: gVisor bot <gvisor-bot@google.com>
Date: Thu, 5 Dec 2019 00:17:24 +0800
Subject: [PATCH] Feature: lru cache add evict callback

---
 common/cache/lrucache.go      | 14 ++++++++++++++
 common/cache/lrucache_test.go | 21 +++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/common/cache/lrucache.go b/common/cache/lrucache.go
index 1cb278d9..a73b0b9e 100644
--- a/common/cache/lrucache.go
+++ b/common/cache/lrucache.go
@@ -11,6 +11,16 @@ import (
 // Option is part of Functional Options Pattern
 type Option func(*LruCache)
 
+// EvictCallback is used to get a callback when a cache entry is evicted
+type EvictCallback func(key interface{}, value interface{})
+
+// WithEvict set the evict callback
+func WithEvict(cb EvictCallback) Option {
+	return func(l *LruCache) {
+		l.onEvict = cb
+	}
+}
+
 // WithUpdateAgeOnGet update expires when Get element
 func WithUpdateAgeOnGet() Option {
 	return func(l *LruCache) {
@@ -42,6 +52,7 @@ type LruCache struct {
 	cache          map[interface{}]*list.Element
 	lru            *list.List // Front is least-recent
 	updateAgeOnGet bool
+	onEvict        EvictCallback
 }
 
 // NewLRUCache creates an LruCache
@@ -148,6 +159,9 @@ func (c *LruCache) deleteElement(le *list.Element) {
 	c.lru.Remove(le)
 	e := le.Value.(*entry)
 	delete(c.cache, e.key)
+	if c.onEvict != nil {
+		c.onEvict(e.key, e.value)
+	}
 }
 
 type entry struct {
diff --git a/common/cache/lrucache_test.go b/common/cache/lrucache_test.go
index 31f9a919..b296d6b9 100644
--- a/common/cache/lrucache_test.go
+++ b/common/cache/lrucache_test.go
@@ -115,3 +115,24 @@ func TestMaxSize(t *testing.T) {
 	_, ok = c.Get("foo")
 	assert.False(t, ok)
 }
+
+func TestExist(t *testing.T) {
+	c := NewLRUCache(WithSize(1))
+	c.Set(1, 2)
+	assert.True(t, c.Exist(1))
+	c.Set(2, 3)
+	assert.False(t, c.Exist(1))
+}
+
+func TestEvict(t *testing.T) {
+	temp := 0
+	evict := func(key interface{}, value interface{}) {
+		temp = key.(int) + value.(int)
+	}
+
+	c := NewLRUCache(WithEvict(evict), WithSize(1))
+	c.Set(1, 2)
+	c.Set(2, 3)
+
+	assert.Equal(t, temp, 3)
+}