From e53a743487f15b30a72daf4e41f89240a1a5f517 Mon Sep 17 00:00:00 2001
From: Martin Disibio <mdisibio@gmail.com>
Date: Thu, 14 Nov 2024 11:42:13 -0500
Subject: [PATCH] Disconnect job time range filtering from step, so that
 results in split backend/recent range is accurate

---
 pkg/traceql/engine_metrics.go      | 46 ++++++++++++++++++++++--------
 pkg/traceql/engine_metrics_test.go | 17 ++++-------
 2 files changed, 40 insertions(+), 23 deletions(-)

diff --git a/pkg/traceql/engine_metrics.go b/pkg/traceql/engine_metrics.go
index 3aba0a74f6c..f997aac8112 100644
--- a/pkg/traceql/engine_metrics.go
+++ b/pkg/traceql/engine_metrics.go
@@ -50,6 +50,9 @@ func DefaultQueryRangeStep(start, end uint64) uint64 {
 
 // IntervalCount is the number of intervals in the range with step.
 func IntervalCount(start, end, step uint64) int {
+	start = alignStart(start, step)
+	end = alignEnd(end, step)
+
 	intervals := (end - start) / step
 	intervals++
 	return int(intervals)
@@ -57,11 +60,15 @@ func IntervalCount(start, end, step uint64) int {
 
 // TimestampOf the given interval with the start and step.
 func TimestampOf(interval, start, step uint64) uint64 {
+	start = alignStart(start, step)
 	return start + interval*step
 }
 
 // IntervalOf the given timestamp within the range and step.
 func IntervalOf(ts, start, end, step uint64) int {
+	start = alignStart(start, step)
+	end = alignEnd(end, step) + step
+
 	if ts < start || ts > end || end == start || step == 0 {
 		// Invalid
 		return -1
@@ -89,10 +96,6 @@ func TrimToOverlap(start1, end1, step, start2, end2 uint64) (uint64, uint64, uin
 	if wasInstant {
 		// Alter step to maintain instant nature
 		step = end1 - start1
-	} else {
-		// Realign after trimming
-		start1 = (start1 / step) * step
-		end1 = (end1/step)*step + step
 	}
 
 	return start1, end1, step
@@ -110,9 +113,6 @@ func TrimToBefore(req *tempopb.QueryRangeRequest, before time.Time) {
 	if wasInstant {
 		// Maintain instant nature of the request
 		req.Step = req.End - req.Start
-	} else {
-		// Realign after trimming
-		AlignRequest(req)
 	}
 }
 
@@ -128,9 +128,6 @@ func TrimToAfter(req *tempopb.QueryRangeRequest, before time.Time) {
 	if wasInstant {
 		// Maintain instant nature of the request
 		req.Step = req.End - req.Start
-	} else {
-		// Realign after trimming
-		AlignRequest(req)
 	}
 }
 
@@ -148,8 +145,30 @@ func AlignRequest(req *tempopb.QueryRangeRequest) {
 	}
 
 	// It doesn't really matter but the request fields are expected to be in nanoseconds.
-	req.Start = req.Start / req.Step * req.Step
-	req.End = req.End / req.Step * req.Step
+	req.Start = alignStart(req.Start, req.Step)
+	req.End = alignEnd(req.End, req.Step)
+}
+
+// Start time is rounded down to next step
+func alignStart(start, step uint64) uint64 {
+	if step == 0 {
+		return 0
+	}
+	return start - start%step
+}
+
+// End time is rounded up to next step
+func alignEnd(end, step uint64) uint64 {
+	if step == 0 {
+		return 0
+	}
+
+	mod := end % step
+	if mod == 0 {
+		return end
+	}
+
+	return end + (step - mod)
 }
 
 type Label struct {
@@ -238,6 +257,9 @@ func (set SeriesSet) ToProtoDiff(req *tempopb.QueryRangeRequest, rangeForLabels
 			continue
 		}
 
+		start = alignStart(start, req.Step)
+		end = alignEnd(end, req.Step)
+
 		intervals := IntervalCount(start, end, req.Step)
 		samples := make([]tempopb.Sample, 0, intervals)
 		for i, value := range s.Values {
diff --git a/pkg/traceql/engine_metrics_test.go b/pkg/traceql/engine_metrics_test.go
index 2fc0e356896..c5a64148196 100644
--- a/pkg/traceql/engine_metrics_test.go
+++ b/pkg/traceql/engine_metrics_test.go
@@ -41,7 +41,7 @@ func TestStepRangeToIntervals(t *testing.T) {
 			start:    0,
 			end:      10,
 			step:     3,
-			expected: 4, // 0, 3, 6, 9
+			expected: 5, // 0, 3, 6, 9, 12
 		},
 	}
 
@@ -60,9 +60,9 @@ func TestTimestampOf(t *testing.T) {
 		},
 		{
 			interval: 2,
-			start:    10,
+			start:    10, // aligned to 9
 			step:     3,
-			expected: 16,
+			expected: 15, // 9, 12, 15 <-- intervals
 		},
 	}
 
@@ -104,21 +104,16 @@ func TestTrimToOverlap(t *testing.T) {
 		expectedStep               time.Duration
 	}{
 		{
-			// Inner range of 33 to 38
-			// gets rounded at 5m intervals to 30 to 40
+			// Fully overlapping
 			"2024-01-01 01:00:00", "2024-01-01 02:00:00", 5 * time.Minute,
 			"2024-01-01 01:33:00", "2024-01-01 01:38:00",
-			"2024-01-01 01:30:00", "2024-01-01 01:40:00", 5 * time.Minute,
+			"2024-01-01 01:33:00", "2024-01-01 01:38:00", 5 * time.Minute,
 		},
 		{
 			// Partially Overlapping
-			// Overlap between 1:01-2:01 and 1:31-2:31
-			// in 5m intervals is only 1:30-2:05
-			// Start is pushed back
-			// and end is pushed out
 			"2024-01-01 01:01:00", "2024-01-01 02:01:00", 5 * time.Minute,
 			"2024-01-01 01:31:00", "2024-01-01 02:31:00",
-			"2024-01-01 01:30:00", "2024-01-01 02:05:00", 5 * time.Minute,
+			"2024-01-01 01:31:00", "2024-01-01 02:01:00", 5 * time.Minute,
 		},
 		{
 			// Instant query