GOIMPORT/PATCHES/arc-cache-evict-callback.patch

diff --git a/arc.go b/arc.go
index 555225a..f4aaf38 100644
--- a/arc.go
+++ b/arc.go
@@ -18,17 +18,23 @@ type ARCCache struct {
 	size int // Size is the total capacity of the cache
 	p    int // P is the dynamic preference towards T1 or T2
 
-	t1 simplelru.LRUCache // T1 is the LRU for recently accessed items
-	b1 simplelru.LRUCache // B1 is the LRU for evictions from t1
+	t1 *simplelru.LRU // T1 is the LRU for recently accessed items
+	b1 *simplelru.LRU // B1 is the LRU for evictions from t1
 
-	t2 simplelru.LRUCache // T2 is the LRU for frequently accessed items
-	b2 simplelru.LRUCache // B2 is the LRU for evictions from t2
+	t2 *simplelru.LRU // T2 is the LRU for frequently accessed items
+	b2 *simplelru.LRU // B2 is the LRU for evictions from t2
 
 	lock sync.RWMutex
 }
 
 // NewARC creates an ARC of the given size
 func NewARC(size int) (*ARCCache, error) {
+	return NewARCWithEvict(size, nil)
+}
+
+// NewARCWithEvict creates an ARC of the given size with onEvict callback.
+func NewARCWithEvict(size int, onEvicted func(key interface{}, value interface{})) (*ARCCache, error) {
+	callback := simplelru.EvictCallback(onEvicted)
 	// Create the sub LRUs
 	b1, err := simplelru.NewLRU(size, nil)
 	if err != nil {
@@ -38,11 +44,11 @@ func NewARC(size int) (*ARCCache, error) {
 	if err != nil {
 		return nil, err
 	}
-	t1, err := simplelru.NewLRU(size, nil)
+	t1, err := simplelru.NewLRU(size, callback)
 	if err != nil {
 		return nil, err
 	}
-	t2, err := simplelru.NewLRU(size, nil)
+	t2, err := simplelru.NewLRU(size, callback)
 	if err != nil {
 		return nil, err
 	}
@@ -67,7 +73,7 @@ func (c *ARCCache) Get(key interface{}) (value interface{}, ok bool) {
 	// If the value is contained in T1 (recent), then
 	// promote it to T2 (frequent)
 	if val, ok := c.t1.Peek(key); ok {
-		c.t1.Remove(key)
+		c.t1.RemoveWithoutEvict(key)
 		c.t2.Add(key, val)
 		return val, ok
 	}
@@ -89,7 +95,7 @@ func (c *ARCCache) Add(key, value interface{}) {
 	// Check if the value is contained in T1 (recent), and potentially
 	// promote it to frequent T2
 	if c.t1.Contains(key) {
-		c.t1.Remove(key)
+		c.t1.RemoveWithoutEvict(key)
 		c.t2.Add(key, value)
 		return
 	}
@@ -209,21 +215,18 @@ func (c *ARCCache) Keys() []interface{} {
 	return append(k1, k2...)
 }
 
+func (c *ARCCache) subcaches() [4]simplelru.LRUCache {
+	return [4]simplelru.LRUCache{c.t1, c.t2, c.b1, c.b2}
+}
+
 // Remove is used to purge a key from the cache
 func (c *ARCCache) Remove(key interface{}) {
 	c.lock.Lock()
 	defer c.lock.Unlock()
-	if c.t1.Remove(key) {
-		return
-	}
-	if c.t2.Remove(key) {
-		return
-	}
-	if c.b1.Remove(key) {
-		return
-	}
-	if c.b2.Remove(key) {
-		return
+	for _, cache := range c.subcaches() {
+		if cache.Remove(key) {
+			return
+		}
 	}
 }
 
@@ -231,10 +234,9 @@ func (c *ARCCache) Remove(key interface{}) {
 func (c *ARCCache) Purge() {
 	c.lock.Lock()
 	defer c.lock.Unlock()
-	c.t1.Purge()
-	c.t2.Purge()
-	c.b1.Purge()
-	c.b2.Purge()
+	for _, cache := range c.subcaches() {
+		cache.Purge()
+	}
 }
 
 // Contains is used to check if the cache contains a key
diff --git a/arc_test.go b/arc_test.go
index e2d9b68..e597eda 100644
--- a/arc_test.go
+++ b/arc_test.go
@@ -4,6 +4,8 @@ import (
 	"math/rand"
 	"testing"
 	"time"
+
+	"google3/third_party/golang/cmp/cmp"
 )
 
 func init() {
@@ -375,3 +377,66 @@ func TestARC_Peek(t *testing.T) {
 		t.Errorf("should not have updated recent-ness of 1")
 	}
 }
+
+func TestARC_OnEvictedCallback(t *testing.T) {
+	type testCase struct {
+		size    int
+		data    []int
+		evicted []int
+	}
+
+	testCases := map[string]testCase{
+		"basic-lru-like": {
+			size:    2,
+			data:    []int{1, 2, 3, 4},
+			evicted: []int{1, 2},
+		},
+		"t1-to-t2-promotion-simple": {
+			size:    4,
+			data:    []int{1, 2, 1, 2, 3, 4, 5, 6},
+			evicted: []int{3, 4},
+		},
+		"adaptive-replacement-simple": {
+			size:    4,
+			data:    []int{1, 2, 3, 4, 5, 6, 1, 2, 7, 8},
+			evicted: []int{1, 2, 3, 4, 1, 5},
+		},
+		"MFU-not-evicted-by-new-entries": {
+			size:    4,
+			data:    []int{4, 4, 1, 2, 3, 4, 5, 6, 7, 8},
+			evicted: []int{1, 2, 3, 5},
+		},
+		"t2-full-b2-empty-new-items": {
+			size: 4,
+			data: []int{
+				// fill t2, b2 should be empty since its target size is zero
+				1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8,
+				// insert "new" items
+				1, 2, 3, 4,
+			},
+			evicted: []int{1, 2, 3, 4, 5, 1, 2, 3},
+		},
+	}
+
+	runTest := func(t *testing.T, tc testCase) {
+		evicted := []int{}
+		c, err := NewARCWithEvict(tc.size, func(key interface{}, _ interface{}) {
+			evicted = append(evicted, key.(int))
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		for _, val := range tc.data {
+			c.Add(val, val)
+		}
+
+		if !cmp.Equal(tc.evicted, evicted) {
+			t.Errorf("Expectations mismatch.\nWant: %+v\nGot:  %+v", tc.evicted, evicted)
+		}
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) { runTest(t, tc) })
+	}
+}
diff --git a/simplelru/lru.go b/simplelru/lru.go
index 5673773..0b40413 100644
--- a/simplelru/lru.go
+++ b/simplelru/lru.go
@@ -1,3 +1,4 @@
+// Package simplelru contains intefrace and implementation of linked-list based LRU cache.
 package simplelru
 
 import (
@@ -98,11 +99,21 @@ func (c *LRU) Peek(key interface{}) (value interface{}, ok bool) {
 // Remove removes the provided key from the cache, returning if the
 // key was contained.
 func (c *LRU) Remove(key interface{}) (present bool) {
-	if ent, ok := c.items[key]; ok {
+	ent, present := c.items[key]
+	if present {
 		c.removeElement(ent)
-		return true
 	}
-	return false
+	return present
+}
+
+// RemoveWithoutEvict removes the provided key from the cache without triggering
+// onEvict callback, returning if the key was contained.
+func (c *LRU) RemoveWithoutEvict(key interface{}) (present bool) {
+	ent, present := c.items[key]
+	if present {
+		c.removeWithoutEvict(ent)
+	}
+	return present
 }
 
 // RemoveOldest removes the oldest item from the cache.
@@ -152,10 +163,17 @@ func (c *LRU) removeOldest() {
 
 // removeElement is used to remove a given list element from the cache
 func (c *LRU) removeElement(e *list.Element) {
-	c.evictList.Remove(e)
-	kv := e.Value.(*entry)
-	delete(c.items, kv.key)
+	kv := c.removeWithoutEvict(e)
 	if c.onEvict != nil {
 		c.onEvict(kv.key, kv.value)
 	}
 }
+
+// removeWithoutEvict is used to remove a given list element from the cache
+// without triggering onEvict callback to be called.
+func (c *LRU) removeWithoutEvict(e *list.Element) *entry {
+	c.evictList.Remove(e)
+	kv := e.Value.(*entry)
+	delete(c.items, kv.key)
+	return kv
+}