[query-frontend] Set write deadline for /active_series requests

CHANGELOG.md

@@ -87,6 +87,7 @@
 * [ENHANCEMENT] API: Use github.com/klauspost/compress for faster gzip and deflate compression of API responses. #7475
 * [ENHANCEMENT] Ingester: Limiting on owned series (`-ingester.use-ingester-owned-series-for-limits`) now prevents discards in cases where a tenant is sharded across all ingesters (or shuffle sharding is disabled) and the ingester count increases. #7411
 * [ENHANCEMENT] Block upload: include converted timestamps in the error message if block is from the future. #7538
+* [ENHANCEMENT] Query-frontend: Introduce `-query-frontend.active-series-write-timeout` to allow configuring the server-side write timeout for active series requests. #7553
 * [BUGFIX] Ingester: don't ignore errors encountered while iterating through chunks or samples in response to a query request. #6451
 * [BUGFIX] Fix issue where queries can fail or omit OOO samples if OOO head compaction occurs between creating a querier and reading chunks #6766
 * [BUGFIX] Fix issue where concatenatingChunkIterator can obscure errors #6766

cmd/mimir/config-descriptor.json

@@ -4687,6 +4687,17 @@
           "fieldType": "boolean",
           "fieldCategory": "advanced"
         },
+        {
+          "kind": "field",
+          "name": "active_series_write_timeout",
+          "required": false,
+          "desc": "Timeout for writing active series responses. 0 means the value from `-server.http-write-timeout` is used.",
+          "fieldValue": null,
+          "fieldDefaultValue": 300000000000,
+          "fieldFlag": "query-frontend.active-series-write-timeout",
+          "fieldType": "duration",
+          "fieldCategory": "experimental"
+        },
         {
           "kind": "field",
           "name": "max_outstanding_per_tenant",

cmd/mimir/help-all.txt.tmpl

@@ -1799,6 +1799,8 @@ Usage of ./cmd/mimir/mimir:
     	[experimental] Number of series to buffer per store-gateway when streaming chunks from store-gateways. (default 256)
   -querier.timeout duration
     	The timeout for a query. This config option should be set on query-frontend too when query sharding is enabled. This also applies to queries evaluated by the ruler (internally or remotely). (default 2m0s)
+  -query-frontend.active-series-write-timeout duration
+    	[experimental] Timeout for writing active series responses. 0 means the value from `-server.http-write-timeout` is used. (default 5m0s)
   -query-frontend.additional-query-queue-dimensions-enabled
     	[experimental] Enqueue query requests with additional queue dimensions to split tenant request queues into subqueues. This enables separate requests to proceed from a tenant's subqueues even when other subqueues are blocked on slow query requests. Must be set on both query-frontend and scheduler to take effect. (default false)
   -query-frontend.align-queries-with-step

docs/sources/mimir/configure/about-versioning.md

@@ -136,6 +136,7 @@ The following features are currently experimental:
   - Query blocking on a per-tenant basis (configured with the limit `blocked_queries`)
   - Max number of tenants that may be queried at once (`-tenant-federation.max-tenants`)
   - Sharding of active series queries (`-query-frontend.shard-active-series-queries`)
+  - Server-side write timeout for responses to active series requests (`-query-frontend.active-series-write-timeout`)
 - Query-scheduler
   - `-query-scheduler.querier-forget-delay`
 - Store-gateway

docs/sources/mimir/configure/configuration-parameters/index.md

@@ -1385,6 +1385,11 @@ The `frontend` block configures the query-frontend.
 # CLI flag: -query-frontend.query-stats-enabled
 [query_stats_enabled: <boolean> | default = true]
 
+# (experimental) Timeout for writing active series responses. 0 means the value
+# from `-server.http-write-timeout` is used.
+# CLI flag: -query-frontend.active-series-write-timeout
+[active_series_write_timeout: <duration> | default = 5m]
+
 # (advanced) Maximum number of outstanding requests per tenant per frontend;
 # requests beyond this error with HTTP 429.
 # CLI flag: -querier.max-outstanding-requests-per-tenant

pkg/frontend/transport/handler.go

@@ -49,17 +49,19 @@ var (
 
 // HandlerConfig is a config for the handler.
 type HandlerConfig struct {
-	LogQueriesLongerThan   time.Duration          `yaml:"log_queries_longer_than"`
-	LogQueryRequestHeaders flagext.StringSliceCSV `yaml:"log_query_request_headers" category:"advanced"`
-	MaxBodySize            int64                  `yaml:"max_body_size" category:"advanced"`
-	QueryStatsEnabled      bool                   `yaml:"query_stats_enabled" category:"advanced"`
+	LogQueriesLongerThan     time.Duration          `yaml:"log_queries_longer_than"`
+	LogQueryRequestHeaders   flagext.StringSliceCSV `yaml:"log_query_request_headers" category:"advanced"`
+	MaxBodySize              int64                  `yaml:"max_body_size" category:"advanced"`
+	QueryStatsEnabled        bool                   `yaml:"query_stats_enabled" category:"advanced"`
+	ActiveSeriesWriteTimeout time.Duration          `yaml:"active_series_write_timeout" category:"experimental"`
 }
 
 func (cfg *HandlerConfig) RegisterFlags(f *flag.FlagSet) {
 	f.DurationVar(&cfg.LogQueriesLongerThan, "query-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.Var(&cfg.LogQueryRequestHeaders, "query-frontend.log-query-request-headers", "Comma-separated list of request header names to include in query logs. Applies to both query stats and slow queries logs.")
 	f.Int64Var(&cfg.MaxBodySize, "query-frontend.max-body-size", 10*1024*1024, "Max body size for downstream prometheus.")
 	f.BoolVar(&cfg.QueryStatsEnabled, "query-frontend.query-stats-enabled", true, "False to disable query statistics tracking. When enabled, a message with some statistics is logged for every query.")
+	f.DurationVar(&cfg.ActiveSeriesWriteTimeout, "query-frontend.active-series-write-timeout", 5*time.Minute, "Timeout for writing active series responses. 0 means the value from `-server.http-write-timeout` is used.")
 }
 
 // Handler accepts queries and forwards them to RoundTripper. It can wait on in-flight requests and log slow queries,
@@ -190,6 +192,16 @@ func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	activityIndex := f.at.Insert(func() string { return httpRequestActivity(r, r.Header.Get("User-Agent"), params) })
 	defer f.at.Delete(activityIndex)
 
+	if isActiveSeriesEndpoint(r) && f.cfg.ActiveSeriesWriteTimeout > 0 {
+		deadline := time.Now().Add(f.cfg.ActiveSeriesWriteTimeout)
+		err = http.NewResponseController(w).SetWriteDeadline(deadline)
+		if err != nil {
+			err := fmt.Errorf("failed to set write deadline for response writer: %w", err)
+			writeError(w, apierror.New(apierror.TypeInternal, err.Error()))
+			return
+		}
+	}
+
 	startTime := time.Now()
 	resp, err := f.roundTripper.RoundTrip(r)
 	queryResponseTime := time.Since(startTime)
@@ -441,3 +453,7 @@ func httpRequestActivity(request *http.Request, userAgent string, requestParams
 	// This doesn't have to be pretty, just useful for debugging, so prioritize efficiency.
 	return fmt.Sprintf("user:%s UA:%s req:%s %s %s", tenantID, userAgent, request.Method, request.URL.Path, params)
 }
+
+func isActiveSeriesEndpoint(r *http.Request) bool {
+	return strings.HasSuffix(r.URL.Path, "api/v1/cardinality/active_series")
+}

pkg/frontend/transport/handler_test.go

@@ -533,6 +533,47 @@ func TestHandler_LogsFormattedQueryDetails(t *testing.T) {
 	}
 }
 
+func TestHandler_ActiveSeriesWriteTimeout(t *testing.T) {
+	for _, tt := range []struct {
+		name      string
+		path      string
+		wantError bool
+	}{
+		{name: "deadline exceeded for non-streaming endpoint", path: "/api/v1/query", wantError: true},
+		{name: "deadline not exceeded for streaming endpoint", path: "/api/v1/cardinality/active_series"},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			const serverWriteTimeout = 50 * time.Millisecond
+
+			roundTripper := roundTripperFunc(func(req *http.Request) (*http.Response, error) {
+				// Simulate a request that takes longer than the server write timeout.
+				time.Sleep(2 * serverWriteTimeout)
+				return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader("{}"))}, nil
+			})
+
+			handler := NewHandler(HandlerConfig{ActiveSeriesWriteTimeout: time.Minute}, roundTripper, log.NewNopLogger(), nil, nil)
+
+			server := httptest.NewUnstartedServer(handler)
+			server.Config.WriteTimeout = serverWriteTimeout
+			server.Start()
+			defer server.Close()
+
+			req, err := http.NewRequest(http.MethodGet, fmt.Sprintf("%s%s", server.URL, tt.path), nil)
+			require.NoError(t, err)
+			resp, err := http.DefaultClient.Do(req)
+			if tt.wantError {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+			defer func() {
+				_, _ = io.Copy(io.Discard, resp.Body)
+				_ = resp.Body.Close()
+			}()
+		})
+	}
+}
+
 type testLogger struct {
 	logMessages []map[string]interface{}
 }

pkg/util/gziphandler/gzip.go

@@ -158,6 +158,12 @@ func (w *GzipResponseWriter) Write(b []byte) (int, error) {
 	return w.startPlainWrite(len(b))
 }
 
+// Unwrap returns the underlying ResponseWriter. This interface is used by
+// http.ResponseController to operate on the underlying ResponseWriter.
+func (w *GzipResponseWriter) Unwrap() http.ResponseWriter {
+	return w.ResponseWriter
+}
+
 func (w *GzipResponseWriter) startPlainWrite(blen int) (int, error) {
 	if err := w.startPlain(); err != nil {
 		return 0, err