Add support for measuring key size in caches (#1781)

Author: StephenButtolph

cache/lru_sized_cache.go

@@ -20,10 +20,10 @@ type sizedLRU[K comparable, V any] struct {
 	elements    linkedhashmap.LinkedHashmap[K, V]
 	maxSize     int
 	currentSize int
-	size        func(V) int
+	size        func(K, V) int
 }
 
-func NewSizedLRU[K comparable, V any](maxSize int, size func(V) int) Cacher[K, V] {
+func NewSizedLRU[K comparable, V any](maxSize int, size func(K, V) int) Cacher[K, V] {
 	return &sizedLRU[K, V]{
 		elements: linkedhashmap.New[K, V](),
 		maxSize:  maxSize,
@@ -67,25 +67,25 @@ func (c *sizedLRU[_, _]) PortionFilled() float64 {
 }
 
 func (c *sizedLRU[K, V]) put(key K, value V) {
-	valueSize := c.size(value)
-	if valueSize > c.maxSize {
+	newEntrySize := c.size(key, value)
+	if newEntrySize > c.maxSize {
 		c.flush()
 		return
 	}
 
 	if oldValue, ok := c.elements.Get(key); ok {
-		c.currentSize -= c.size(oldValue)
+		c.currentSize -= c.size(key, oldValue)
 	}
 
 	// Remove elements until the size of elements in the cache <= [c.maxSize].
-	for c.currentSize > c.maxSize-valueSize {
-		oldestKey, value, _ := c.elements.Oldest()
+	for c.currentSize > c.maxSize-newEntrySize {
+		oldestKey, oldestValue, _ := c.elements.Oldest()
 		c.elements.Delete(oldestKey)
-		c.currentSize -= c.size(value)
+		c.currentSize -= c.size(oldestKey, oldestValue)
 	}
 
 	c.elements.Put(key, value)
-	c.currentSize += valueSize
+	c.currentSize += newEntrySize
 }
 
 func (c *sizedLRU[K, V]) get(key K) (V, bool) {
@@ -101,7 +101,7 @@ func (c *sizedLRU[K, V]) get(key K) (V, bool) {
 func (c *sizedLRU[K, _]) evict(key K) {
 	if value, ok := c.elements.Get(key); ok {
 		c.elements.Delete(key)
-		c.currentSize -= c.size(value)
+		c.currentSize -= c.size(key, value)
 	}
 }
 

cache/lru_sized_cache_test.go

@@ -6,6 +6,8 @@ package cache
 import (
 	"testing"
 
+	"github.com/stretchr/testify/require"
+
 	"github.com/ava-labs/avalanchego/ids"
 )
 
@@ -20,3 +22,31 @@ func TestSizedLRUEviction(t *testing.T) {
 
 	TestEviction(t, cache)
 }
+
+func TestSizedLRUWrongKeyEvictionRegression(t *testing.T) {
+	require := require.New(t)
+
+	cache := NewSizedLRU[string, struct{}](
+		3,
+		func(key string, _ struct{}) int {
+			return len(key)
+		},
+	)
+
+	cache.Put("a", struct{}{})
+	cache.Put("b", struct{}{})
+	cache.Put("c", struct{}{})
+	cache.Put("dd", struct{}{})
+
+	_, ok := cache.Get("a")
+	require.False(ok)
+
+	_, ok = cache.Get("b")
+	require.False(ok)
+
+	_, ok = cache.Get("c")
+	require.True(ok)
+
+	_, ok = cache.Get("dd")
+	require.True(ok)
+}

cache/test_cacher.go

@@ -11,9 +11,9 @@ import (
 	"github.com/ava-labs/avalanchego/ids"
 )
 
-const TestIntSize = 8
+const TestIntSize = ids.IDLen + 8
 
-func TestIntSizeFunc(int64) int {
+func TestIntSizeFunc(ids.ID, int64) int {
 	return TestIntSize
 }
 

vms/platformvm/state/state.go

@@ -381,25 +381,25 @@ type txAndStatus struct {
 	status status.Status
 }
 
-func txSize(tx *txs.Tx) int {
+func txSize(_ ids.ID, tx *txs.Tx) int {
 	if tx == nil {
-		return pointerOverhead
+		return ids.IDLen + pointerOverhead
 	}
-	return len(tx.Bytes()) + pointerOverhead
+	return ids.IDLen + len(tx.Bytes()) + pointerOverhead
 }
 
-func txAndStatusSize(t *txAndStatus) int {
+func txAndStatusSize(_ ids.ID, t *txAndStatus) int {
 	if t == nil {
-		return pointerOverhead
+		return ids.IDLen + pointerOverhead
 	}
-	return len(t.tx.Bytes()) + wrappers.IntLen + pointerOverhead
+	return ids.IDLen + len(t.tx.Bytes()) + wrappers.IntLen + pointerOverhead
 }
 
-func blockSize(blk blocks.Block) int {
+func blockSize(_ ids.ID, blk blocks.Block) int {
 	if blk == nil {
-		return pointerOverhead
+		return ids.IDLen + pointerOverhead
 	}
-	return len(blk.Bytes()) + pointerOverhead
+	return ids.IDLen + len(blk.Bytes()) + pointerOverhead
 }
 
 func New(