refactor: implement cache abstraction and unit test

cache/cache.go

@@ -0,0 +1,108 @@
+package cache
+
+import (
+	"container/list"
+)
+
+// Node represents a node eligible for caching.
+type Node interface {
+	GetKey() []byte
+}
+
+// Cache is an in-memory structure to persist nodes for quick access.
+// Please see lruCache for more details about why we need a custom
+// cache implementation.
+type Cache interface {
+	// Adds node to cache. If full and had to remove the oldest element,
+	// returns the oldest, otherwise nil.
+	// CONTRACT: node can never be nil. Otherwise, cache panics.
+	Add(node Node) Node
+
+	// Returns Node for the key, if exists. nil otherwise.
+	Get(key []byte) Node
+
+	// Has returns true if node with key exists in cache, false otherwise.
+	Has(key []byte) bool
+
+	// Remove removes node with key from cache. The removed node is returned.
+	// if not in cache, return nil.
+	Remove(key []byte) Node
+
+	// Len returns the cache length.
+	Len() int
+}
+
+// lruCache is an LRU cache implementation.
+// The motivation for using a custom cache implementation is to
+// allow for a custom max policy.
+//
+// Currently, the cache maximum is implemented in terms of the
+// number of nodes which is not intuitive to configure.
+// Instead, we are planning to add a byte maximum.
+// The alternative implementations do not allow for
+// customization and the ability to estimate the byte
+// size of the cache.
+type lruCache struct {
+	dict            map[string]*list.Element // FastNode cache.
+	maxElementCount int                      // FastNode the maximum number of nodes in the cache.
+	ll              *list.List               // LRU queue of cache elements. Used for deletion.
+}
+
+var _ Cache = (*lruCache)(nil)
+
+func New(maxElementCount int) Cache {
+	return &lruCache{
+		dict:            make(map[string]*list.Element),
+		maxElementCount: maxElementCount,
+		ll:              list.New(),
+	}
+}
+
+func (c *lruCache) Add(node Node) Node {
+	if e, exists := c.dict[string(node.GetKey())]; exists {
+		c.ll.MoveToFront(e)
+		old := e.Value
+		e.Value = node
+		return old.(Node)
+	}
+
+	elem := c.ll.PushFront(node)
+	c.dict[string(node.GetKey())] = elem
+
+	if c.ll.Len() > c.maxElementCount {
+		oldest := c.ll.Back()
+
+		return c.remove(oldest)
+	}
+	return nil
+}
+
+func (nc *lruCache) Get(key []byte) Node {
+	if ele, hit := nc.dict[string(key)]; hit {
+		nc.ll.MoveToFront(ele)
+		return ele.Value.(Node)
+	}
+	return nil
+}
+
+func (c *lruCache) Has(key []byte) bool {
+	_, exists := c.dict[string(key)]
+	return exists
+}
+
+func (nc *lruCache) Len() int {
+	return nc.ll.Len()
+}
+
+func (c *lruCache) Remove(key []byte) Node {
+	if elem, exists := c.dict[string(key)]; exists {
+		return c.remove(elem)
+	}
+	return nil
+}
+
+func (c *lruCache) remove(e *list.Element) Node {
+	removed := c.ll.Remove(e).(Node)
+	delete(c.dict, string(removed.GetKey()))
+	return removed
+}

cache/cache_bench_test.go

@@ -0,0 +1,69 @@
+package cache_test
+
+import (
+	"math/rand"
+	"testing"
+
+	"github.com/cosmos/iavl/cache"
+)
+
+func BenchmarkAdd(b *testing.B) {
+	b.ReportAllocs()
+	testcases := map[string]struct {
+		cacheMax int
+		keySize  int
+	}{
+		"small - max: 10K, key size - 10b": {
+			cacheMax: 10000,
+			keySize:  10,
+		},
+		"med - max: 100K, key size 20b": {
+			cacheMax: 100000,
+			keySize:  20,
+		},
+		"large - max: 1M, key size 30b": {
+			cacheMax: 1000000,
+			keySize:  30,
+		},
+	}
+
+	for name, tc := range testcases {
+		cache := cache.New(tc.cacheMax)
+		b.Run(name, func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				b.StopTimer()
+				key := randBytes(tc.keySize)
+				b.StartTimer()
+
+				_ = cache.Add(&testNode{
+					key: key,
+				})
+			}
+		})
+	}
+}
+
+func BenchmarkRemove(b *testing.B) {
+	b.ReportAllocs()
+
+	cache := cache.New(1000)
+	existentKeyMirror := [][]byte{}
+	// Populate cache
+	for i := 0; i < 50; i++ {
+		key := randBytes(1000)
+
+		existentKeyMirror = append(existentKeyMirror, key)
+
+		cache.Add(&testNode{
+			key: key,
+		})
+	}
+
+	randSeed := 498727689 // For deterministic tests
+	r := rand.New(rand.NewSource(int64(randSeed)))
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		key := existentKeyMirror[r.Intn(len(existentKeyMirror))]
+		_ = cache.Remove(key)
+	}
+}