fix avg_over_time panic with max series

Author: ie-pham

pkg/traceql/engine_metrics_average.go

@@ -273,14 +273,19 @@ func (b *averageOverTimeSeriesAggregator) Combine(in []*tempopb.TimeSeries) {
 			// This is a counter series, we can skip it
 			continue
 		}
+		countIndex, ok := countPosMapper[ts.PromLabels]
+		if !ok {
+			// The count series might have been truncated, skip this value
+			continue
+		}
 		for i, sample := range ts.Samples {
 			pos := IntervalOfMs(sample.TimestampMs, b.start, b.end, b.step)
 			if pos < 0 || pos >= len(b.weightedAverageSeries[ts.PromLabels].values) {
 				continue
 			}
 
 			incomingMean := sample.Value
-			incomingWeight := in[countPosMapper[ts.PromLabels]].Samples[i].Value
+			incomingWeight := in[countIndex].Samples[i].Value
 			existing.addWeigthedMean(pos, incomingMean, incomingWeight)
 			b.aggregateExemplars(ts, b.weightedAverageSeries[ts.PromLabels])
 		}

pkg/traceql/engine_metrics_test.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"math"
 	"math/rand/v2"
+	"strings"
 	"testing"
 	"time"
 
@@ -1022,6 +1023,79 @@ func TestObserveSeriesAverageOverTimeForSpanAttribute(t *testing.T) {
 	assert.Equal(t, 100.0, fooBar.Values[2])
 }
 
+func TestObserveSeriesAverageOverTimeForSpanAttributeWithTruncation(t *testing.T) {
+	req := &tempopb.QueryRangeRequest{
+		Start: uint64(1 * time.Second),
+		End:   uint64(3 * time.Second),
+		Step:  uint64(1 * time.Second),
+		Query: "{ } | avg_over_time(span.http.status_code) by (span.foo)",
+	}
+
+	// A variety of spans across times, durations, and series. All durations are powers of 2 for simplicity
+	in := []Span{
+		newMockSpan(nil).WithStartTime(uint64(1*time.Second)).WithSpanString("foo", "bar").WithSpanInt("http.status_code", 200),
+		newMockSpan(nil).WithStartTime(uint64(1*time.Second)).WithSpanString("foo", "bar").WithSpanInt("http.status_code", 300),
+		newMockSpan(nil).WithStartTime(uint64(1*time.Second)).WithSpanString("foo", "bar").WithSpanInt("http.status_code", 400),
+
+		newMockSpan(nil).WithStartTime(uint64(2*time.Second)).WithSpanString("foo", "bar").WithSpanInt("http.status_code", 200),
+		newMockSpan(nil).WithStartTime(uint64(2*time.Second)).WithSpanString("foo", "bar").WithSpanInt("http.status_code", 200),
+		newMockSpan(nil).WithStartTime(uint64(2*time.Second)).WithSpanString("foo", "bar").WithSpanInt("http.status_code", 100),
+		newMockSpan(nil).WithStartTime(uint64(2*time.Second)).WithSpanString("foo", "bar").WithSpanInt("http.status_code", 100),
+
+		newMockSpan(nil).WithStartTime(uint64(3*time.Second)).WithSpanString("foo", "baz").WithSpanInt("http.status_code", 200),
+		newMockSpan(nil).WithStartTime(uint64(3*time.Second)).WithSpanString("foo", "baz").WithSpanInt("http.status_code", 400),
+		newMockSpan(nil).WithStartTime(uint64(3*time.Second)).WithSpanString("foo", "baz").WithSpanInt("http.status_code", 500),
+	}
+
+	in2 := []Span{
+		newMockSpan(nil).WithStartTime(uint64(1*time.Second)).WithSpanString("foo", "bar").WithSpanInt("http.status_code", 100),
+		newMockSpan(nil).WithStartTime(uint64(1*time.Second)).WithSpanString("foo", "bar").WithSpanInt("http.status_code", 300),
+
+		newMockSpan(nil).WithStartTime(uint64(2*time.Second)).WithSpanString("foo", "bar").WithSpanInt("http.status_code", 400),
+		newMockSpan(nil).WithStartTime(uint64(2*time.Second)).WithSpanString("foo", "bar").WithSpanInt("http.status_code", 200),
+
+		newMockSpan(nil).WithStartTime(uint64(3*time.Second)).WithSpanString("foo", "baz").WithSpanInt("http.status_code", 100),
+		newMockSpan(nil).WithStartTime(uint64(3*time.Second)).WithSpanString("foo", "bar").WithSpanInt("http.status_code", 100),
+	}
+
+	e := NewEngine()
+	layer1A, _ := e.CompileMetricsQueryRange(req, 0, 0, false)
+	layer1B, _ := e.CompileMetricsQueryRange(req, 0, 0, false)
+	layer2A, _ := e.CompileMetricsQueryRangeNonRaw(req, AggregateModeSum)
+	layer2B, _ := e.CompileMetricsQueryRangeNonRaw(req, AggregateModeSum)
+	layer3, _ := e.CompileMetricsQueryRangeNonRaw(req, AggregateModeFinal)
+
+	for _, s := range in {
+		layer1A.metricsPipeline.observe(s)
+	}
+
+	layer2A.ObserveSeries(layer1A.Results().ToProto(req))
+
+	for _, s := range in2 {
+		layer1B.metricsPipeline.observe(s)
+	}
+
+	layer2B.ObserveSeries(layer1B.Results().ToProto(req))
+
+	layer3.ObserveSeries(layer2A.Results().ToProto(req))
+	layer2bResults := layer2B.Results().ToProto(req)
+	truncated2bResults := make([]*tempopb.TimeSeries, 0, len(layer2bResults)-1)
+	for _, ts := range layer2bResults {
+		if !strings.Contains(ts.PromLabels, internalLabelMetaType) {
+			// add all values series
+			truncated2bResults = append(truncated2bResults, ts)
+		} else {
+			// the panic appears when the count series with 3 samples is missing
+			if len(ts.Samples) != 3 {
+				truncated2bResults = append(truncated2bResults, ts)
+			}
+		}
+	}
+	assert.NotPanics(t, func() {
+		layer3.ObserveSeries(truncated2bResults)
+	}, "should not panic on truncation")
+}
+
 func TestMaxOverTimeForDuration(t *testing.T) {
 	req := &tempopb.QueryRangeRequest{
 		Start: uint64(1 * time.Second),