fix: added cost() to cloneable interface

pkg/cache/memorycache/provider.go

@@ -112,11 +112,13 @@ func (provider *provider) Set(ctx context.Context, orgID valuer.UUID, cacheKey s
 	}
 
 	if cloneable, ok := data.(cachetypes.Cloneable); ok {
+		cost := cloneable.Size()
+		// Clamp to a minimum of 1: ristretto treats cost 0 specially and we
+		// never want zero-size entries to bypass admission accounting.
 		span.SetAttributes(attribute.Bool("memory.cloneable", true))
-		span.SetAttributes(attribute.Int64("memory.cost", 1))
+		span.SetAttributes(attribute.Int64("memory.cost", cost))
 		toCache := cloneable.Clone()
-		// In case of contention we are choosing to evict the cloneable entries first hence cost is set to 1
-		if ok := provider.cc.SetWithTTL(strings.Join([]string{orgID.StringValue(), cacheKey}, "::"), toCache, 1, ttl); !ok {
+		if ok := provider.cc.SetWithTTL(strings.Join([]string{orgID.StringValue(), cacheKey}, "::"), toCache, max(cost, 1), ttl); !ok {
 			return errors.New(errors.TypeInternal, errors.CodeInternal, "error writing to cache")
 		}
 
@@ -125,15 +127,15 @@ func (provider *provider) Set(ctx context.Context, orgID valuer.UUID, cacheKey s
 	}
 
 	toCache, err := provider.marshalBinary(ctx, data)
-	cost := int64(len(toCache))
 	if err != nil {
 		return err
 	}
+	cost := int64(len(toCache))
 
 	span.SetAttributes(attribute.Bool("memory.cloneable", false))
 	span.SetAttributes(attribute.Int64("memory.cost", cost))
 
-	if ok := provider.cc.SetWithTTL(strings.Join([]string{orgID.StringValue(), cacheKey}, "::"), toCache, 1, ttl); !ok {
+	if ok := provider.cc.SetWithTTL(strings.Join([]string{orgID.StringValue(), cacheKey}, "::"), toCache, max(cost, 1), ttl); !ok {
 		return errors.New(errors.TypeInternal, errors.CodeInternal, "error writing to cache")
 	}
 

pkg/cache/memorycache/provider_test.go

@@ -31,6 +31,10 @@ func (cloneable *CloneableA) Clone() cachetypes.Cacheable {
 	}
 }
 
+func (cloneable *CloneableA) Size() int64 {
+	return int64(len(cloneable.Key)) + 16
+}
+
 func (cloneable *CloneableA) MarshalBinary() ([]byte, error) {
 	return json.Marshal(cloneable)
 }
@@ -165,6 +169,50 @@ func TestSetGetWithDifferentTypes(t *testing.T) {
 	assert.Error(t, err)
 }
 
+// LargeCloneable reports a large byte cost so we can test ristretto eviction
+// without allocating the full payload in memory.
+type LargeCloneable struct {
+	Key  string
+	Cost int64
+}
+
+func (c *LargeCloneable) Clone() cachetypes.Cacheable {
+	return &LargeCloneable{Key: c.Key, Cost: c.Cost}
+}
+
+func (c *LargeCloneable) Size() int64 { return c.Cost }
+
+func (c *LargeCloneable) MarshalBinary() ([]byte, error) { return json.Marshal(c) }
+
+func (c *LargeCloneable) UnmarshalBinary(data []byte) error { return json.Unmarshal(data, c) }
+
+func TestCloneableCostTriggersEviction(t *testing.T) {
+	const maxCost int64 = 1 << 20 // 1 MiB
+	const perEntry int64 = 256 * 1024
+	const entries = 32 // 32 * 256 KiB = 8 MiB, well over MaxCost
+
+	c, err := New(context.Background(), factorytest.NewSettings(), cache.Config{Provider: "memory", Memory: cache.Memory{
+		NumCounters: 10 * 1000,
+		MaxCost:     maxCost,
+	}})
+	require.NoError(t, err)
+
+	orgID := valuer.GenerateUUID()
+	for i := 0; i < entries; i++ {
+		item := &LargeCloneable{Key: fmt.Sprintf("key-%d", i), Cost: perEntry}
+		assert.NoError(t, c.Set(context.Background(), orgID, fmt.Sprintf("key-%d", i), item, time.Minute))
+	}
+
+	metrics := c.(*provider).cc.Metrics
+	// Eviction (or admission rejection) must have kicked in: we wrote 32 entries
+	// each costing 256 KiB into a 1 MiB cache.
+	assert.Greater(t, metrics.KeysEvicted()+metrics.SetsRejected(), uint64(0),
+		"expected eviction or admission rejection once total cost exceeds MaxCost; got evicted=%d rejected=%d",
+		metrics.KeysEvicted(), metrics.SetsRejected())
+	// Net retained cost should not exceed MaxCost.
+	assert.LessOrEqual(t, int64(metrics.CostAdded()-metrics.CostEvicted()), maxCost)
+}
+
 func TestCloneableConcurrentSetGet(t *testing.T) {
 	cache, err := New(context.Background(), factorytest.NewSettings(), cache.Config{Provider: "memory", Memory: cache.Memory{
 		NumCounters: 10 * 1000,

pkg/cache/rediscache/provider_test.go

@@ -29,6 +29,10 @@ func (cacheable *CacheableA) Clone() cachetypes.Cacheable {
 	}
 }
 
+func (cacheable *CacheableA) Size() int64 {
+	return int64(len(cacheable.Key)) + 16
+}
+
 func (cacheable *CacheableA) MarshalBinary() ([]byte, error) {
 	return json.Marshal(cacheable)
 }

pkg/query-service/model/cacheable.go

@@ -41,6 +41,11 @@ func (c *GetWaterfallSpansForTraceWithMetadataCache) Clone() cachetypes.Cacheabl
 	}
 }
 
+func (c *GetWaterfallSpansForTraceWithMetadataCache) Size() int64 {
+	const perSpanBytes = 256
+	return int64(c.TotalSpans) * perSpanBytes
+}
+
 func (c *GetWaterfallSpansForTraceWithMetadataCache) MarshalBinary() (data []byte, err error) {
 	return json.Marshal(c)
 }
@@ -66,6 +71,16 @@ func (c *GetFlamegraphSpansForTraceCache) Clone() cachetypes.Cacheable {
 	}
 }
 
+func (c *GetFlamegraphSpansForTraceCache) Size() int64 {
+	const perSpanBytes = 128
+	var spans int64
+	for _, row := range c.SelectedSpans {
+		spans += int64(len(row))
+	}
+	spans += int64(len(c.TraceRoots))
+	return spans * perSpanBytes
+}
+
 func (c *GetFlamegraphSpansForTraceCache) MarshalBinary() (data []byte, err error) {
 	return json.Marshal(c)
 }

pkg/types/cachetypes/cacheable.go

@@ -18,6 +18,9 @@ type Cloneable interface {
 	// Creates a deep copy of the Cacheable. This method is useful for memory caches to avoid the need for serialization/deserialization. It also prevents
 	// race conditions in the memory cache.
 	Clone() Cacheable
+	// Size returns the approximate retained byte cost of this entry. Memory-backed
+	// caches use it as the ristretto cost so MaxCost-based eviction works.
+	Size() int64
 }
 
 func NewSha1CacheKey(val string) string {

pkg/types/querybuildertypes/querybuildertypesv5/cached.go

@@ -59,3 +59,21 @@ func (c *CachedData) Clone() cachetypes.Cacheable {
 
 	return clonedCachedData
 }
+
+// Size approximates the retained bytes of this CachedData. The dominant cost is
+// the serialized bucket values (json.RawMessage); other fields are fixed-size
+// or small strings.
+func (c *CachedData) Size() int64 {
+	var size int64
+	for _, b := range c.Buckets {
+		if b == nil {
+			continue
+		}
+		// Value is the bulk of the payload
+		size += int64(len(b.Value))
+	}
+	for _, w := range c.Warnings {
+		size += int64(len(w))
+	}
+	return size
+}

pkg/types/tracedetailtypes/waterfall.go

@@ -238,6 +238,13 @@ func (c *WaterfallCache) Clone() cachetypes.Cacheable {
 	}
 }
 
+// Size approximates the retained bytes. Each span dominates; use a per-span
+// constant rather than reflecting field-by-field.
+func (c *WaterfallCache) Size() int64 {
+	const perSpanBytes = 256
+	return int64(c.TotalSpans) * perSpanBytes
+}
+
 func (c *WaterfallCache) MarshalBinary() (data []byte, err error) {
 	return json.Marshal(c)
 }