query-frontend: Refactor cache key generation

Signed-off-by: Michał Mazur <mmazur.box@gmail.com>
thanos-io · Nov 5, 2024 · d4725c3 · d4725c3
1 parent 8f9d0d4
commit d4725c3
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 26 deletions.
diff --git a/pkg/queryfrontend/cache.go b/pkg/queryfrontend/cache.go
@@ -24,26 +24,26 @@ func newThanosCacheKeyGenerator() thanosCacheKeyGenerator {
 // GenerateCacheKey generates a cache key based on the Request and interval.
 // TODO(yeya24): Add other request params as request key.
 func (t thanosCacheKeyGenerator) GenerateCacheKey(userID string, r queryrange.Request) string {
-	switch tr := r.(type) {
-	case *ThanosQueryRangeRequest:
-		i := 0
-		for ; i < len(t.resolutions) && t.resolutions[i] > tr.MaxSourceResolution; i++ {
-		}
-		shardInfoKey := generateShardInfoKey(tr)
-		splitInterval := tr.SplitInterval.Milliseconds()
-		currentInterval := r.GetStart() / splitInterval
-		return fmt.Sprintf("fe:%s:%s:%d:%d:%d:%d:%s:%d:%s", userID, tr.Query, tr.Step, splitInterval, currentInterval, i, shardInfoKey, tr.LookbackDelta, tr.Engine)
-	case *ThanosLabelsRequest:
-		splitInterval := tr.SplitInterval.Milliseconds()
-		currentInterval := r.GetStart() / splitInterval
-		return fmt.Sprintf("fe:%s:%s:%s:%d:%d", userID, tr.Label, tr.Matchers, splitInterval, currentInterval)
-	case *ThanosSeriesRequest:
-		splitInterval := tr.SplitInterval.Milliseconds()
+	if sr, ok := r.(SplitRequest); ok {
+		splitInterval := sr.GetSplitInterval().Milliseconds()
 		currentInterval := r.GetStart() / splitInterval
-		return fmt.Sprintf("fe:%s:%s:%d:%d", userID, tr.Matchers, splitInterval, currentInterval)
+
+		switch tr := r.(type) {
+		case *ThanosQueryRangeRequest:
+			i := 0
+			for ; i < len(t.resolutions) && t.resolutions[i] > tr.MaxSourceResolution; i++ {
+			}
+			shardInfoKey := generateShardInfoKey(tr)
+			return fmt.Sprintf("fe:%s:%s:%d:%d:%d:%d:%s:%d:%s", userID, tr.Query, tr.Step, splitInterval, currentInterval, i, shardInfoKey, tr.LookbackDelta, tr.Engine)
+		case *ThanosLabelsRequest:
+			return fmt.Sprintf("fe:%s:%s:%s:%d:%d", userID, tr.Label, tr.Matchers, splitInterval, currentInterval)
+		case *ThanosSeriesRequest:
+			return fmt.Sprintf("fe:%s:%s:%d:%d", userID, tr.Matchers, splitInterval, currentInterval)
+		}
 	}
 
-	return "request_type_not_supported"
+	// all possible request types are already covered
+	panic("request type not supported")
 }
 
 func generateShardInfoKey(r *ThanosQueryRangeRequest) string {

diff --git a/pkg/queryfrontend/cache_test.go b/pkg/queryfrontend/cache_test.go
@@ -22,15 +22,6 @@ func TestGenerateCacheKey(t *testing.T) {
 		req      queryrange.Request
 		expected string
 	}{
-		{
-			name: "non thanos req",
-			req: &queryrange.PrometheusRequest{
-				Query: "up",
-				Start: 0,
-				Step:  60 * seconds,
-			},
-			expected: "request_type_not_supported",
-		},
 		{
 			name: "non downsampling resolution specified",
 			req: &ThanosQueryRangeRequest{
@@ -164,3 +155,23 @@ func TestGenerateCacheKey(t *testing.T) {
 		})
 	}
 }
+
+func TestGenerateCacheKey_UnsupportedRequest(t *testing.T) {
+	splitter := newThanosCacheKeyGenerator()
+
+	req := &queryrange.PrometheusRequest{
+		Query: "up",
+		Start: 0,
+		Step:  60 * seconds,
+	}
+
+	defer func() {
+		if r := recover(); r == nil {
+			t.Fatal("expected panic")
+		} else {
+			testutil.Assert(t, r == "request type not supported", "unexpected panic: %v", r)
+		}
+	}()
+
+	splitter.GenerateCacheKey("", req)
+}