Refactor cache implementations (#3239)

Author: StephenButtolph

cache/cache.go

@@ -25,19 +25,3 @@ type Cacher[K comparable, V any] interface {
 	// Returns fraction of cache currently filled (0 --> 1)
 	PortionFilled() float64
 }
-
-// Evictable allows the object to be notified when it is evicted
-type Evictable[K comparable] interface {
-	Key() K
-	Evict()
-}
-
-// Deduplicator acts as a best effort deduplication service
-type Deduplicator[K comparable, V Evictable[K]] interface {
-	// Deduplicate returns either the provided value, or a previously provided
-	// value with the same ID that hasn't yet been evicted
-	Deduplicate(V) V
-
-	// Flush removes all entries from the cache
-	Flush()
-}

cache/cachetest/cacher.go

@@ -23,11 +23,11 @@ var Tests = []struct {
 	Size int
 	Func func(t *testing.T, c cache.Cacher[ids.ID, int64])
 }{
-	{Size: 1, Func: TestBasic},
-	{Size: 2, Func: TestEviction},
+	{Size: 1, Func: Basic},
+	{Size: 2, Func: Eviction},
 }
 
-func TestBasic(t *testing.T, cache cache.Cacher[ids.ID, int64]) {
+func Basic(t *testing.T, cache cache.Cacher[ids.ID, int64]) {
 	require := require.New(t)
 
 	id1 := ids.ID{1}
@@ -62,7 +62,7 @@ func TestBasic(t *testing.T, cache cache.Cacher[ids.ID, int64]) {
 	require.Equal(expectedValue2, value)
 }
 
-func TestEviction(t *testing.T, cache cache.Cacher[ids.ID, int64]) {
+func Eviction(t *testing.T, cache cache.Cacher[ids.ID, int64]) {
 	require := require.New(t)
 
 	id1 := ids.ID{1}

cache/empty_cache.go cache/empty.gocache/empty_cache.go renamed to cache/empty.go

@@ -7,6 +7,7 @@ import "github.com/ava-labs/avalanchego/utils"
 
 var _ Cacher[struct{}, struct{}] = (*Empty[struct{}, struct{}])(nil)
 
+// Empty is a cache that doesn't store anything.
 type Empty[K any, V any] struct{}
 
 func (*Empty[K, V]) Put(K, V) {}

cache/lru/cache.go

@@ -0,0 +1,81 @@
+// Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package lru
+
+import (
+	"sync"
+
+	"github.com/ava-labs/avalanchego/cache"
+	"github.com/ava-labs/avalanchego/utils"
+	"github.com/ava-labs/avalanchego/utils/linked"
+)
+
+var _ cache.Cacher[struct{}, struct{}] = (*Cache[struct{}, struct{}])(nil)
+
+// Cache is a key value store with bounded size. If the size is attempted to be
+// exceeded, then an element is removed from the cache before the insertion is
+// done, based on evicting the least recently used value.
+type Cache[K comparable, V any] struct {
+	lock     sync.Mutex
+	elements *linked.Hashmap[K, V]
+	size     int
+}
+
+func NewCache[K comparable, V any](size int) *Cache[K, V] {
+	return &Cache[K, V]{
+		elements: linked.NewHashmap[K, V](),
+		size:     max(size, 1),
+	}
+}
+
+func (c *Cache[K, V]) Put(key K, value V) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	if c.elements.Len() == c.size {
+		oldestKey, _, _ := c.elements.Oldest()
+		c.elements.Delete(oldestKey)
+	}
+	c.elements.Put(key, value)
+}
+
+func (c *Cache[K, V]) Get(key K) (V, bool) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	val, ok := c.elements.Get(key)
+	if !ok {
+		return utils.Zero[V](), false
+	}
+	c.elements.Put(key, val) // Mark [k] as MRU.
+	return val, true
+}
+
+func (c *Cache[K, _]) Evict(key K) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	c.elements.Delete(key)
+}
+
+func (c *Cache[_, _]) Flush() {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	c.elements.Clear()
+}
+
+func (c *Cache[_, _]) Len() int {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	return c.elements.Len()
+}
+
+func (c *Cache[_, _]) PortionFilled() float64 {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	return float64(c.elements.Len()) / float64(c.size)
+}

cache/lru/cache_test.go

@@ -0,0 +1,21 @@
+// Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package lru
+
+import (
+	"testing"
+
+	"github.com/ava-labs/avalanchego/cache/cachetest"
+	"github.com/ava-labs/avalanchego/ids"
+)
+
+func TestCache(t *testing.T) {
+	c := NewCache[ids.ID, int64](1)
+	cachetest.Basic(t, c)
+}
+
+func TestCacheEviction(t *testing.T) {
+	c := NewCache[ids.ID, int64](2)
+	cachetest.Eviction(t, c)
+}

cache/lru/deduplicator.go

@@ -0,0 +1,82 @@
+// Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package lru
+
+import (
+	"sync"
+
+	"github.com/ava-labs/avalanchego/utils/linked"
+)
+
+// Evictable allows the object to be notified when it is evicted
+//
+// Deprecated: Remove this once the vertex state no longer uses it.
+type Evictable[K comparable] interface {
+	Key() K
+	Evict()
+}
+
+// Deduplicator is an LRU cache that notifies the objects when they are evicted.
+//
+// Deprecated: Remove this once the vertex state no longer uses it.
+type Deduplicator[K comparable, V Evictable[K]] struct {
+	lock      sync.Mutex
+	entryMap  map[K]*linked.ListElement[V]
+	entryList *linked.List[V]
+	size      int
+}
+
+// Deprecated: Remove this once the vertex state no longer uses it.
+func NewDeduplicator[K comparable, V Evictable[K]](size int) *Deduplicator[K, V] {
+	return &Deduplicator[K, V]{
+		entryMap:  make(map[K]*linked.ListElement[V]),
+		entryList: linked.NewList[V](),
+		size:      max(size, 1),
+	}
+}
+
+// Deduplicate returns either the provided value, or a previously provided value
+// with the same ID that hasn't yet been evicted
+func (d *Deduplicator[_, V]) Deduplicate(value V) V {
+	d.lock.Lock()
+	defer d.lock.Unlock()
+
+	key := value.Key()
+	if e, ok := d.entryMap[key]; !ok {
+		if d.entryList.Len() >= d.size {
+			e = d.entryList.Front()
+			d.entryList.MoveToBack(e)
+
+			delete(d.entryMap, e.Value.Key())
+			e.Value.Evict()
+
+			e.Value = value
+		} else {
+			e = &linked.ListElement[V]{
+				Value: value,
+			}
+			d.entryList.PushBack(e)
+		}
+		d.entryMap[key] = e
+	} else {
+		d.entryList.MoveToBack(e)
+
+		value = e.Value
+	}
+	return value
+}
+
+// Flush removes all entries from the cache
+func (d *Deduplicator[_, _]) Flush() {
+	d.lock.Lock()
+	defer d.lock.Unlock()
+
+	for d.entryList.Len() > 0 {
+		e := d.entryList.Front()
+		d.entryList.Remove(e)
+
+		delete(d.entryMap, e.Value.Key())
+		e.Value.Evict()
+	}
+}

cache/unique_cache_test.go cache/lru/deduplicator_test.gocache/unique_cache_test.go renamed to cache/lru/deduplicator_test.go

@@ -1,7 +1,7 @@
 // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
 // See the file LICENSE for licensing terms.
 
-package cache
+package lru
 
 import (
 	"testing"
@@ -24,10 +24,10 @@ func (e *evictable[_]) Evict() {
 	e.evicted++
 }
 
-func TestEvictableLRU(t *testing.T) {
+func TestDeduplicator(t *testing.T) {
 	require := require.New(t)
 
-	cache := EvictableLRU[ids.ID, *evictable[ids.ID]]{}
+	cache := NewDeduplicator[ids.ID, *evictable[ids.ID]](1)
 
 	expectedValue1 := &evictable[ids.ID]{id: ids.ID{1}}
 	require.Equal(expectedValue1, cache.Deduplicate(expectedValue1))
@@ -40,17 +40,4 @@ func TestEvictableLRU(t *testing.T) {
 	require.Equal(expectedValue2, returnedValue)
 	require.Equal(1, expectedValue1.evicted)
 	require.Zero(expectedValue2.evicted)
-
-	cache.Size = 2
-
-	expectedValue3 := &evictable[ids.ID]{id: ids.ID{2}}
-	returnedValue = cache.Deduplicate(expectedValue3)
-	require.Equal(expectedValue2, returnedValue)
-	require.Equal(1, expectedValue1.evicted)
-	require.Zero(expectedValue2.evicted)
-
-	cache.Flush()
-	require.Equal(1, expectedValue1.evicted)
-	require.Equal(1, expectedValue2.evicted)
-	require.Zero(expectedValue3.evicted)
 }