Resource based throttling bugfix 2 (#7013)

* Make error from RBT query rejection retryable in querier

Signed-off-by: Justin Jung <[email protected]>

* Add tests

Signed-off-by: Justin Jung <[email protected]>

* Lint

Signed-off-by: Justin Jung <[email protected]>

* Lint import

Signed-off-by: Justin Jung <[email protected]>

* Added test in error translate queryable

Signed-off-by: Justin Jung <[email protected]>

---------

Signed-off-by: Justin Jung <[email protected]>

Author: justinjung04

pkg/frontend/transport/handler_test.go

@@ -388,7 +388,7 @@ func TestHandler_ServeHTTP(t *testing.T) {
 			cfg:             HandlerConfig{QueryStatsEnabled: true},
 			expectedMetrics: 6,
 			roundTripperFunc: roundTripperFunc(func(req *http.Request) (*http.Response, error) {
-				resourceLimitReachedErr := &limiter.ResourceLimitReachedError{}
+				resourceLimitReachedErr := limiter.ErrResourceLimitReached
 				return &http.Response{
 					StatusCode: http.StatusServiceUnavailable,
 					Body:       io.NopCloser(strings.NewReader(resourceLimitReachedErr.Error())),

pkg/ingester/ingester.go

@@ -2265,7 +2265,7 @@ func (i *Ingester) trackInflightQueryRequest() (func(), error) {
 	if i.resourceBasedLimiter != nil {
 		if err := i.resourceBasedLimiter.AcceptNewRequest(); err != nil {
 			level.Warn(i.logger).Log("msg", "failed to accept request", "err", err)
-			return nil, httpgrpc.Errorf(http.StatusServiceUnavailable, "failed to query: %s", limiter.ErrResourceLimitReachedStr)
+			return nil, limiter.ErrResourceLimitReached
 		}
 	}
 

pkg/ingester/ingester_test.go

@@ -3230,8 +3230,9 @@ func Test_Ingester_Query_ResourceThresholdBreached(t *testing.T) {
 	s := &mockQueryStreamServer{ctx: ctx}
 	err = i.QueryStream(rreq, s)
 	require.Error(t, err)
-	exhaustedErr := limiter.ResourceLimitReachedError{}
-	require.ErrorContains(t, err, exhaustedErr.Error())
+
+	// Expected error from isRetryableError in blocks_store_queryable.go
+	require.ErrorIs(t, err, limiter.ErrResourceLimitReached)
 }
 
 func TestIngester_LabelValues_ShouldNotCreateTSDBIfDoesNotExists(t *testing.T) {

pkg/querier/blocks_store_queryable.go

@@ -1204,17 +1204,11 @@ func countSamplesAndChunks(series ...*storepb.Series) (samplesCount, chunksCount
 
 // only retry connection issues
 func isRetryableError(err error) bool {
-	// retry upon resource exhaustion error from resource monitor
-	var resourceExhaustedErr *limiter.ResourceLimitReachedError
-	if errors.As(err, &resourceExhaustedErr) {
-		return true
-	}
-
 	switch status.Code(err) {
 	case codes.Unavailable:
 		return true
 	case codes.ResourceExhausted:
-		return errors.Is(err, storegateway.ErrTooManyInflightRequests)
+		return errors.Is(err, storegateway.ErrTooManyInflightRequests) || errors.Is(err, limiter.ErrResourceLimitReached)
 	// 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:

pkg/querier/blocks_store_queryable_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"net/http"
 	"sort"
 	"strings"
 	"testing"
@@ -31,6 +32,7 @@ import (
 	"github.com/thanos-io/thanos/pkg/store/hintspb"
 	"github.com/thanos-io/thanos/pkg/store/labelpb"
 	"github.com/thanos-io/thanos/pkg/store/storepb"
+	"github.com/weaveworks/common/httpgrpc"
 	"github.com/weaveworks/common/user"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
@@ -44,7 +46,6 @@ import (
 	"github.com/cortexproject/cortex/pkg/util"
 	"github.com/cortexproject/cortex/pkg/util/limiter"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
-	"github.com/cortexproject/cortex/pkg/util/resource"
 	"github.com/cortexproject/cortex/pkg/util/services"
 	"github.com/cortexproject/cortex/pkg/util/validation"
 )
@@ -1527,7 +1528,7 @@ func TestBlocksStoreQuerier_Select(t *testing.T) {
 				map[BlocksStoreClient][]ulid.ULID{
 					&storeGatewayClientMock{
 						remoteAddr:      "1.1.1.1",
-						mockedSeriesErr: &limiter.ResourceLimitReachedError{},
+						mockedSeriesErr: limiter.ErrResourceLimitReached,
 					}: {block1},
 				},
 				map[BlocksStoreClient][]ulid.ULID{
@@ -2499,20 +2500,18 @@ func TestBlocksStoreQuerier_PromQLExecution(t *testing.T) {
 	}
 }
 
-func TestBlocksStoreQuerier_ShouldRetryResourceBasedThrottlingError(t *testing.T) {
-	limits := map[resource.Type]float64{
-		resource.CPU:  0.5,
-		resource.Heap: 0.5,
-	}
-
-	resourceBasedLimiter, err := limiter.NewResourceBasedLimiter(&limiter.MockMonitor{
-		CpuUtilization:  0.7,
-		HeapUtilization: 0.7,
-	}, limits, prometheus.DefaultRegisterer, "ingester")
-	require.NoError(t, err)
-
-	err = resourceBasedLimiter.AcceptNewRequest()
-	require.True(t, isRetryableError(err))
+func TestBlocksStoreQuerier_isRetryableError(t *testing.T) {
+	require.True(t, isRetryableError(status.Error(codes.Unavailable, "")))
+	require.True(t, isRetryableError(storegateway.ErrTooManyInflightRequests))
+	require.True(t, isRetryableError(limiter.ErrResourceLimitReached))
+	require.True(t, isRetryableError(status.Error(codes.Canceled, "grpc: the client connection is closing")))
+	require.True(t, isRetryableError(errors.New("pool exhausted")))
+
+	require.False(t, isRetryableError(status.Error(codes.ResourceExhausted, "some other error")))
+	require.False(t, isRetryableError(status.Error(codes.Canceled, "some other error")))
+	require.False(t, isRetryableError(errors.New("some other error")))
+	require.False(t, isRetryableError(fmt.Errorf("some other error")))
+	require.False(t, isRetryableError(httpgrpc.Errorf(http.StatusServiceUnavailable, "some other error")))
 }
 
 type blocksStoreSetMock struct {

pkg/querier/error_translate_queryable_test.go

@@ -24,6 +24,8 @@ import (
 	"github.com/weaveworks/common/httpgrpc"
 	"github.com/weaveworks/common/user"
 
+	"github.com/cortexproject/cortex/pkg/storegateway"
+	"github.com/cortexproject/cortex/pkg/util/limiter"
 	"github.com/cortexproject/cortex/pkg/util/validation"
 )
 
@@ -113,6 +115,16 @@ func TestApiStatusCodes(t *testing.T) {
 			expectedString: "test string",
 			expectedCode:   422,
 		},
+		{
+			err:            storegateway.ErrTooManyInflightRequests,
+			expectedString: "too many inflight requests in store gateway",
+			expectedCode:   500,
+		},
+		{
+			err:            limiter.ErrResourceLimitReached,
+			expectedString: limiter.ErrResourceLimitReachedStr,
+			expectedCode:   500,
+		},
 	} {
 		for k, q := range map[string]storage.SampleAndChunkQueryable{
 			"error from queryable": errorTestQueryable{err: tc.err},

pkg/storegateway/gateway.go

@@ -16,7 +16,6 @@ import (
 	"github.com/thanos-io/objstore"
 	"github.com/thanos-io/thanos/pkg/extprom"
 	"github.com/thanos-io/thanos/pkg/store/storepb"
-	"github.com/weaveworks/common/httpgrpc"
 	"github.com/weaveworks/common/logging"
 
 	"github.com/cortexproject/cortex/pkg/configs"
@@ -437,7 +436,7 @@ func (g *StoreGateway) checkResourceUtilization() error {
 
 	if err := g.resourceBasedLimiter.AcceptNewRequest(); err != nil {
 		level.Warn(g.logger).Log("msg", "failed to accept request", "err", err)
-		return httpgrpc.Errorf(http.StatusServiceUnavailable, "failed to query: %s", util_limiter.ErrResourceLimitReachedStr)
+		return util_limiter.ErrResourceLimitReached
 	}
 
 	return nil

pkg/storegateway/gateway_test.go

@@ -1239,8 +1239,9 @@ func TestStoreGateway_SeriesThrottledByResourceMonitor(t *testing.T) {
 	srv := newBucketStoreSeriesServer(setUserIDToGRPCContext(ctx, userID))
 	err = g.Series(req, srv)
 	require.Error(t, err)
-	exhaustedErr := util_limiter.ResourceLimitReachedError{}
-	require.ErrorContains(t, err, exhaustedErr.Error())
+
+	// Expected error from isRetryableError in blocks_store_queryable.go
+	require.ErrorIs(t, err, util_limiter.ErrResourceLimitReached)
 }
 
 func mockGatewayConfig() Config {

pkg/util/limiter/resource_based_limiter.go

@@ -5,17 +5,15 @@ import (
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 
 	"github.com/cortexproject/cortex/pkg/util/resource"
 )
 
 const ErrResourceLimitReachedStr = "resource limit reached"
 
-type ResourceLimitReachedError struct{}
-
-func (e *ResourceLimitReachedError) Error() string {
-	return ErrResourceLimitReachedStr
-}
+var ErrResourceLimitReached = status.Error(codes.ResourceExhausted, ErrResourceLimitReachedStr)
 
 type ResourceBasedLimiter struct {
 	resourceMonitor    resource.IMonitor
@@ -64,7 +62,7 @@ func (l *ResourceBasedLimiter) AcceptNewRequest() error {
 
 		if utilization >= limit {
 			l.limitBreachedCount.WithLabelValues(string(resType)).Inc()
-			return fmt.Errorf("%s utilization limit reached (limit: %.3f, utilization: %.3f): %w", resType, limit, utilization, &ResourceLimitReachedError{})
+			return fmt.Errorf("%s utilization limit reached (limit: %.3f, utilization: %.3f): %s", resType, limit, utilization, ErrResourceLimitReachedStr)
 		}
 	}
 

pkg/util/limiter/resource_based_limiter_test.go

@@ -5,6 +5,8 @@ import (
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 
 	"github.com/cortexproject/cortex/pkg/util/resource"
 )
@@ -24,3 +26,8 @@ func Test_ResourceBasedLimiter(t *testing.T) {
 	err = limiter.AcceptNewRequest()
 	require.NoError(t, err)
 }
+
+func Test_ResourceBasedLimiter_ErrResourceLimitReached(t *testing.T) {
+	// Expected error code from isRetryableError in blocks_store_queryable.go
+	require.Equal(t, codes.ResourceExhausted, status.Code(ErrResourceLimitReached))
+}