Cache vertical shards in query frontend

pkg/queryfrontend/cache.go

@@ -33,11 +33,19 @@ func (t thanosCacheKeyGenerator) GenerateCacheKey(userID string, r queryrange.Re
 		i := 0
 		for ; i < len(t.resolutions) && t.resolutions[i] > tr.MaxSourceResolution; i++ {
 		}
-		return fmt.Sprintf("fe:%s:%s:%d:%d:%d", userID, tr.Query, tr.Step, currentInterval, i)
+		shardInfoKey := generateShardInfoKey(tr)
+		return fmt.Sprintf("fe:%s:%s:%d:%d:%d:%s", userID, tr.Query, tr.Step, currentInterval, i, shardInfoKey)
 	case *ThanosLabelsRequest:
 		return fmt.Sprintf("fe:%s:%s:%s:%d", userID, tr.Label, tr.Matchers, currentInterval)
 	case *ThanosSeriesRequest:
 		return fmt.Sprintf("fe:%s:%s:%d", userID, tr.Matchers, currentInterval)
 	}
 	return fmt.Sprintf("fe:%s:%s:%d:%d", userID, r.GetQuery(), r.GetStep(), currentInterval)
 }
+
+func generateShardInfoKey(r *ThanosQueryRangeRequest) string {
+	if r.ShardInfo == nil {
+		return "-"
+	}
+	return fmt.Sprintf("%d:%d", r.ShardInfo.TotalShards, r.ShardInfo.ShardIndex)
+}

pkg/queryfrontend/cache_test.go

@@ -37,7 +37,7 @@ func TestGenerateCacheKey(t *testing.T) {
 				Start: 0,
 				Step:  60 * seconds,
 			},
-			expected: "fe::up:60000:0:2",
+			expected: "fe::up:60000:0:2:-",
 		},
 		{
 			name: "10s step",
@@ -46,7 +46,7 @@ func TestGenerateCacheKey(t *testing.T) {
 				Start: 0,
 				Step:  10 * seconds,
 			},
-			expected: "fe::up:10000:0:2",
+			expected: "fe::up:10000:0:2:-",
 		},
 		{
 			name: "1m downsampling resolution",
@@ -56,7 +56,7 @@ func TestGenerateCacheKey(t *testing.T) {
 				Step:                10 * seconds,
 				MaxSourceResolution: 60 * seconds,
 			},
-			expected: "fe::up:10000:0:2",
+			expected: "fe::up:10000:0:2:-",
 		},
 		{
 			name: "5m downsampling resolution, different cache key",
@@ -66,7 +66,7 @@ func TestGenerateCacheKey(t *testing.T) {
 				Step:                10 * seconds,
 				MaxSourceResolution: 300 * seconds,
 			},
-			expected: "fe::up:10000:0:1",
+			expected: "fe::up:10000:0:1:-",
 		},
 		{
 			name: "1h downsampling resolution, different cache key",
@@ -76,7 +76,7 @@ func TestGenerateCacheKey(t *testing.T) {
 				Step:                10 * seconds,
 				MaxSourceResolution: hour,
 			},
-			expected: "fe::up:10000:0:0",
+			expected: "fe::up:10000:0:0:-",
 		},
 		{
 			name: "label names, no matcher",
@@ -134,7 +134,9 @@ func TestGenerateCacheKey(t *testing.T) {
 			expected: `fe::up:[[foo="bar"] [baz="qux"]]:0`,
 		},
 	} {
-		key := splitter.GenerateCacheKey("", tc.req)
-		testutil.Equals(t, tc.expected, key)
+		t.Run(tc.name, func(t *testing.T) {
+			key := splitter.GenerateCacheKey("", tc.req)
+			testutil.Equals(t, tc.expected, key)
+		})
 	}
 }

pkg/queryfrontend/roundtrip.go

@@ -191,6 +191,13 @@ func newQueryRangeTripperware(
 		)
 	}
 
+	if numShards > 0 {
+		queryRangeMiddleware = append(
+			queryRangeMiddleware,
+			PromQLShardingMiddleware(querysharding.NewQueryAnalyzer(), numShards, limits, codec, reg),
+		)
+	}
+
 	if config.ResultsCacheConfig != nil {
 		queryCacheMiddleware, _, err := queryrange.NewResultsCacheMiddleware(
 			logger,
@@ -222,13 +229,6 @@ func newQueryRangeTripperware(
 		)
 	}
 
-	if numShards > 0 {
-		queryRangeMiddleware = append(
-			queryRangeMiddleware,
-			PromQLShardingMiddleware(querysharding.NewQueryAnalyzer(), numShards, limits, codec, reg),
-		)
-	}
-
 	return func(next http.RoundTripper) http.RoundTripper {
 		rt := queryrange.NewRoundTripper(next, codec, forwardHeaders, queryRangeMiddleware...)
 		return queryrange.RoundTripFunc(func(r *http.Request) (*http.Response, error) {

pkg/queryfrontend/roundtrip_test.go

@@ -464,7 +464,6 @@ func TestRoundTripQueryRangeCacheMiddleware(t *testing.T) {
 		},
 	} {
 		if !t.Run(tc.name, func(t *testing.T) {
-
 			ctx := user.InjectOrgID(context.Background(), "1")
 			httpReq, err := NewThanosQueryRangeCodec(true).EncodeRequest(ctx, tc.req)
 			testutil.Ok(t, err)
@@ -479,6 +478,96 @@ func TestRoundTripQueryRangeCacheMiddleware(t *testing.T) {
 	}
 }
 
+func TestRoundTripQueryCacheWithShardingMiddleware(t *testing.T) {
+	testRequest := &ThanosQueryRangeRequest{
+		Path:    "/api/v1/query_range",
+		Start:   0,
+		End:     2 * hour,
+		Step:    10 * seconds,
+		Dedup:   true,
+		Query:   "sum by (pod) (memory_usage)",
+		Timeout: hour,
+	}
+
+	cacheConf := &queryrange.ResultsCacheConfig{
+		CacheConfig: cortexcache.Config{
+			EnableFifoCache: true,
+			Fifocache: cortexcache.FifoCacheConfig{
+				MaxSizeBytes: "1MiB",
+				MaxSizeItems: 1000,
+				Validity:     time.Hour,
+			},
+		},
+	}
+
+	tpw, err := NewTripperware(
+		Config{
+			NumShards: 2,
+			QueryRangeConfig: QueryRangeConfig{
+				Limits:                 defaultLimits,
+				ResultsCacheConfig:     cacheConf,
+				SplitQueriesByInterval: day,
+			},
+		}, nil, log.NewNopLogger(),
+	)
+	testutil.Ok(t, err)
+
+	rt, err := newFakeRoundTripper()
+	testutil.Ok(t, err)
+	defer rt.Close()
+	res, handler := promqlResultsWithFailures(3)
+	rt.setHandler(handler)
+
+	for _, tc := range []struct {
+		name     string
+		req      queryrange.Request
+		err      bool
+		expected int
+	}{
+		{
+			name:     "query with vertical sharding",
+			req:      testRequest,
+			err:      true,
+			expected: 2,
+		},
+		{
+			name:     "same query as before, both requests are executed",
+			req:      testRequest,
+			err:      true,
+			expected: 4,
+		},
+		{
+			name:     "same query as before, one request is executed",
+			req:      testRequest,
+			err:      false,
+			expected: 5,
+		},
+		{
+			name:     "same query as before again, no requests are executed",
+			req:      testRequest,
+			err:      false,
+			expected: 5,
+		},
+	} {
+		if !t.Run(tc.name, func(t *testing.T) {
+			ctx := user.InjectOrgID(context.Background(), "1")
+			httpReq, err := NewThanosQueryRangeCodec(true).EncodeRequest(ctx, tc.req)
+			testutil.Ok(t, err)
+
+			_, err = tpw(rt).RoundTrip(httpReq)
+			if tc.err {
+				testutil.NotOk(t, err)
+			} else {
+				testutil.Ok(t, err)
+			}
+
+			testutil.Equals(t, tc.expected, *res)
+		}) {
+			break
+		}
+	}
+}
+
 // TestRoundTripLabelsCacheMiddleware tests the cache middleware for labels requests.
 func TestRoundTripLabelsCacheMiddleware(t *testing.T) {
 	testRequest := &ThanosLabelsRequest{
@@ -730,6 +819,57 @@ func promqlResults(fail bool) (*int, http.Handler) {
 	})
 }
 
+// promqlResultsWithFailures is a mock handler used to test split and cache middleware.
+// it will return a failed response numFailures times.
+func promqlResultsWithFailures(numFailures int) (*int, http.Handler) {
+	count := 0
+	var lock sync.Mutex
+	q := queryrange.PrometheusResponse{
+		Status: "success",
+		Data: queryrange.PrometheusData{
+			ResultType: string(parser.ValueTypeMatrix),
+			Result: []queryrange.SampleStream{
+				{
+					Labels: []cortexpb.LabelAdapter{},
+					Samples: []cortexpb.Sample{
+						{Value: 0, TimestampMs: 0},
+						{Value: 1, TimestampMs: 1},
+					},
+				},
+			},
+		},
+	}
+
+	cond := sync.NewCond(&sync.Mutex{})
+	cond.L.Lock()
+	return &count, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		lock.Lock()
+		defer lock.Unlock()
+
+		// Set fail in the response code to test retry.
+		if numFailures > 0 {
+			numFailures--
+
+			// Wait for a successful request.
+			// Release the lock to allow other requests to execute.
+			if numFailures == 0 {
+				lock.Unlock()
+				cond.Wait()
+				<-time.After(500 * time.Millisecond)
+				lock.Lock()
+			}
+			w.WriteHeader(500)
+		}
+		if err := json.NewEncoder(w).Encode(q); err != nil {
+			panic(err)
+		}
+		if numFailures == 0 {
+			cond.Broadcast()
+		}
+		count++
+	})
+}
+
 // labelsResults is a mock handler used to test split and cache middleware for label names and label values requests.
 func labelsResults(fail bool) (*int, http.Handler) {
 	count := 0