From 16e7e9f654b734ff3d3c5944001db5945079d5e7 Mon Sep 17 00:00:00 2001
From: pawarpranav83 <pawarpranav@gmail.com>
Date: Tue, 26 Dec 2023 23:25:18 +0530
Subject: [PATCH] Added negative offset check for caching queries

Signed-off-by: pawarpranav83 <pawarpranav@gmail.com>
---
 .../querier/queryrange/results_cache.go       | 42 +++++++++++++++++++
 .../querier/queryrange/results_cache_test.go  | 39 +++++++++++++++++
 2 files changed, 81 insertions(+)

diff --git a/internal/cortex/querier/queryrange/results_cache.go b/internal/cortex/querier/queryrange/results_cache.go
index d29f5c33050..2534a0e11be 100644
--- a/internal/cortex/querier/queryrange/results_cache.go
+++ b/internal/cortex/querier/queryrange/results_cache.go
@@ -282,6 +282,9 @@ func (s resultsCache) shouldCacheResponse(ctx context.Context, req Request, r Re
 	if !s.isAtModifierCachable(req, maxCacheTime) {
 		return false
 	}
+	if !s.isOffsetCachable(req) {
+		return false
+	}
 
 	if s.cacheGenNumberLoader == nil {
 		return true
@@ -357,6 +360,45 @@ func (s resultsCache) isAtModifierCachable(r Request, maxCacheTime int64) bool {
 	return atModCachable
 }
 
+var negativeOffset = errors.New("negative offset")
+
+func (s resultsCache) isOffsetCachable(r Request) bool {
+	query := r.GetQuery()
+	if !strings.Contains(query, "offset") {
+		return true
+	}
+	expr, err := parser.ParseExpr(query)
+	if err != nil {
+		level.Warn(s.logger).Log("msg", "failed to parse query, considering @ modifier as not cachable", "query", query, "err", err)
+		return false
+	}
+
+	offsetCachable := true
+	parser.Inspect(expr, func(n parser.Node, _ []parser.Node) error {
+		switch e := n.(type) {
+		case *parser.VectorSelector:
+			if e.OriginalOffset < 0 {
+				offsetCachable = false
+				return negativeOffset
+			}
+		case *parser.MatrixSelector:
+			offset := e.VectorSelector.(*parser.VectorSelector).OriginalOffset
+			if offset < 0 {
+				offsetCachable = false
+				return negativeOffset
+			}
+		case *parser.SubqueryExpr:
+			if e.OriginalOffset < 0 {
+				offsetCachable = false
+				return negativeOffset
+			}
+		}
+		return nil
+	})
+
+	return offsetCachable
+}
+
 func getHeaderValuesWithName(r Response, headerName string) (headerValues []string) {
 	for _, hv := range r.GetHeaders() {
 		if hv.GetName() != headerName {
diff --git a/internal/cortex/querier/queryrange/results_cache_test.go b/internal/cortex/querier/queryrange/results_cache_test.go
index cb7d04fc95b..cb12d5cb231 100644
--- a/internal/cortex/querier/queryrange/results_cache_test.go
+++ b/internal/cortex/querier/queryrange/results_cache_test.go
@@ -532,6 +532,45 @@ func TestShouldCache(t *testing.T) {
 			input:    Response(&PrometheusResponse{}),
 			expected: false,
 		},
+		// offset on vector selectors.
+		{
+			name:     "positive offset on vector selector",
+			request:  &PrometheusRequest{Query: "metric offset 10ms", End: 125000},
+			input:    Response(&PrometheusResponse{}),
+			expected: true,
+		},
+		{
+			name:     "negative offset on vector selector",
+			request:  &PrometheusRequest{Query: "metric offset -10ms", End: 125000},
+			input:    Response(&PrometheusResponse{}),
+			expected: false,
+		},
+		// offset on matrix selectors.
+		{
+			name:     "positive offset on matrix selector",
+			request:  &PrometheusRequest{Query: "rate(metric[5m] offset 10ms)", End: 125000},
+			input:    Response(&PrometheusResponse{}),
+			expected: true,
+		},
+		{
+			name:     "negative offset on matrix selector",
+			request:  &PrometheusRequest{Query: "rate(metric[5m] offset -10ms)", End: 125000},
+			input:    Response(&PrometheusResponse{}),
+			expected: false,
+		},
+		// offset on subqueries.
+		{
+			name:     "positive offset on subqueries",
+			request:  &PrometheusRequest{Query: "sum_over_time(rate(metric[1m])[10m:1m] offset 10ms)", End: 125000},
+			input:    Response(&PrometheusResponse{}),
+			expected: true,
+		},
+		{
+			name:     "negative offset on subqueries",
+			request:  &PrometheusRequest{Query: "sum_over_time(rate(metric[1m])[10m:1m] offset -10ms)", End: 125000},
+			input:    Response(&PrometheusResponse{}),
+			expected: false,
+		},
 	} {
 		{
 			t.Run(tc.name, func(t *testing.T) {