Add retries for instant query

Signed-off-by: 🌲 Harry 🌊 John 🏔 <[email protected]>

Author: harry671003

CHANGELOG.md

@@ -35,6 +35,7 @@
 * [FEATURE] Ruler: Support for filtering rules in the API. #5417
 * [FEATURE] Compactor: Add `-compactor.ring.tokens-file-path` to store generated tokens locally. #5432
 * [FEATURE] Query Frontend: Add `-frontend.retry-on-too-many-outstanding-requests` to re-enqueue 429 requests if there are multiple query-schedulers available. #5496
+* [FEATURE] Store Gateway: Add `-blocks-storage.bucket-store.max-inflight-requests`for store gateways to reject further requests upon reaching the limit. #5553
 * [ENHANCEMENT] Distributor/Ingester: Add span on push path #5319
 * [ENHANCEMENT] Support object storage backends for runtime configuration file. #5292
 * [ENHANCEMENT] Query Frontend: Reject subquery with too small step size. #5323
@@ -59,6 +60,8 @@
 * [ENHANCEMENT] Store Gateway: add metric `cortex_bucket_store_chunk_refetches_total` for number of chunk refetches. #5532
 * [ENHANCEMENT] BasicLifeCycler: allow final-sleep during shutdown #5517
 * [ENHANCEMENT] All: Handling CMK Access Denied errors. #5420 #5542
+* [ENHANCEMENT] Querier: Retry store gateway client connection closing gRPC error. #5558
+* [ENHANCEMENT] Query Frontend: Add retries for instant query. #5560
 * [BUGFIX] Ruler: Validate if rule group can be safely converted back to rule group yaml from protobuf message #5265
 * [BUGFIX] Querier: Convert gRPC `ResourceExhausted` status code from store gateway to 422 limit error. #5286
 * [BUGFIX] Alertmanager: Route web-ui requests to the alertmanager distributor when sharding is enabled. #5293

docs/blocks-storage/querier.md

@@ -499,6 +499,11 @@ blocks_storage:
     # CLI flag: -blocks-storage.bucket-store.max-concurrent
     [max_concurrent: <int> | default = 100]
 
+    # Max number of inflight queries to execute against the long-term storage.
+    # The limit is shared across all tenants. 0 to disable.
+    # CLI flag: -blocks-storage.bucket-store.max-inflight-requests
+    [max_inflight_requests: <int> | default = 0]
+
     # Maximum number of concurrent tenants synching blocks.
     # CLI flag: -blocks-storage.bucket-store.tenant-sync-concurrency
     [tenant_sync_concurrency: <int> | default = 10]

docs/blocks-storage/store-gateway.md

@@ -602,6 +602,11 @@ blocks_storage:
     # CLI flag: -blocks-storage.bucket-store.max-concurrent
     [max_concurrent: <int> | default = 100]
 
+    # Max number of inflight queries to execute against the long-term storage.
+    # The limit is shared across all tenants. 0 to disable.
+    # CLI flag: -blocks-storage.bucket-store.max-inflight-requests
+    [max_inflight_requests: <int> | default = 0]
+
     # Maximum number of concurrent tenants synching blocks.
     # CLI flag: -blocks-storage.bucket-store.tenant-sync-concurrency
     [tenant_sync_concurrency: <int> | default = 10]

docs/configuration/config-file-reference.md

@@ -1042,6 +1042,11 @@ bucket_store:
   # CLI flag: -blocks-storage.bucket-store.max-concurrent
   [max_concurrent: <int> | default = 100]
 
+  # Max number of inflight queries to execute against the long-term storage. The
+  # limit is shared across all tenants. 0 to disable.
+  # CLI flag: -blocks-storage.bucket-store.max-inflight-requests
+  [max_inflight_requests: <int> | default = 0]
+
   # Maximum number of concurrent tenants synching blocks.
   # CLI flag: -blocks-storage.bucket-store.tenant-sync-concurrency
   [tenant_sync_concurrency: <int> | default = 10]

pkg/cortex/modules.go

@@ -455,6 +455,7 @@ func (t *Cortex) initQueryFrontendTripperware() (serv services.Service, err erro
 	prometheusCodec := queryrange.NewPrometheusCodec(false)
 	// ShardedPrometheusCodec is same as PrometheusCodec but to be used on the sharded queries (it sum up the stats)
 	shardedPrometheusCodec := queryrange.NewPrometheusCodec(true)
+	retryMiddlewareMetrics := queryrange.NewRetryMiddlewareMetrics(prometheus.DefaultRegisterer)
 
 	queryRangeMiddlewares, cache, err := queryrange.Middlewares(
 		t.Cfg.QueryRange,
@@ -466,12 +467,13 @@ func (t *Cortex) initQueryFrontendTripperware() (serv services.Service, err erro
 		queryAnalyzer,
 		prometheusCodec,
 		shardedPrometheusCodec,
+		retryMiddlewareMetrics,
 	)
 	if err != nil {
 		return nil, err
 	}
 
-	instantQueryMiddlewares, err := instantquery.Middlewares(util_log.Logger, t.Overrides, queryAnalyzer)
+	instantQueryMiddlewares, err := instantquery.Middlewares(util_log.Logger, t.Overrides, retryMiddlewareMetrics, t.Cfg.QueryRange.MaxRetries, queryAnalyzer)
 	if err != nil {
 		return nil, err
 	}

pkg/querier/blocks_store_queryable.go

@@ -1116,6 +1116,12 @@ func isRetryableError(err error) bool {
 	switch status.Code(err) {
 	case codes.Unavailable:
 		return true
+	case codes.ResourceExhausted:
+		return errors.Is(err, storegateway.ErrTooManyInflightRequests)
+	// Client side connection closing, this error happens during store gateway deployment.
+	// https://github.com/grpc/grpc-go/blob/03172006f5d168fc646d87928d85cb9c4a480291/clientconn.go#L67
+	case codes.Canceled:
+		return strings.Contains(err.Error(), "grpc: the client connection is closing")
 	default:
 		return false
 	}

pkg/querier/blocks_store_queryable_test.go

@@ -34,6 +34,7 @@ import (
 	"google.golang.org/grpc/status"
 
 	"github.com/cortexproject/cortex/pkg/storage/tsdb/bucketindex"
+	"github.com/cortexproject/cortex/pkg/storegateway"
 	"github.com/cortexproject/cortex/pkg/storegateway/storegatewaypb"
 	"github.com/cortexproject/cortex/pkg/util"
 	"github.com/cortexproject/cortex/pkg/util/limiter"
@@ -638,6 +639,35 @@ func TestBlocksStoreQuerier_Select(t *testing.T) {
 				},
 			},
 		},
+		"multiple store-gateways has the block, but one of them fails to return due to clientconn closing": {
+			finderResult: bucketindex.Blocks{
+				{ID: block1},
+			},
+			storeSetResponses: []interface{}{
+				map[BlocksStoreClient][]ulid.ULID{
+					&storeGatewayClientMock{
+						remoteAddr:      "1.1.1.1",
+						mockedSeriesErr: status.Error(codes.Canceled, "grpc: the client connection is closing"),
+					}: {block1},
+				},
+				map[BlocksStoreClient][]ulid.ULID{
+					&storeGatewayClientMock{remoteAddr: "2.2.2.2", mockedSeriesResponses: []*storepb.SeriesResponse{
+						mockSeriesResponse(labels.Labels{metricNameLabel, series1Label}, minT, 2),
+						mockHintsResponse(block1),
+					}}: {block1},
+				},
+			},
+			limits:       &blocksStoreLimitsMock{},
+			queryLimiter: noOpQueryLimiter,
+			expectedSeries: []seriesResult{
+				{
+					lbls: labels.New(metricNameLabel, series1Label),
+					values: []valueResult{
+						{t: minT, v: 2},
+					},
+				},
+			},
+		},
 		"all store-gateways return PermissionDenied": {
 			finderResult: bucketindex.Blocks{
 				{ID: block1},
@@ -708,6 +738,56 @@ func TestBlocksStoreQuerier_Select(t *testing.T) {
 				},
 			},
 		},
+		"multiple store-gateways has the block, but one of them had too many inflight requests": {
+			finderResult: bucketindex.Blocks{
+				{ID: block1},
+			},
+			storeSetResponses: []interface{}{
+				map[BlocksStoreClient][]ulid.ULID{
+					&storeGatewayClientMock{
+						remoteAddr:      "1.1.1.1",
+						mockedSeriesErr: storegateway.ErrTooManyInflightRequests,
+					}: {block1},
+				},
+				map[BlocksStoreClient][]ulid.ULID{
+					&storeGatewayClientMock{remoteAddr: "2.2.2.2", mockedSeriesResponses: []*storepb.SeriesResponse{
+						mockSeriesResponse(labels.Labels{metricNameLabel, series1Label}, minT, 2),
+						mockHintsResponse(block1),
+					}}: {block1},
+				},
+			},
+			limits:       &blocksStoreLimitsMock{},
+			queryLimiter: noOpQueryLimiter,
+			expectedSeries: []seriesResult{
+				{
+					lbls: labels.New(metricNameLabel, series1Label),
+					values: []valueResult{
+						{t: minT, v: 2},
+					},
+				},
+			},
+		},
+		"store gateway returns resource exhausted error other than max inflight request": {
+			finderResult: bucketindex.Blocks{
+				{ID: block1},
+			},
+			storeSetResponses: []interface{}{
+				map[BlocksStoreClient][]ulid.ULID{
+					&storeGatewayClientMock{
+						remoteAddr:      "1.1.1.1",
+						mockedSeriesErr: status.Error(codes.ResourceExhausted, "some other resource"),
+					}: {block1},
+				},
+				map[BlocksStoreClient][]ulid.ULID{
+					&storeGatewayClientMock{remoteAddr: "2.2.2.2", mockedSeriesResponses: []*storepb.SeriesResponse{
+						mockSeriesResponse(labels.Labels{metricNameLabel, series1Label}, minT, 2),
+						mockHintsResponse(block1),
+					}}: {block1},
+				},
+			},
+			limits:      &blocksStoreLimitsMock{},
+			expectedErr: errors.Wrapf(status.Error(codes.ResourceExhausted, "some other resource"), "failed to fetch series from 1.1.1.1"),
+		},
 	}
 
 	for testName, testData := range tests {

pkg/querier/tripperware/instantquery/instant_query_middlewares.go

@@ -5,15 +5,21 @@ import (
 	"github.com/thanos-io/thanos/pkg/querysharding"
 
 	"github.com/cortexproject/cortex/pkg/querier/tripperware"
+	"github.com/cortexproject/cortex/pkg/querier/tripperware/queryrange"
 )
 
 func Middlewares(
 	log log.Logger,
 	limits tripperware.Limits,
+	retryMiddlewareMetrics *queryrange.RetryMiddlewareMetrics,
+	maxRetries int,
 	queryAnalyzer querysharding.Analyzer,
 ) ([]tripperware.Middleware, error) {
 	var m []tripperware.Middleware
 
+	if maxRetries > 0 {
+		m = append(m, queryrange.NewRetryMiddleware(log, maxRetries, retryMiddlewareMetrics))
+	}
 	m = append(m, tripperware.ShardByMiddleware(log, limits, InstantQueryCodec, queryAnalyzer))
 	return m, nil
 }

pkg/querier/tripperware/instantquery/instant_query_middlewares_test.go

@@ -0,0 +1,118 @@
+package instantquery
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/go-kit/log"
+	"github.com/stretchr/testify/require"
+	"github.com/thanos-io/thanos/pkg/querysharding"
+	"github.com/weaveworks/common/middleware"
+	"github.com/weaveworks/common/user"
+	"go.uber.org/atomic"
+
+	"github.com/cortexproject/cortex/pkg/querier/tripperware"
+)
+
+var (
+	query        = "/api/v1/query?time=1536716898&query=sum by (label) (up)&stats=all"
+	responseBody = `{"status":"success","data":{"resultType":"vector","result":[]}}`
+)
+
+func TestRoundTrip(t *testing.T) {
+	t.Parallel()
+	var try atomic.Int32
+	s := httptest.NewServer(
+		middleware.AuthenticateUser.Wrap(
+			http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				var err error
+				if try.Inc() > 2 {
+					_, err = w.Write([]byte(responseBody))
+				} else {
+					http.Error(w, `{"status":"error"}`, http.StatusInternalServerError)
+				}
+				if err != nil {
+					t.Fatal(err)
+				}
+			}),
+		),
+	)
+	defer s.Close()
+
+	u, err := url.Parse(s.URL)
+	require.NoError(t, err)
+
+	downstream := singleHostRoundTripper{
+		host: u.Host,
+		next: http.DefaultTransport,
+	}
+	limits := tripperware.MockLimits{
+		ShardSize: 2,
+	}
+	qa := querysharding.NewQueryAnalyzer()
+	instantQueryMiddlewares, err := Middlewares(
+		log.NewNopLogger(),
+		limits,
+		nil,
+		3,
+		qa)
+	require.NoError(t, err)
+
+	tw := tripperware.NewQueryTripperware(
+		log.NewNopLogger(),
+		nil,
+		nil,
+		nil,
+		instantQueryMiddlewares,
+		nil,
+		InstantQueryCodec,
+		limits,
+		qa,
+		time.Minute,
+	)
+
+	for i, tc := range []struct {
+		path, expectedBody string
+	}{
+		{query, responseBody},
+	} {
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			//parallel testing causes data race
+			req, err := http.NewRequest("GET", tc.path, http.NoBody)
+			require.NoError(t, err)
+
+			// query-frontend doesn't actually authenticate requests, we rely on
+			// the queriers to do this.  Hence we ensure the request doesn't have a
+			// org ID in the ctx, but does have the header.
+			ctx := user.InjectOrgID(context.Background(), "1")
+			req = req.WithContext(ctx)
+			err = user.InjectOrgIDIntoHTTPRequest(ctx, req)
+			require.NoError(t, err)
+
+			resp, err := tw(downstream).RoundTrip(req)
+			require.NoError(t, err)
+			require.Equal(t, 200, resp.StatusCode)
+
+			bs, err := io.ReadAll(resp.Body)
+			require.NoError(t, err)
+			require.Equal(t, tc.expectedBody, string(bs))
+		})
+	}
+}
+
+type singleHostRoundTripper struct {
+	host string
+	next http.RoundTripper
+}
+
+func (s singleHostRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
+	r.URL.Scheme = "http"
+	r.URL.Host = s.host
+	return s.next.RoundTrip(r)
+}

pkg/querier/tripperware/queryrange/limits_test.go

@@ -75,7 +75,7 @@ func TestLimitsMiddleware_MaxQueryLookback(t *testing.T) {
 				End:   util.TimeToMillis(testData.reqEndTime),
 			}
 
-			limits := mockLimits{maxQueryLookback: testData.maxQueryLookback}
+			limits := tripperware.MockLimits{QueryLookback: testData.maxQueryLookback}
 			middleware := NewLimitsMiddleware(limits)
 
 			innerRes := NewEmptyPrometheusResponse()
@@ -163,7 +163,7 @@ func TestLimitsMiddleware_MaxQueryLength(t *testing.T) {
 				End:   util.TimeToMillis(testData.reqEndTime),
 			}
 
-			limits := mockLimits{maxQueryLength: testData.maxQueryLength}
+			limits := tripperware.MockLimits{QueryLength: testData.maxQueryLength}
 			middleware := NewLimitsMiddleware(limits)
 
 			innerRes := NewEmptyPrometheusResponse()
@@ -193,32 +193,6 @@ func TestLimitsMiddleware_MaxQueryLength(t *testing.T) {
 	}
 }
 
-type mockLimits struct {
-	maxQueryLookback  time.Duration
-	maxQueryLength    time.Duration
-	maxCacheFreshness time.Duration
-}
-
-func (m mockLimits) MaxQueryLookback(string) time.Duration {
-	return m.maxQueryLookback
-}
-
-func (m mockLimits) MaxQueryLength(string) time.Duration {
-	return m.maxQueryLength
-}
-
-func (mockLimits) MaxQueryParallelism(string) int {
-	return 14 // Flag default.
-}
-
-func (m mockLimits) MaxCacheFreshness(string) time.Duration {
-	return m.maxCacheFreshness
-}
-
-func (m mockLimits) QueryVerticalShardSize(userID string) int {
-	return 0
-}
-
 type mockHandler struct {
 	mock.Mock
 }