Update total cache size when updating the item

Signed-off-by: alanprot <[email protected]>

Author: alanprot

pkg/storage/tsdb/expanded_postings_cache.go

@@ -304,7 +304,6 @@ func (c *fifoCache[V]) getPromiseForKey(k string, fetch func() (V, int64, error)
 	r := &cacheEntryPromise[V]{
 		done: make(chan struct{}),
 	}
-
 	defer close(r.done)
 
 	if !c.cfg.Enabled {
@@ -326,8 +325,11 @@ func (c *fifoCache[V]) getPromiseForKey(k string, fetch func() (V, int64, error)
 	if ok && loaded.(*cacheEntryPromise[V]).isExpired(c.cfg.Ttl, c.timeNow()) {
 		if c.cachedValues.CompareAndSwap(k, loaded, r) {
 			r.v, r.sizeBytes, r.err = fetch()
+			r.sizeBytes += int64(len(k))
+			c.updateSize(loaded.(*cacheEntryPromise[V]).sizeBytes, r.sizeBytes)
 			loaded = r
 			r.ts = c.timeNow()
+			ok = false
 		}
 	}
 
@@ -377,6 +379,17 @@ func (c *fifoCache[V]) created(key string, sizeBytes int64) {
 	c.cachedBytes += sizeBytes
 }
 
+func (c *fifoCache[V]) updateSize(oldSize, newSizeBytes int64) {
+
+	if oldSize == newSizeBytes {
+		return
+	}
+
+	c.cachedMtx.Lock()
+	defer c.cachedMtx.Unlock()
+	c.cachedBytes += newSizeBytes - oldSize
+}
+
 type cacheEntryPromise[V any] struct {
 	ts        time.Time
 	sizeBytes int64

pkg/storage/tsdb/expanded_postings_cache_test.go

@@ -104,19 +104,23 @@ func TestFifoCacheExpire(t *testing.T) {
 			require.Equal(t, c.expectedFinalItems, totalCacheSize)
 
 			if c.ttlExpire {
+				cache.timeNow = func() time.Time {
+					return timeNow().Add(2 * c.cfg.Ttl)
+				}
+
 				for i := 0; i < numberOfKeys; i++ {
 					key := RepeatStringIfNeeded(fmt.Sprintf("key%d", i), keySize)
-					cache.timeNow = func() time.Time {
-						return timeNow().Add(2 * c.cfg.Ttl)
-					}
-					p, _ := cache.getPromiseForKey(key, func() (int, int64, error) {
-						return 2, 8, nil
+					originalSize := cache.cachedBytes
+					p, loaded := cache.getPromiseForKey(key, func() (int, int64, error) {
+						return 2, 18, nil
 					})
-					//require.False(t, loaded)
+					require.False(t, loaded)
 					v, err := p.result(context.Background())
 					require.NoError(t, err)
 					// New value
 					require.Equal(t, 2, v)
+					// Total Size Updated
+					require.Equal(t, originalSize+10, cache.cachedBytes)
 				}
 			}
 		})