chore(querier): note dropped stats when merging cache buckets

Assisted-by: Claude Opus 5.5

pkg/cache/memorycache/provider.go

@@ -120,7 +120,7 @@ func (provider *provider) Set(ctx context.Context, orgID valuer.UUID, cacheKey s
 		span.SetAttributes(attribute.Bool("memory.cloneable", true))
 		span.SetAttributes(attribute.Int64("memory.cost", cost))
 		toCache := cloneable.Clone()
-		if ok := provider.cc.SetWithTTL(strings.Join([]string{orgID.StringValue(), cacheKey}, "::"), toCache, cost, ttl); !ok {
+		if ok := provider.setWithEviction(strings.Join([]string{orgID.StringValue(), cacheKey}, "::"), toCache, cost, ttl); !ok {
 			return errors.New(errors.TypeInternal, errors.CodeInternal, "error writing to cache")
 		}
 
@@ -137,7 +137,7 @@ func (provider *provider) Set(ctx context.Context, orgID valuer.UUID, cacheKey s
 	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, cost, ttl); !ok {
+	if ok := provider.setWithEviction(strings.Join([]string{orgID.StringValue(), cacheKey}, "::"), toCache, cost, ttl); !ok {
 		return errors.New(errors.TypeInternal, errors.CodeInternal, "error writing to cache")
 	}
 
@@ -228,3 +228,10 @@ func (provider *provider) unmarshalBinary(ctx context.Context, dest cachetypes.C
 	defer span.End()
 	return dest.UnmarshalBinary(fromCache)
 }
+
+// setWithEviction writes the key as a new item: ristretto evicts only on
+// insert, so updating an existing key in place can grow the cache past MaxCost.
+func (provider *provider) setWithEviction(key string, value any, cost int64, ttl time.Duration) bool {
+	provider.cc.Del(key)
+	return provider.cc.SetWithTTL(key, value, cost, ttl)
+}

pkg/cache/memorycache/provider_test.go

@@ -260,3 +260,46 @@ func TestCloneableConcurrentSetGet(t *testing.T) {
 		assert.NotSame(t, cachedCloneable, cloneables[i])
 	}
 }
+
+func TestSetGrowingExistingKeysStaysWithinMaxCost(t *testing.T) {
+	testCases := []struct {
+		name     string
+		newValue func(key string, size int) cachetypes.Cacheable
+	}{
+		{
+			name: "Cloneable",
+			newValue: func(key string, size int) cachetypes.Cacheable {
+				return &LargeCloneable{Key: key, CostHint: int64(size)}
+			},
+		},
+		{
+			name: "Cacheable",
+			newValue: func(key string, size int) cachetypes.Cacheable {
+				return &CacheableB{Key: strings.Repeat(key, size/len(key))}
+			},
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			var maxCost int64 = 1 << 20
+			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()
+			// 20 keys rewritten up to 100 KB each: 2 MB in total if updates never evict.
+			for round := 1; round <= 50; round++ {
+				for k := 0; k < 20; k++ {
+					key := fmt.Sprintf("key-%02d", k)
+					_ = c.Set(context.Background(), orgID, key, testCase.newValue(key, round*2000), time.Hour)
+				}
+			}
+
+			metrics := c.(*provider).cc.Metrics
+			assert.LessOrEqual(t, int64(metrics.CostAdded())-int64(metrics.CostEvicted()), maxCost)
+		})
+	}
+}

pkg/querier/bucket_cache.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"log/slog"
+	"math/bits"
 	"slices"
 	"time"
 
@@ -16,6 +17,8 @@ import (
 	"github.com/SigNoz/signoz/pkg/valuer"
 )
 
+const stepsPerBucket = 300
+
 // BucketCache is the interface for bucket-based caching.
 type BucketCache interface {
 	// cached portion + list of gaps to fetch
@@ -175,7 +178,8 @@ func (bc *bucketCache) Put(ctx context.Context, orgID valuer.UUID, q qbtypes.Que
 	}
 
 	// Convert trimmed result to buckets with adjusted boundaries
-	freshBuckets := bc.resultToBuckets(ctx, trimmedResult, cachableStartMs, cachableEndMs)
+	bucketWidthMs := bucketWidth(stepMs)
+	freshBuckets := bc.resultToBuckets(ctx, trimmedResult, cachableStartMs, cachableEndMs, bucketWidthMs)
 
 	// If no fresh buckets and no existing data, don't cache
 	if len(freshBuckets) == 0 && len(existingData.Buckets) == 0 {
@@ -183,7 +187,7 @@ func (bc *bucketCache) Put(ctx context.Context, orgID valuer.UUID, q qbtypes.Que
 	}
 
 	// Merge with existing buckets
-	mergedBuckets := bc.mergeAndDeduplicateBuckets(existingData.Buckets, freshBuckets)
+	mergedBuckets := bc.mergeAndDeduplicateBuckets(ctx, existingData.Buckets, freshBuckets, bucketWidthMs)
 
 	// Update warnings
 	allWarnings := append(existingData.Warnings, trimmedResult.Warnings...)
@@ -646,7 +650,7 @@ func (bc *bucketCache) isEmptyResult(result *qbtypes.Result) (isEmpty bool, isFi
 }
 
 // resultToBuckets converts a query result into time-based buckets.
-func (bc *bucketCache) resultToBuckets(ctx context.Context, result *qbtypes.Result, startMs, endMs uint64) []*qbtypes.CachedBucket {
+func (bc *bucketCache) resultToBuckets(ctx context.Context, result *qbtypes.Result, startMs, endMs, widthMs uint64) []*qbtypes.CachedBucket {
 	// Check if result is empty
 	isEmpty, isFiltered := bc.isEmptyResult(result)
 
@@ -657,48 +661,87 @@ func (bc *bucketCache) resultToBuckets(ctx context.Context, result *qbtypes.Resu
 		return nil
 	}
 
-	// For now, create a single bucket for the entire range
-	// In the future, we could split large ranges into smaller buckets
-	valueBytes, err := json.Marshal(result.Value)
-	if err != nil {
-		bc.logger.ErrorContext(ctx, "failed to marshal result value", errors.Attr(err))
+	tsData, ok := result.Value.(*qbtypes.TimeSeriesData)
+	if !ok {
 		return nil
 	}
 
-	// Always create a bucket, even for empty filtered results
+	// Always create buckets, even for empty filtered results
 	// This ensures we don't re-query for data that doesn't exist
-	return []*qbtypes.CachedBucket{
-		{
-			StartMs: startMs,
-			EndMs:   endMs,
-			Type:    result.Type,
-			Value:   valueBytes,
-			Stats:   result.Stats,
-		},
+	return bc.splitIntoBuckets(ctx, tsData, result.Type, result.Stats, startMs, endMs, widthMs)
+}
+
+// splitIntoBuckets cuts [startMs, endMs) at multiples of widthMs, one bucket per piece.
+func (bc *bucketCache) splitIntoBuckets(ctx context.Context, tsData *qbtypes.TimeSeriesData, resultType qbtypes.RequestType, stats qbtypes.ExecStats, startMs, endMs, widthMs uint64) []*qbtypes.CachedBucket {
+	buckets := make([]*qbtypes.CachedBucket, 0, 1)
+	for chunkStartMs := startMs; chunkStartMs < endMs; {
+		chunkEndMs := endMs
+		if widthMs > 0 {
+			chunkEndMs = min(endMs, (chunkStartMs/widthMs+1)*widthMs)
+		}
+
+		// A PromQL value at endMs is inside the window and travels with the last piece.
+		upToMs := chunkEndMs
+		if chunkEndMs == endMs {
+			upToMs = endMs + 1
+		}
+		value, err := json.Marshal(filterTimeSeriesData(tsData, chunkStartMs, upToMs))
+		if err != nil {
+			bc.logger.ErrorContext(ctx, "failed to marshal result value", errors.Attr(err))
+			return nil
+		}
+
+		buckets = append(buckets, &qbtypes.CachedBucket{
+			StartMs: chunkStartMs,
+			EndMs:   chunkEndMs,
+			Type:    resultType,
+			Value:   value,
+			Stats:   statsShare(stats, startMs, endMs, chunkStartMs, chunkEndMs),
+		})
+		chunkStartMs = chunkEndMs
+	}
+	return buckets
+}
+
+// bucketWidth depends only on the step, which is part of the cache key, so every window of a query
+// cuts at the same edges. Zero means no cutting.
+func bucketWidth(stepMs uint64) uint64 {
+	return stepMs * stepsPerBucket
+}
+
+// statsShare is the duration-weighted part of stats for [fromMs, toMs) within [startMs, endMs).
+// Shares are differences of running totals, so the pieces of a split add back up to stats exactly.
+func statsShare(stats qbtypes.ExecStats, startMs, endMs, fromMs, toMs uint64) qbtypes.ExecStats {
+	totalMs := endMs - startMs
+	if totalMs == 0 {
+		return stats
+	}
+	upTo := func(v, atMs uint64) uint64 {
+		hi, lo := bits.Mul64(v, atMs-startMs)
+		q, _ := bits.Div64(hi, lo, totalMs)
+		return q
+	}
+	return qbtypes.ExecStats{
+		RowsScanned:   upTo(stats.RowsScanned, toMs) - upTo(stats.RowsScanned, fromMs),
+		BytesScanned:  upTo(stats.BytesScanned, toMs) - upTo(stats.BytesScanned, fromMs),
+		DurationMS:    upTo(stats.DurationMS, toMs) - upTo(stats.DurationMS, fromMs),
+		StepIntervals: stats.StepIntervals,
 	}
 }
 
 // mergeAndDeduplicateBuckets combines and deduplicates bucket lists.
-func (bc *bucketCache) mergeAndDeduplicateBuckets(existing, fresh []*qbtypes.CachedBucket) []*qbtypes.CachedBucket {
-	// Create a map to deduplicate by time range
-	bucketMap := make(map[string]*qbtypes.CachedBucket)
+// Buckets stay disjoint and never cross a multiple of widthMs: existing buckets overlapping a
+// fresh one are folded into it, keeping their values outside the fresh range, and the union is
+// cut back into width-sized pieces.
+func (bc *bucketCache) mergeAndDeduplicateBuckets(ctx context.Context, existing, fresh []*qbtypes.CachedBucket, widthMs uint64) []*qbtypes.CachedBucket {
+	result := make([]*qbtypes.CachedBucket, 0, len(existing)+len(fresh))
 
 	// Add existing buckets
-	for _, bucket := range existing {
-		key := fmt.Sprintf("%d-%d", bucket.StartMs, bucket.EndMs)
-		bucketMap[key] = bucket
-	}
+	result = append(result, existing...)
 
 	// Add/update with fresh buckets
 	for _, bucket := range fresh {
-		key := fmt.Sprintf("%d-%d", bucket.StartMs, bucket.EndMs)
-		bucketMap[key] = bucket
-	}
-
-	// Convert back to slice with pre-allocated capacity
-	result := make([]*qbtypes.CachedBucket, 0, len(bucketMap))
-	for _, bucket := range bucketMap {
-		result = append(result, bucket)
+		result = bc.upsertBucket(ctx, result, bucket, widthMs)
 	}
 
 	// Sort by start time
@@ -715,6 +758,67 @@ func (bc *bucketCache) mergeAndDeduplicateBuckets(existing, fresh []*qbtypes.Cac
 	return result
 }
 
+// upsertBucket adds fresh to buckets, absorbing every bucket it overlaps.
+func (bc *bucketCache) upsertBucket(ctx context.Context, buckets []*qbtypes.CachedBucket, fresh *qbtypes.CachedBucket, widthMs uint64) []*qbtypes.CachedBucket {
+	kept := make([]*qbtypes.CachedBucket, 0, len(buckets)+1)
+	merged := []*qbtypes.CachedBucket{fresh}
+	for _, bucket := range buckets {
+		if bucket.StartMs >= fresh.EndMs || fresh.StartMs >= bucket.EndMs {
+			kept = append(kept, bucket)
+			continue
+		}
+		if fresh.StartMs <= bucket.StartMs && bucket.EndMs <= fresh.EndMs {
+			continue
+		}
+		if trimmed := bc.excludeRange(ctx, bucket, fresh.StartMs, fresh.EndMs); trimmed != nil {
+			merged = append(merged, trimmed)
+		}
+	}
+	if len(merged) == 1 {
+		return append(kept, fresh)
+	}
+
+	startMs, endMs := fresh.StartMs, fresh.EndMs
+	for _, bucket := range merged[1:] {
+		startMs = min(startMs, bucket.StartMs)
+		endMs = max(endMs, bucket.EndMs)
+	}
+	// TODO(tvats): stats of the absorbed buckets outside fresh's range are dropped; add their share.
+	pieces := bc.splitIntoBuckets(ctx, bc.mergeTimeSeriesValues(ctx, merged), fresh.Type, fresh.Stats, startMs, endMs, widthMs)
+	if pieces == nil {
+		return append(kept, fresh)
+	}
+	return append(kept, pieces...)
+}
+
+// excludeRange returns bucket without its values in [startMs, endMs), or nil if it cannot be decoded.
+func (bc *bucketCache) excludeRange(ctx context.Context, bucket *qbtypes.CachedBucket, startMs, endMs uint64) *qbtypes.CachedBucket {
+	var tsData *qbtypes.TimeSeriesData
+	if err := json.Unmarshal(bucket.Value, &tsData); err != nil || tsData == nil {
+		bc.logger.ErrorContext(ctx, "failed to unmarshal time series data", errors.Attr(err))
+		return nil
+	}
+	for _, aggBucket := range tsData.Aggregations {
+		for _, series := range aggBucket.Series {
+			series.Values = slices.DeleteFunc(series.Values, func(v *qbtypes.TimeSeriesValue) bool {
+				return uint64(v.Timestamp) >= startMs && uint64(v.Timestamp) < endMs
+			})
+		}
+	}
+	value, err := json.Marshal(tsData)
+	if err != nil {
+		bc.logger.ErrorContext(ctx, "failed to marshal time series data", errors.Attr(err))
+		return nil
+	}
+	return &qbtypes.CachedBucket{
+		StartMs: bucket.StartMs,
+		EndMs:   bucket.EndMs,
+		Type:    bucket.Type,
+		Value:   value,
+		Stats:   bucket.Stats,
+	}
+}
+
 // deduplicateWarnings removes duplicate warnings.
 func (bc *bucketCache) deduplicateWarnings(warnings []string) []string {
 	return dedupeWarnings(warnings)
@@ -817,41 +921,7 @@ func (bc *bucketCache) filterResultToTimeRange(result *qbtypes.Result, q qbtypes
 	switch result.Type {
 	case qbtypes.RequestTypeTimeSeries, qbtypes.RequestTypeHeatmap:
 		if tsData, ok := result.Value.(*qbtypes.TimeSeriesData); ok {
-			filteredData := &qbtypes.TimeSeriesData{
-				Aggregations: make([]*qbtypes.AggregationBucket, 0, len(tsData.Aggregations)),
-			}
-
-			for _, aggBucket := range tsData.Aggregations {
-				filteredBucket := &qbtypes.AggregationBucket{
-					Index:  aggBucket.Index,
-					Alias:  aggBucket.Alias,
-					Meta:   aggBucket.Meta,
-					Series: make([]*qbtypes.TimeSeries, 0, len(aggBucket.Series)),
-				}
-
-				for _, series := range aggBucket.Series {
-					filteredSeries := &qbtypes.TimeSeries{
-						Labels: series.Labels,
-						Values: make([]*qbtypes.TimeSeriesValue, 0, len(series.Values)),
-					}
-
-					// Filter values to only include those within the requested time range
-					for _, value := range series.Values {
-						timestampMs := uint64(value.Timestamp)
-						if timestampMs >= startMs && timestampMs <= maxTimestampMs {
-							filteredSeries.Values = append(filteredSeries.Values, value)
-						}
-					}
-
-					// Always add series to preserve structure (even if empty)
-					filteredBucket.Series = append(filteredBucket.Series, filteredSeries)
-				}
-
-				// Only add bucket if it has series
-				if len(filteredBucket.Series) > 0 {
-					filteredData.Aggregations = append(filteredData.Aggregations, filteredBucket)
-				}
-			}
+			filteredData := filterTimeSeriesData(tsData, startMs, maxTimestampMs+1)
 
 			bc.trimHeatmapAxisToTheWindow(q, filteredData)
 
@@ -890,3 +960,44 @@ func (bc *bucketCache) trimHeatmapAxisToTheWindow(q qbtypes.Query, tsData *qbtyp
 		aggBucket.TrimAxisToCountedBuckets()
 	}
 }
+
+// filterTimeSeriesData copies tsData keeping only values in [startMs, endMs); series survive empty, aggregations without series do not.
+func filterTimeSeriesData(tsData *qbtypes.TimeSeriesData, startMs, endMs uint64) *qbtypes.TimeSeriesData {
+	filteredData := &qbtypes.TimeSeriesData{
+		Aggregations: make([]*qbtypes.AggregationBucket, 0, len(tsData.Aggregations)),
+	}
+
+	for _, aggBucket := range tsData.Aggregations {
+		filteredBucket := &qbtypes.AggregationBucket{
+			Index:  aggBucket.Index,
+			Alias:  aggBucket.Alias,
+			Meta:   aggBucket.Meta,
+			Series: make([]*qbtypes.TimeSeries, 0, len(aggBucket.Series)),
+		}
+
+		for _, series := range aggBucket.Series {
+			filteredSeries := &qbtypes.TimeSeries{
+				Labels: series.Labels,
+				Values: make([]*qbtypes.TimeSeriesValue, 0, len(series.Values)),
+			}
+
+			// Filter values to only include those within the requested time range
+			for _, value := range series.Values {
+				timestampMs := uint64(value.Timestamp)
+				if timestampMs >= startMs && timestampMs < endMs {
+					filteredSeries.Values = append(filteredSeries.Values, value)
+				}
+			}
+
+			// Always add series to preserve structure (even if empty)
+			filteredBucket.Series = append(filteredBucket.Series, filteredSeries)
+		}
+
+		// Only add bucket if it has series
+		if len(filteredBucket.Series) > 0 {
+			filteredData.Aggregations = append(filteredData.Aggregations, filteredBucket)
+		}
+	}
+
+	return filteredData
+}

pkg/querier/bucket_cache_test.go

@@ -1465,3 +1465,80 @@ func TestBucketCache_NoCache(t *testing.T) {
 	// The actual NoCache logic is implemented in querier.run(), not in bucket cache
 	// This test verifies that the cache works normally and NoCache bypasses it at a higher level
 }
+
+func TestBucketCache_Put_OverlappingBuckets(t *testing.T) {
+	// Buckets are cut at multiples of 300 steps, so a 10ms step fills 3s-aligned buckets.
+	testCases := []struct {
+		name            string
+		windows         []qbtypes.TimeRange
+		expectedBuckets []qbtypes.TimeRange
+	}{
+		{
+			name:            "SlidingWindows_FillAlignedBuckets",
+			windows:         []qbtypes.TimeRange{{From: 1000, To: 5000}, {From: 2000, To: 6000}, {From: 3000, To: 7000}},
+			expectedBuckets: []qbtypes.TimeRange{{From: 1000, To: 3000}, {From: 3000, To: 6000}, {From: 6000, To: 7000}},
+		},
+		{
+			name:            "DisjointWindows_KeepAll",
+			windows:         []qbtypes.TimeRange{{From: 1000, To: 3000}, {From: 5000, To: 7000}},
+			expectedBuckets: []qbtypes.TimeRange{{From: 1000, To: 3000}, {From: 5000, To: 6000}, {From: 6000, To: 7000}},
+		},
+		{
+			name:            "PartialOverlap_MergeOnlyOverlapping",
+			windows:         []qbtypes.TimeRange{{From: 1000, To: 3000}, {From: 5000, To: 7000}, {From: 2000, To: 4000}},
+			expectedBuckets: []qbtypes.TimeRange{{From: 1000, To: 3000}, {From: 3000, To: 4000}, {From: 5000, To: 6000}, {From: 6000, To: 7000}},
+		},
+		{
+			name:            "NarrowerWindow_KeepsBucketEdges",
+			windows:         []qbtypes.TimeRange{{From: 1000, To: 7000}, {From: 3000, To: 5000}},
+			expectedBuckets: []qbtypes.TimeRange{{From: 1000, To: 3000}, {From: 3000, To: 6000}, {From: 6000, To: 7000}},
+		},
+		{
+			name:            "WiderWindow_AbsorbsCoveredBuckets",
+			windows:         []qbtypes.TimeRange{{From: 2000, To: 3000}, {From: 4000, To: 5000}, {From: 1000, To: 7000}},
+			expectedBuckets: []qbtypes.TimeRange{{From: 1000, To: 3000}, {From: 3000, To: 6000}, {From: 6000, To: 7000}},
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			memCache := createTestCache(t)
+			bc := NewBucketCache(instrumentationtest.New().ToProviderSettings(), memCache, cacheTTL, defaultFluxInterval)
+			step := qbtypes.Step{Duration: 10 * time.Millisecond}
+
+			// Each window stamps its values with its 1-based run number.
+			expectedValues := map[int64]float64{}
+			for i, window := range testCase.windows {
+				tsData := createTestTimeSeries("A", window.From, window.To, 1000)
+				for _, value := range tsData.Aggregations[0].Series[0].Values {
+					value.Value = float64(i + 1)
+					expectedValues[value.Timestamp] = value.Value
+				}
+				query := &mockQuery{fingerprint: testCase.name, startMs: window.From, endMs: window.To}
+				bc.Put(context.Background(), valuer.UUID{}, query, step, &qbtypes.Result{Type: qbtypes.RequestTypeTimeSeries, Value: tsData})
+			}
+
+			var data qbtypes.CachedData
+			require.NoError(t, memCache.Get(context.Background(), valuer.UUID{}, "v5:query:"+testCase.name, &data))
+			buckets := make([]qbtypes.TimeRange, 0, len(data.Buckets))
+			for _, bucket := range data.Buckets {
+				buckets = append(buckets, qbtypes.TimeRange{From: bucket.StartMs, To: bucket.EndMs})
+			}
+			assert.Equal(t, testCase.expectedBuckets, buckets)
+
+			// Every cached point carries the value of the last window that covered it.
+			span := &mockQuery{fingerprint: testCase.name, startMs: testCase.expectedBuckets[0].From, endMs: testCase.expectedBuckets[len(testCase.expectedBuckets)-1].To}
+			cached, _ := bc.GetMissRanges(context.Background(), valuer.UUID{}, span, step)
+			require.NotNil(t, cached)
+			tsData, ok := cached.Value.(*qbtypes.TimeSeriesData)
+			require.True(t, ok)
+			require.Len(t, tsData.Aggregations, 1)
+			require.Len(t, tsData.Aggregations[0].Series, 1)
+			cachedValues := map[int64]float64{}
+			for _, value := range tsData.Aggregations[0].Series[0].Values {
+				cachedValues[value.Timestamp] = value.Value
+			}
+			assert.Equal(t, expectedValues, cachedValues)
+		})
+	}
+}