QFE: Add active query tracker to query frontend

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 * [FEATURE] Ruler: Add experimental support for federated rule groups. A rule group listing tenants in its `source_tenants` field is evaluated against those tenants while the resulting series and alerts are written to the tenant owning the rule group. Enabled with `-ruler.enable-federated-rules` (requires `-tenant-federation.enabled`), and restricted to selected tenants with `-ruler.allowed-federated-tenants` and `-ruler.disallowed-federated-tenants`. #7828
 * [FEATURE] Distributor: Add experimental OTLP metrics ingestion over gRPC (`opentelemetry.proto.collector.metrics.v1.MetricsService/Export`) on the gRPC server port. Enable it with `-distributor.otlp.grpc-enabled`. The tenant is read from the `X-Scope-OrgID` gRPC metadata, and the request size is limited by `-server.grpc-max-recv-msg-size-bytes`. #7873
+* [FEATURE] Query Frontend: Add an active query tracker, enabled with `-frontend.active-query-tracker-dir`, that logs on startup the queries which were in flight when the query frontend last stopped. Add `-frontend.max-concurrent` to limit the number of inflight queries per query frontend, enforced by the tracker. Queries above the limit are rejected with 429. Both are disabled by default. #6935
 * [ENHANCEMENT] Ruler: Add rule information to query parameters sent to the Query Frontend when `-ruler.frontend-address` is configured, so the Query Frontend query stats logs (`-frontend.enabled-ruler-query-stats`) can identify the evaluated rule. In this mode, the Ruler query stats log (`-ruler.query-stats-enabled`) omits the fetched statistics fields, which are always zero. #6539
 * [ENHANCEMENT] gRPC: When `-auth.enabled=true`, a gRPC request without the `X-Scope-OrgID` metadata is now rejected with the `UNAUTHENTICATED` status code instead of `UNKNOWN`. #7873
 * [ENHANCEMENT] Upgrade promql-engine and Thanos to latest. #7889

docs/configuration/config-file-reference.md

@@ -5578,6 +5578,19 @@ The `query_frontend_config` configures the Cortex query-frontend.
 # CLI flag: -frontend.enabled-ruler-query-stats
 [enabled_ruler_query_stats_log: <boolean> | default = false]
 
+# Active query tracker monitors queries in flight in the query frontend, and
+# writes them to a file in the given directory. If Cortex discovers any queries
+# in this file during startup, it will log them. Empty value disables the active
+# query tracker.
+# CLI flag: -frontend.active-query-tracker-dir
+[active_query_tracker_dir: <string> | default = ""]
+
+# The maximum number of inflight queries in the query frontend, enforced by the
+# active query tracker. Queries above the limit are rejected with 429. Requires
+# -frontend.active-query-tracker-dir. 0 to disable.
+# CLI flag: -frontend.max-concurrent
+[max_concurrent: <int> | default = 0]
+
 # If a querier disconnects without sending notification about graceful shutdown,
 # the query-frontend will keep the querier in the tenant's shard until the
 # forget delay has passed. This feature is useful to reduce the blast radius

pkg/cortex/cortex.go

@@ -285,6 +285,9 @@ func (c *Config) Validate(log log.Logger) error {
 	if err := c.Querier.Validate(c.ResourceMonitor.Resources); err != nil {
 		return errors.Wrap(err, "invalid querier config")
 	}
+	if err := c.Frontend.Handler.Validate(); err != nil {
+		return errors.Wrap(err, "invalid frontend config")
+	}
 	if c.Querier.TimeoutClassificationEnabled && !c.Frontend.Handler.QueryStatsEnabled {
 		return errTimeoutClassificationRequiresQueryStats
 	}

pkg/frontend/transport/handler.go

@@ -32,6 +32,7 @@ import (
 	util_api "github.com/cortexproject/cortex/pkg/util/api"
 	"github.com/cortexproject/cortex/pkg/util/limiter"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
+	"github.com/cortexproject/cortex/pkg/util/request_tracker"
 	"github.com/cortexproject/cortex/pkg/util/requestmeta"
 	"github.com/cortexproject/cortex/pkg/util/users"
 )
@@ -40,19 +41,27 @@ const (
 	// StatusClientClosedRequest is the status code for when a client request cancellation of a http request
 	StatusClientClosedRequest = 499
 	ServiceTimingHeaderName   = "Server-Timing"
+
+	activeQueryTrackerFileName = "frontend_queries.active"
+	// Number of tracked queries when -frontend.max-concurrent is disabled. Queries above it are served but not tracked.
+	defaultActiveQueryTrackerSize = 500
 )
 
+var errMaxConcurrentRequiresActiveQueryTracker = errors.New("-frontend.max-concurrent requires -frontend.active-query-tracker-dir to be set")
+
 var (
-	errCanceled              = httpgrpc.Errorf(StatusClientClosedRequest, "%s", context.Canceled.Error())
-	errDeadlineExceeded      = httpgrpc.Errorf(http.StatusGatewayTimeout, "%s", context.DeadlineExceeded.Error())
-	errRequestEntityTooLarge = httpgrpc.Errorf(http.StatusRequestEntityTooLarge, "%s", "http: request body too large")
+	errCanceled               = httpgrpc.Errorf(StatusClientClosedRequest, "%s", context.Canceled.Error())
+	errDeadlineExceeded       = httpgrpc.Errorf(http.StatusGatewayTimeout, "%s", context.DeadlineExceeded.Error())
+	errRequestEntityTooLarge  = httpgrpc.Errorf(http.StatusRequestEntityTooLarge, "%s", "http: request body too large")
+	errTooManyInflightQueries = httpgrpc.Errorf(http.StatusTooManyRequests, "%s", "too many inflight queries in the query frontend")
 )
 
 const (
 	reasonTooManyTenants           = "too_many_tenants"
 	reasonRequestBodySizeExceeded  = "request_body_size_exceeded"
 	reasonResponseBodySizeExceeded = "response_body_size_exceeded"
 	reasonTooManyRequests          = "too_many_requests"
+	reasonTooManyInflightQueries   = "too_many_inflight_queries"
 	reasonResourceExhausted        = "resource_exhausted"
 	reasonTimeRangeExceeded        = "time_range_exceeded"
 	reasonResponseSizeExceeded     = "response_size_exceeded"
@@ -89,13 +98,27 @@ type HandlerConfig struct {
 	MaxBodySize               int64         `yaml:"max_body_size"`
 	QueryStatsEnabled         bool          `yaml:"query_stats_enabled"`
 	EnabledRulerQueryStatsLog bool          `yaml:"enabled_ruler_query_stats_log"`
+	ActiveQueryTrackerDir     string        `yaml:"active_query_tracker_dir"`
+	MaxConcurrent             int           `yaml:"max_concurrent"`
 }
 
 func (cfg *HandlerConfig) RegisterFlags(f *flag.FlagSet) {
 	f.DurationVar(&cfg.LogQueriesLongerThan, "frontend.log-queries-longer-than", 0, "Log queries that are slower than the specified duration. Set to 0 to disable. Set to < 0 to enable on all queries.")
 	f.Int64Var(&cfg.MaxBodySize, "frontend.max-body-size", 10*1024*1024, "Max body size for downstream prometheus.")
 	f.BoolVar(&cfg.QueryStatsEnabled, "frontend.query-stats-enabled", false, "True to enable query statistics tracking. When enabled, a message with some statistics is logged for every query.")
 	f.BoolVar(&cfg.EnabledRulerQueryStatsLog, "frontend.enabled-ruler-query-stats", false, "If enabled, report the query stats log for queries coming from the ruler to evaluate rules. It only takes effect when '-ruler.frontend-address' is configured.")
+	f.StringVar(&cfg.ActiveQueryTrackerDir, "frontend.active-query-tracker-dir", "", "Active query tracker monitors queries in flight in the query frontend, and writes them to a file in the given directory. If Cortex discovers any queries in this file during startup, it will log them. Empty value disables the active query tracker.")
+	f.IntVar(&cfg.MaxConcurrent, "frontend.max-concurrent", 0, "The maximum number of inflight queries in the query frontend, enforced by the active query tracker. Queries above the limit are rejected with 429. Requires -frontend.active-query-tracker-dir. 0 to disable.")
+}
+
+func (cfg *HandlerConfig) Validate() error {
+	if cfg.MaxConcurrent < 0 {
+		return errors.New("-frontend.max-concurrent must be non-negative")
+	}
+	if cfg.MaxConcurrent > 0 && cfg.ActiveQueryTrackerDir == "" {
+		return errMaxConcurrentRequiresActiveQueryTracker
+	}
+	return nil
 }
 
 // Handler accepts queries and forwards them to RoundTripper. It can log slow queries,
@@ -118,6 +141,7 @@ type Handler struct {
 	rejectedQueries     *prometheus.CounterVec
 	slowQueries         *prometheus.CounterVec
 	activeUsers         *users.ActiveUsersCleanupService
+	activeQueryTracker  *request_tracker.RequestTracker
 
 	initSlowQueryMetric sync.Once
 	reg                 prometheus.Registerer
@@ -186,6 +210,17 @@ func NewHandler(cfg HandlerConfig, tenantFederationCfg tenantfederation.Config,
 		_ = h.activeUsers.StartAsync(context.Background())
 	}
 
+	if cfg.ActiveQueryTrackerDir != "" {
+		trackerSize := defaultActiveQueryTrackerSize
+		if cfg.MaxConcurrent > 0 {
+			trackerSize = cfg.MaxConcurrent
+		}
+		h.activeQueryTracker = request_tracker.NewRequestTracker(cfg.ActiveQueryTrackerDir, activeQueryTrackerFileName, trackerSize, util_log.GoKitLogToSlog(log))
+		if h.activeQueryTracker == nil && cfg.MaxConcurrent > 0 {
+			level.Error(log).Log("msg", "failed to create the active query tracker, -frontend.max-concurrent is not enforced", "dir", cfg.ActiveQueryTrackerDir)
+		}
+	}
+
 	return h
 }
 
@@ -309,6 +344,19 @@ func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		r.Body = io.NopCloser(&buf)
 	}
 
+	if f.activeQueryTracker != nil {
+		index, ok := f.activeQueryTracker.TryInsert(request_tracker.ExtractorForPath(r.URL.Path).Extract(r))
+		if ok {
+			defer f.activeQueryTracker.Delete(index)
+		} else if f.cfg.MaxConcurrent > 0 {
+			if f.cfg.QueryStatsEnabled {
+				f.rejectedQueries.WithLabelValues(reasonTooManyInflightQueries, source, userID).Inc()
+			}
+			writeError(util_log.WithContext(r.Context(), f.log), w, errTooManyInflightQueries, nil)
+			return
+		}
+	}
+
 	// Log request if the request is not remote read.
 	// We need to parse remote read proto to be properly log it so skip it.
 	if f.cfg.QueryStatsEnabled && !isRemoteRead {

pkg/frontend/transport/handler_test.go

@@ -8,6 +8,8 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"net/url"
+	"os"
+	"path/filepath"
 	"strings"
 	"testing"
 	"time"
@@ -29,6 +31,7 @@ import (
 	querier_stats "github.com/cortexproject/cortex/pkg/querier/stats"
 	"github.com/cortexproject/cortex/pkg/querier/tenantfederation"
 	util_api "github.com/cortexproject/cortex/pkg/util/api"
+	"github.com/cortexproject/cortex/pkg/util/flagext"
 	"github.com/cortexproject/cortex/pkg/util/limiter"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
 	"github.com/cortexproject/cortex/pkg/util/requestmeta"
@@ -1043,3 +1046,125 @@ func TestHandler_RemoteReadRequest_DoesNotParseQueryString(t *testing.T) {
 	// Verify that the request body is still readable (not replaced with empty buffer)
 	require.NotEmpty(t, string(bodyBytes))
 }
+
+func TestHandler_ActiveQueryTrackerAndConcurrencyDisabledByDefault(t *testing.T) {
+	cfg := HandlerConfig{}
+	flagext.DefaultValues(&cfg)
+
+	handler := NewHandler(cfg, tenantfederation.Config{}, nil, http.DefaultTransport, log.NewNopLogger(), nil)
+	assert.Nil(t, handler.activeQueryTracker)
+	assert.NoError(t, cfg.Validate())
+}
+
+func TestHandlerConfig_Validate(t *testing.T) {
+	assert.NoError(t, (&HandlerConfig{}).Validate())
+	assert.NoError(t, (&HandlerConfig{ActiveQueryTrackerDir: "dir"}).Validate())
+	assert.NoError(t, (&HandlerConfig{ActiveQueryTrackerDir: "dir", MaxConcurrent: 10}).Validate())
+	assert.ErrorIs(t, (&HandlerConfig{MaxConcurrent: 10}).Validate(), errMaxConcurrentRequiresActiveQueryTracker)
+	assert.Error(t, (&HandlerConfig{ActiveQueryTrackerDir: "dir", MaxConcurrent: -1}).Validate())
+}
+
+func TestHandler_ActiveQueryTracker(t *testing.T) {
+	dir := t.TempDir()
+	trackerFile := filepath.Join(dir, activeQueryTrackerFileName)
+	query := `sum(rate(http_requests_total{job="api"}[5m]))`
+
+	inFlight := make(chan string, 1)
+	roundTripper := roundTripperFunc(func(req *http.Request) (*http.Response, error) {
+		content, err := os.ReadFile(trackerFile)
+		require.NoError(t, err)
+		inFlight <- string(content)
+		return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader("{}"))}, nil
+	})
+
+	handler := NewHandler(HandlerConfig{ActiveQueryTrackerDir: dir, MaxBodySize: 1024}, tenantfederation.Config{}, nil, roundTripper, log.NewNopLogger(), nil)
+	require.NotNil(t, handler.activeQueryTracker)
+
+	form := url.Values{"query": []string{query}, "start": []string{"1000"}, "end": []string{"2000"}, "step": []string{"15"}}
+	req := httptest.NewRequest(http.MethodPost, "/prometheus/api/v1/query_range", strings.NewReader(form.Encode()))
+	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
+	req = req.WithContext(user.InjectOrgID(context.Background(), "user-1"))
+
+	resp := httptest.NewRecorder()
+	handler.ServeHTTP(resp, req)
+	require.Equal(t, http.StatusOK, resp.Code)
+
+	tracked := <-inFlight
+	assert.Contains(t, tracked, `"TenantID":"user-1"`)
+	assert.Contains(t, tracked, `"Path":"/prometheus/api/v1/query_range"`)
+	assert.Contains(t, tracked, `"step":"15"`)
+	assert.Contains(t, tracked, `sum(rate(http_requests_total{job=\"api\"}[5m]))`)
+
+	content, err := os.ReadFile(trackerFile)
+	require.NoError(t, err)
+	assert.NotContains(t, string(content), "user-1")
+}
+
+func TestHandler_MaxConcurrent(t *testing.T) {
+	started := make(chan struct{})
+	release := make(chan struct{})
+	roundTripper := roundTripperFunc(func(req *http.Request) (*http.Response, error) {
+		started <- struct{}{}
+		<-release
+		return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader("{}"))}, nil
+	})
+
+	reg := prometheus.NewPedanticRegistry()
+	cfg := HandlerConfig{ActiveQueryTrackerDir: t.TempDir(), MaxConcurrent: 1, MaxBodySize: 1024, QueryStatsEnabled: true}
+	handler := NewHandler(cfg, tenantfederation.Config{}, nil, roundTripper, log.NewNopLogger(), reg)
+
+	newRequest := func() *http.Request {
+		req := httptest.NewRequest(http.MethodGet, "/api/v1/query?query=up", nil)
+		return req.WithContext(user.InjectOrgID(context.Background(), "user-1"))
+	}
+
+	firstDone := make(chan int)
+	go func() {
+		resp := httptest.NewRecorder()
+		handler.ServeHTTP(resp, newRequest())
+		firstDone <- resp.Code
+	}()
+	<-started
+
+	resp := httptest.NewRecorder()
+	handler.ServeHTTP(resp, newRequest())
+	assert.Equal(t, http.StatusTooManyRequests, resp.Code)
+	assert.Equal(t, 1.0, promtest.ToFloat64(handler.rejectedQueries.WithLabelValues(reasonTooManyInflightQueries, requestmeta.SourceAPI, "user-1")))
+
+	close(release)
+	assert.Equal(t, http.StatusOK, <-firstDone)
+
+	go func() { <-started }()
+	resp = httptest.NewRecorder()
+	handler.ServeHTTP(resp, newRequest())
+	assert.Equal(t, http.StatusOK, resp.Code)
+}
+
+func TestHandler_ActiveQueryTrackerWithoutLimitDoesNotReject(t *testing.T) {
+	started := make(chan struct{}, defaultActiveQueryTrackerSize+1)
+	release := make(chan struct{})
+	roundTripper := roundTripperFunc(func(req *http.Request) (*http.Response, error) {
+		started <- struct{}{}
+		<-release
+		return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader("{}"))}, nil
+	})
+
+	handler := NewHandler(HandlerConfig{ActiveQueryTrackerDir: t.TempDir(), MaxBodySize: 1024}, tenantfederation.Config{}, nil, roundTripper, log.NewNopLogger(), nil)
+
+	codes := make(chan int, defaultActiveQueryTrackerSize+1)
+	for range defaultActiveQueryTrackerSize + 1 {
+		go func() {
+			req := httptest.NewRequest(http.MethodGet, "/api/v1/query?query=up", nil)
+			resp := httptest.NewRecorder()
+			handler.ServeHTTP(resp, req.WithContext(user.InjectOrgID(context.Background(), "user-1")))
+			codes <- resp.Code
+		}()
+	}
+	for range defaultActiveQueryTrackerSize + 1 {
+		<-started
+	}
+	close(release)
+	for range defaultActiveQueryTrackerSize + 1 {
+		assert.Equal(t, http.StatusOK, <-codes)
+	}
+}

pkg/util/request_tracker/request_extractor.go

@@ -3,6 +3,7 @@ package request_tracker
 import (
 	"encoding/json"
 	"net/http"
+	"net/url"
 	"strings"
 	"time"
 	"unicode/utf8"
@@ -21,6 +22,33 @@ type InstantQueryExtractor struct{}
 
 type RangedQueryExtractor struct{}
 
+var (
+	apiExtractor          = &ApiExtractor{}
+	instantQueryExtractor = &InstantQueryExtractor{}
+	rangedQueryExtractor  = &RangedQueryExtractor{}
+)
+
+// ExtractorForPath returns the extractor matching the Prometheus API endpoint of the given path.
+func ExtractorForPath(path string) Extractor {
+	switch {
+	case strings.HasSuffix(path, "/api/v1/query"):
+		return instantQueryExtractor
+	case strings.HasSuffix(path, "/api/v1/query_range"):
+		return rangedQueryExtractor
+	default:
+		return apiExtractor
+	}
+}
+
+// requestParams returns the parsed form when available, so parameters sent in a POST body are
+// included, and falls back to the URL query otherwise.
+func requestParams(r *http.Request) url.Values {
+	if r.Form != nil {
+		return r.Form
+	}
+	return r.URL.Query()
+}
+
 func generateCommonMap(r *http.Request) map[string]any {
 	ctx := r.Context()
 	entryMap := make(map[string]any)
@@ -38,11 +66,12 @@ func generateCommonMap(r *http.Request) map[string]any {
 
 func (e *ApiExtractor) Extract(r *http.Request) []byte {
 	entryMap := generateCommonMap(r)
-	entryMap["limit"] = r.URL.Query().Get("limit")
-	entryMap["start"] = r.URL.Query().Get("start")
-	entryMap["end"] = r.URL.Query().Get("end")
+	params := requestParams(r)
+	entryMap["limit"] = params.Get("limit")
+	entryMap["start"] = params.Get("start")
+	entryMap["end"] = params.Get("end")
 
-	matches := r.URL.Query()["match[]"]
+	matches := params["match[]"]
 	entryMap["numberOfMatches"] = len(matches)
 	matchesStr := strings.Join(matches, ",")
 
@@ -51,16 +80,18 @@ func (e *ApiExtractor) Extract(r *http.Request) []byte {
 
 func (e *InstantQueryExtractor) Extract(r *http.Request) []byte {
 	entryMap := generateCommonMap(r)
-	entryMap["time"] = r.URL.Query().Get("time")
-	return generateJSONEntryWithTruncatedField(entryMap, "query", r.URL.Query().Get("query"))
+	params := requestParams(r)
+	entryMap["time"] = params.Get("time")
+	return generateJSONEntryWithTruncatedField(entryMap, "query", params.Get("query"))
 }
 
 func (e *RangedQueryExtractor) Extract(r *http.Request) []byte {
 	entryMap := generateCommonMap(r)
-	entryMap["start"] = r.URL.Query().Get("start")
-	entryMap["end"] = r.URL.Query().Get("end")
-	entryMap["step"] = r.URL.Query().Get("step")
-	return generateJSONEntryWithTruncatedField(entryMap, "query", r.URL.Query().Get("query"))
+	params := requestParams(r)
+	entryMap["start"] = params.Get("start")
+	entryMap["end"] = params.Get("end")
+	entryMap["step"] = params.Get("step")
+	return generateJSONEntryWithTruncatedField(entryMap, "query", params.Get("query"))
 }
 
 func generateJSONEntry(entryMap map[string]any) []byte {