resource based throttling: reject only adhoc queries

CHANGELOG.md

@@ -69,6 +69,7 @@
 * [ENHANCEMENT] Query Frontend: Add a `format_query` label value to the `op` label at `cortex_query_frontend_queries_total` metric. #6925
 * [ENHANCEMENT] API: add request ID injection to context to enable tracking requests across downstream services. #6895
 * [ENHANCEMENT] gRPC: Add gRPC Channelz monitoring. #6950
+* [ENHANCEMENT] Add source metadata to requests(api vs ruler) and used in resource based throttling to only reject adhoc queries. #6947
 * [BUGFIX] Ingester: Avoid error or early throttling when READONLY ingesters are present in the ring #6517
 * [BUGFIX] Ingester: Fix labelset data race condition. #6573
 * [BUGFIX] Compactor: Cleaner should not put deletion marker for blocks with no-compact marker. #6576

pkg/api/middlewares.go

@@ -37,6 +37,7 @@ func (h HTTPHeaderMiddleware) injectRequestContext(r *http.Request) *http.Reques
 		reqId = uuid.NewString()
 	}
 	requestContextMap[requestmeta.RequestIdKey] = reqId
+	requestContextMap[requestmeta.RequestSourceKey] = requestmeta.SourceApi
 
 	ctx := requestmeta.ContextWithRequestMetadataMap(r.Context(), requestContextMap)
 	return r.WithContext(ctx)

pkg/ingester/ingester.go

@@ -63,6 +63,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/util/limiter"
 	logutil "github.com/cortexproject/cortex/pkg/util/log"
 	util_math "github.com/cortexproject/cortex/pkg/util/math"
+	"github.com/cortexproject/cortex/pkg/util/requestmeta"
 	"github.com/cortexproject/cortex/pkg/util/resource"
 	"github.com/cortexproject/cortex/pkg/util/services"
 	"github.com/cortexproject/cortex/pkg/util/spanlogger"
@@ -1696,7 +1697,7 @@ func (i *Ingester) QueryExemplars(ctx context.Context, req *client.ExemplarQuery
 	}
 
 	// We will report *this* request in the error too.
-	c, err := i.trackInflightQueryRequest()
+	c, err := i.trackInflightQueryRequest(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -1804,7 +1805,7 @@ func (i *Ingester) labelsValuesCommon(ctx context.Context, req *client.LabelValu
 		q.Close()
 	}
 
-	c, err := i.trackInflightQueryRequest()
+	c, err := i.trackInflightQueryRequest(ctx)
 	if err != nil {
 		return nil, cleanup, err
 	}
@@ -1901,7 +1902,7 @@ func (i *Ingester) labelNamesCommon(ctx context.Context, req *client.LabelNamesR
 		q.Close()
 	}
 
-	c, err := i.trackInflightQueryRequest()
+	c, err := i.trackInflightQueryRequest(ctx)
 	if err != nil {
 		return nil, cleanup, err
 	}
@@ -2252,7 +2253,7 @@ func (i *Ingester) QueryStream(req *client.QueryRequest, stream client.Ingester_
 	return nil
 }
 
-func (i *Ingester) trackInflightQueryRequest() (func(), error) {
+func (i *Ingester) trackInflightQueryRequest(ctx context.Context) (func(), error) {
 	gl := i.getInstanceLimits()
 	if gl != nil && gl.MaxInflightQueryRequests > 0 {
 		if i.inflightQueryRequests.Load() >= gl.MaxInflightQueryRequests {
@@ -2262,7 +2263,7 @@ func (i *Ingester) trackInflightQueryRequest() (func(), error) {
 
 	i.maxInflightQueryRequests.Track(i.inflightQueryRequests.Inc())
 
-	if i.resourceBasedLimiter != nil {
+	if i.resourceBasedLimiter != nil && !requestmeta.RequestFromRuler(ctx) {
 		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)
@@ -2282,7 +2283,7 @@ func (i *Ingester) queryStreamChunks(ctx context.Context, db *userTSDB, from, th
 	}
 	defer q.Close()
 
-	c, err := i.trackInflightQueryRequest()
+	c, err := i.trackInflightQueryRequest(ctx)
 	if err != nil {
 		return 0, 0, 0, 0, err
 	}

pkg/ingester/ingester_test.go

@@ -61,6 +61,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/util"
 	"github.com/cortexproject/cortex/pkg/util/chunkcompat"
 	"github.com/cortexproject/cortex/pkg/util/limiter"
+	"github.com/cortexproject/cortex/pkg/util/requestmeta"
 	"github.com/cortexproject/cortex/pkg/util/resource"
 	"github.com/cortexproject/cortex/pkg/util/services"
 	"github.com/cortexproject/cortex/pkg/util/test"
@@ -3227,11 +3228,18 @@ func Test_Ingester_Query_ResourceThresholdBreached(t *testing.T) {
 	}
 
 	rreq := &client.QueryRequest{}
+	ctx = requestmeta.ContextWithRequestSource(ctx, requestmeta.SourceApi)
 	s := &mockQueryStreamServer{ctx: ctx}
 	err = i.QueryStream(rreq, s)
 	require.Error(t, err)
 	exhaustedErr := limiter.ResourceLimitReachedError{}
 	require.ErrorContains(t, err, exhaustedErr.Error())
+
+	// we shouldn't reject queries from ruler
+	ctx = requestmeta.ContextWithRequestSource(ctx, requestmeta.SourceRuler)
+	s = &mockQueryStreamServer{ctx: ctx}
+	err = i.QueryStream(rreq, s)
+	require.Nil(t, err)
 }
 
 func TestIngester_LabelValues_ShouldNotCreateTSDBIfDoesNotExists(t *testing.T) {

pkg/ruler/compat.go

@@ -188,6 +188,7 @@ func EngineQueryFunc(engine promql.QueryEngine, frontendClient *frontendClient,
 
 		// Add request ID to the context so that it can be used in logs and metrics for split queries.
 		ctx = requestmeta.ContextWithRequestId(ctx, uuid.NewString())
+		ctx = requestmeta.ContextWithRequestSource(ctx, requestmeta.SourceRuler)
 
 		if frontendClient != nil {
 			v, err := frontendClient.InstantQuery(ctx, qs, t)

pkg/storegateway/gateway.go

@@ -28,6 +28,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/util"
 	"github.com/cortexproject/cortex/pkg/util/flagext"
 	util_limiter "github.com/cortexproject/cortex/pkg/util/limiter"
+	"github.com/cortexproject/cortex/pkg/util/requestmeta"
 	"github.com/cortexproject/cortex/pkg/util/resource"
 	"github.com/cortexproject/cortex/pkg/util/services"
 	"github.com/cortexproject/cortex/pkg/util/validation"
@@ -408,30 +409,30 @@ func (g *StoreGateway) syncStores(ctx context.Context, reason string) {
 }
 
 func (g *StoreGateway) Series(req *storepb.SeriesRequest, srv storegatewaypb.StoreGateway_SeriesServer) error {
-	if err := g.checkResourceUtilization(); err != nil {
+	if err := g.checkResourceUtilization(srv.Context()); err != nil {
 		return err
 	}
 	return g.stores.Series(req, srv)
 }
 
 // LabelNames implements the Storegateway proto service.
 func (g *StoreGateway) LabelNames(ctx context.Context, req *storepb.LabelNamesRequest) (*storepb.LabelNamesResponse, error) {
-	if err := g.checkResourceUtilization(); err != nil {
+	if err := g.checkResourceUtilization(ctx); err != nil {
 		return nil, err
 	}
 	return g.stores.LabelNames(ctx, req)
 }
 
 // LabelValues implements the Storegateway proto service.
 func (g *StoreGateway) LabelValues(ctx context.Context, req *storepb.LabelValuesRequest) (*storepb.LabelValuesResponse, error) {
-	if err := g.checkResourceUtilization(); err != nil {
+	if err := g.checkResourceUtilization(ctx); err != nil {
 		return nil, err
 	}
 	return g.stores.LabelValues(ctx, req)
 }
 
-func (g *StoreGateway) checkResourceUtilization() error {
-	if g.resourceBasedLimiter == nil {
+func (g *StoreGateway) checkResourceUtilization(ctx context.Context) error {
+	if g.resourceBasedLimiter == nil || requestmeta.RequestFromRuler(ctx) {
 		return nil
 	}
 

pkg/storegateway/gateway_test.go

@@ -43,6 +43,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/util"
 	"github.com/cortexproject/cortex/pkg/util/flagext"
 	util_limiter "github.com/cortexproject/cortex/pkg/util/limiter"
+	"github.com/cortexproject/cortex/pkg/util/requestmeta"
 	"github.com/cortexproject/cortex/pkg/util/resource"
 	"github.com/cortexproject/cortex/pkg/util/services"
 	"github.com/cortexproject/cortex/pkg/util/test"
@@ -1236,11 +1237,17 @@ func TestStoreGateway_SeriesThrottledByResourceMonitor(t *testing.T) {
 	g.resourceBasedLimiter, err = util_limiter.NewResourceBasedLimiter(&mockResourceMonitor{cpu: 0.4, heap: 0.6}, limits, nil, "store-gateway")
 	require.NoError(t, err)
 
+	ctx = requestmeta.ContextWithRequestSource(ctx, requestmeta.SourceApi)
 	srv := newBucketStoreSeriesServer(setUserIDToGRPCContext(ctx, userID))
 	err = g.Series(req, srv)
 	require.Error(t, err)
 	exhaustedErr := util_limiter.ResourceLimitReachedError{}
 	require.ErrorContains(t, err, exhaustedErr.Error())
+
+	ctx = requestmeta.ContextWithRequestSource(ctx, requestmeta.SourceRuler)
+	srv = newBucketStoreSeriesServer(setUserIDToGRPCContext(ctx, userID))
+	err = g.Series(req, srv)
+	require.Nil(t, err)
 }
 
 func mockGatewayConfig() Config {

pkg/util/requestmeta/context.go

@@ -40,6 +40,7 @@ func ContextWithRequestMetadataMapFromHeaders(ctx context.Context, headers map[s
 		headerKeys = append(headerKeys, LoggingHeadersKey)
 	}
 	headerKeys = append(headerKeys, RequestIdKey)
+	headerKeys = append(headerKeys, RequestSourceKey)
 	for _, header := range headerKeys {
 		if v, ok := headers[textproto.CanonicalMIMEHeaderKey(header)]; ok {
 			headerMap[header] = v

pkg/util/requestmeta/source.go

@@ -0,0 +1,27 @@
+package requestmeta
+
+import "context"
+
+const RequestSourceKey = "x-cortex-request-source"
+
+const (
+	SourceApi   = "api"
+	SourceRuler = "ruler"
+)
+
+func ContextWithRequestSource(ctx context.Context, source string) context.Context {
+	metadataMap := MapFromContext(ctx)
+	if metadataMap == nil {
+		metadataMap = make(map[string]string)
+	}
+	metadataMap[RequestSourceKey] = source
+	return ContextWithRequestMetadataMap(ctx, metadataMap)
+}
+
+func RequestFromRuler(ctx context.Context) bool {
+	metadataMap := MapFromContext(ctx)
+	if metadataMap == nil {
+		return false
+	}
+	return metadataMap[RequestSourceKey] == SourceRuler
+}