Prevent sized LRU misuse by calculating the size only once (#3634)

cache/lru_sized_cache.go

@@ -12,20 +12,30 @@ import (
 
 var _ Cacher[struct{}, any] = (*sizedLRU[struct{}, any])(nil)
 
+// sizedElement is used to store the element with its size, so we don't
+// calculate the size multiple times.
+//
+// This ensures that any inconsistencies returned by the size function can not
+// corrupt the cache.
+type sizedElement[V any] struct {
+	value V
+	size  int
+}
+
 // sizedLRU is a key value store with bounded size. If the size is attempted to
 // be exceeded, then elements are removed from the cache until the bound is
 // honored, based on evicting the least recently used value.
 type sizedLRU[K comparable, V any] struct {
 	lock        sync.Mutex
-	elements    *linked.Hashmap[K, V]
+	elements    *linked.Hashmap[K, *sizedElement[V]]
 	maxSize     int
 	currentSize int
 	size        func(K, V) int
 }
 
 func NewSizedLRU[K comparable, V any](maxSize int, size func(K, V) int) Cacher[K, V] {
 	return &sizedLRU[K, V]{
-		elements: linked.NewHashmap[K, V](),
+		elements: linked.NewHashmap[K, *sizedElement[V]](),
 		maxSize:  maxSize,
 		size:     size,
 	}
@@ -80,35 +90,38 @@ func (c *sizedLRU[K, V]) put(key K, value V) {
 		return
 	}
 
-	if oldValue, ok := c.elements.Get(key); ok {
-		c.currentSize -= c.size(key, oldValue)
+	if oldElement, ok := c.elements.Get(key); ok {
+		c.currentSize -= oldElement.size
 	}
 
 	// Remove elements until the size of elements in the cache <= [c.maxSize].
 	for c.currentSize > c.maxSize-newEntrySize {
-		oldestKey, oldestValue, _ := c.elements.Oldest()
+		oldestKey, oldestElement, _ := c.elements.Oldest()
 		c.elements.Delete(oldestKey)
-		c.currentSize -= c.size(oldestKey, oldestValue)
+		c.currentSize -= oldestElement.size
 	}
 
-	c.elements.Put(key, value)
+	c.elements.Put(key, &sizedElement[V]{
+		value: value,
+		size:  newEntrySize,
+	})
 	c.currentSize += newEntrySize
 }
 
 func (c *sizedLRU[K, V]) get(key K) (V, bool) {
-	value, ok := c.elements.Get(key)
+	element, ok := c.elements.Get(key)
 	if !ok {
 		return utils.Zero[V](), false
 	}
 
-	c.elements.Put(key, value) // Mark [k] as MRU.
-	return value, true
+	c.elements.Put(key, element) // Mark [k] as MRU.
+	return element.value, true
 }
 
 func (c *sizedLRU[K, _]) evict(key K) {
-	if value, ok := c.elements.Get(key); ok {
+	if element, ok := c.elements.Get(key); ok {
 		c.elements.Delete(key)
-		c.currentSize -= c.size(key, value)
+		c.currentSize -= element.size
 	}
 }
 

cache/lru_sized_cache_test.go

@@ -53,3 +53,42 @@ func TestSizedLRUWrongKeyEvictionRegression(t *testing.T) {
 	_, ok = cache.Get("dd")
 	require.True(ok)
 }
+
+func TestSizedLRUSizeAlteringRegression(t *testing.T) {
+	require := require.New(t)
+
+	cache := NewSizedLRU[string, *string](
+		5,
+		func(key string, val *string) int {
+			if val != nil {
+				return len(key) + len(*val)
+			}
+
+			return len(key)
+		},
+	)
+
+	// put first value
+	expectedPortionFilled := 0.6
+	valueA := "ab"
+	cache.Put("a", &valueA)
+
+	require.Equal(expectedPortionFilled, cache.PortionFilled())
+
+	// mutate first value
+	valueA = "abcd"
+	require.Equal(expectedPortionFilled, cache.PortionFilled(), "after value A mutation, portion filled should be the same")
+
+	// put second value
+	expectedPortionFilled = 0.8
+	valueB := "bcd"
+	cache.Put("b", &valueB)
+
+	require.Equal(expectedPortionFilled, cache.PortionFilled())
+
+	_, ok := cache.Get("a")
+	require.False(ok, "key a shouldn't exist after b is put")
+
+	_, ok = cache.Get("b")
+	require.True(ok, "key b should exist")
+}