query-frontend: add experimental header for tracking throughput

CHANGELOG.md

@@ -40,6 +40,7 @@
 * [ENHANCEMENT] Query-scheduler: Introduce `query-scheduler.use-multi-algorithm-query-queue`, which allows use of an experimental queue structure, with no change in external queue behavior. #7873
 * [ENHANCEMENT] Expose a new `s3.trace.enabled` configuration option to enable detailed logging of operations against S3-compatible object stores. #8690
 * [ENHANCEMENT] memberlist: locally-generated messages (e.g. ring updates) are sent to gossip network before forwarded messages. Introduced `-memberlist.broadcast-timeout-for-local-updates-on-shutdown` option to modify how long to wait until queue with locally-generated messages is empty when shutting down. Previously this was hard-coded to 10s, and wait included all messages (locally-generated and forwarded). Now it defaults to 10s, 0 means no timeout. Increasing this value may help to avoid problem when ring updates on shutdown are not propagated to other nodes, and ring entry is left in a wrong state. #8761
+* [ENGANCEMENT] Query-frontend: experimental support for the `X-Mimir-Query-Stats` header to provide some out of bound statistics. #7966
 * [BUGFIX] Ruler: add support for draining any outstanding alert notifications before shutting down. This can be enabled with the `-ruler.drain-notification-queue-on-shutdown=true` CLI flag. #8346
 * [BUGFIX] Query-frontend: fix `-querier.max-query-lookback` enforcement when `-compactor.blocks-retention-period` is not set, and viceversa. #8388
 * [BUGFIX] Ingester: fix sporadic `not found` error causing an internal server error if label names are queried with matchers during head compaction. #8391

pkg/api/handlers.go

@@ -277,7 +277,7 @@ func NewQuerierHandler(
 		// This is used for the stats API which we should not support. Or find other ways to.
 		prometheus.GathererFunc(func() ([]*dto.MetricFamily, error) { return nil, nil }),
 		reg,
-		nil,
+		querier.StatsRenderer,
 		remoteWriteEnabled,
 		nil,
 		oltpEnabled,

pkg/frontend/transport/handler.go

@@ -43,6 +43,7 @@ const (
 	ServiceTimingHeaderName   = "Server-Timing"
 	cacheControlHeader        = "Cache-Control"
 	cacheControlLogField      = "header_cache_control"
+	MimirQueryStatsHeaderName = "X-Mimir-Query-Stats"
 )
 
 var (
@@ -246,6 +247,7 @@ func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	if f.cfg.QueryStatsEnabled {
 		writeServiceTimingHeader(queryResponseTime, hs, queryDetails.QuerierStats)
+		writeStatsHeader(hs, queryDetails.QuerierStats)
 	}
 
 	w.WriteHeader(resp.StatusCode)
@@ -331,6 +333,7 @@ func (f *Handler) reportQueryStats(
 		"split_queries", stats.LoadSplitQueries(),
 		"estimated_series_count", stats.GetEstimatedSeriesCount(),
 		"queue_time_seconds", stats.LoadQueueTime().Seconds(),
+		"total_samples", stats.LoadTotalSamples(),
 	}, formatQueryString(details, queryString)...)
 
 	if details != nil {
@@ -487,6 +490,12 @@ func writeServiceTimingHeader(queryResponseTime time.Duration, headers http.Head
 	}
 }
 
+func writeStatsHeader(headers http.Header, stats *querier_stats.Stats) {
+	if stats != nil {
+		headers.Set(MimirQueryStatsHeaderName, fmt.Sprintf("total_samples=%d", stats.LoadTotalSamples()))
+	}
+}
+
 func statsValue(name string, d time.Duration) string {
 	durationInMs := strconv.FormatFloat(float64(d)/float64(time.Millisecond), 'f', -1, 64)
 	return name + ";dur=" + durationInMs

pkg/frontend/transport/handler_test.go

@@ -84,16 +84,17 @@ func TestHandler_ServeHTTP(t *testing.T) {
 	}
 
 	for _, tt := range []struct {
-		name                    string
-		cfg                     HandlerConfig
-		request                 func() *http.Request
-		downstreamResponse      *http.Response
-		downstreamErr           error
-		expectedStatusCode      int
-		expectedParams          url.Values
-		expectedMetrics         int
-		expectedActivity        string
-		expectedReadConsistency string
+		name                     string
+		cfg                      HandlerConfig
+		request                  func() *http.Request
+		downstreamResponse       *http.Response
+		downstreamErr            error
+		expectedStatusCode       int
+		expectedParams           url.Values
+		expectedMetrics          int
+		expectedActivity         string
+		expectedReadConsistency  string
+		expectedExtraStatsHeader string
 	}{
 		{
 			name: "handler with stats enabled, POST request with params",
@@ -114,9 +115,10 @@ func TestHandler_ServeHTTP(t *testing.T) {
 				"query": []string{"some_metric"},
 				"time":  []string{"42"},
 			},
-			expectedMetrics:         5,
-			expectedActivity:        "user:12345 UA:test-user-agent req:POST /api/v1/query query=some_metric&time=42",
-			expectedReadConsistency: "",
+			expectedMetrics:          5,
+			expectedActivity:         "user:12345 UA:test-user-agent req:POST /api/v1/query query=some_metric&time=42",
+			expectedReadConsistency:  "",
+			expectedExtraStatsHeader: "total_samples=0",
 		},
 		{
 			name: "handler with stats enabled, GET request with params",
@@ -132,9 +134,10 @@ func TestHandler_ServeHTTP(t *testing.T) {
 				"query": []string{"some_metric"},
 				"time":  []string{"42"},
 			},
-			expectedMetrics:         5,
-			expectedActivity:        "user:12345 UA:test-user-agent req:GET /api/v1/query query=some_metric&time=42",
-			expectedReadConsistency: "",
+			expectedMetrics:          5,
+			expectedActivity:         "user:12345 UA:test-user-agent req:GET /api/v1/query query=some_metric&time=42",
+			expectedReadConsistency:  "",
+			expectedExtraStatsHeader: "total_samples=0",
 		},
 		{
 			name: "handler with stats enabled, GET request with params and read consistency specified",
@@ -150,9 +153,10 @@ func TestHandler_ServeHTTP(t *testing.T) {
 				"query": []string{"some_metric"},
 				"time":  []string{"42"},
 			},
-			expectedMetrics:         5,
-			expectedActivity:        "user:12345 UA:test-user-agent req:GET /api/v1/query query=some_metric&time=42",
-			expectedReadConsistency: api.ReadConsistencyStrong,
+			expectedMetrics:          5,
+			expectedActivity:         "user:12345 UA:test-user-agent req:GET /api/v1/query query=some_metric&time=42",
+			expectedReadConsistency:  api.ReadConsistencyStrong,
+			expectedExtraStatsHeader: "total_samples=0",
 		},
 		{
 			name: "handler with stats enabled, GET request without params",
@@ -162,12 +166,13 @@ func TestHandler_ServeHTTP(t *testing.T) {
 				r.Header.Add("User-Agent", "test-user-agent")
 				return r
 			},
-			downstreamResponse:      makeSuccessfulDownstreamResponse(),
-			expectedStatusCode:      200,
-			expectedParams:          url.Values{},
-			expectedMetrics:         5,
-			expectedActivity:        "user:12345 UA:test-user-agent req:GET /api/v1/query (no params)",
-			expectedReadConsistency: "",
+			downstreamResponse:       makeSuccessfulDownstreamResponse(),
+			expectedStatusCode:       200,
+			expectedParams:           url.Values{},
+			expectedMetrics:          5,
+			expectedActivity:         "user:12345 UA:test-user-agent req:GET /api/v1/query (no params)",
+			expectedReadConsistency:  "",
+			expectedExtraStatsHeader: "total_samples=0",
 		},
 		{
 			name: "handler with stats disabled, GET request with params",
@@ -235,6 +240,7 @@ func TestHandler_ServeHTTP(t *testing.T) {
 				"end_1":      []string{"20"},
 				"hints_1":    []string{"{\"step_ms\":1000}"},
 			},
+			expectedExtraStatsHeader: "total_samples=0",
 		},
 		{
 			name: "downstream returns an apierror with 4xx status code",
@@ -320,6 +326,12 @@ func TestHandler_ServeHTTP(t *testing.T) {
 			responseData, _ := io.ReadAll(resp.Body)
 			require.Equal(t, tt.expectedStatusCode, resp.Code)
 
+			if tt.expectedExtraStatsHeader == "" {
+				assert.NotContains(t, resp.Header(), "X-Mimir-Query-Stats")
+			} else {
+				assert.Equal(t, tt.expectedExtraStatsHeader, resp.Header().Get("X-Mimir-Query-Stats"))
+			}
+
 			count, err := promtest.GatherAndCount(
 				reg,
 				"cortex_query_seconds_total",

pkg/querier/remote_read_test.go

@@ -53,6 +53,10 @@ type mockQuerier struct {
 	selectFn func(ctx context.Context, sorted bool, hints *storage.SelectHints, matchers ...*labels.Matcher) storage.SeriesSet
 }
 
+func (m mockQuerier) Close() error {
+	return nil
+}
+
 func (m mockQuerier) Select(ctx context.Context, sorted bool, hints *storage.SelectHints, matchers ...*labels.Matcher) storage.SeriesSet {
 	if m.selectFn != nil {
 		return m.selectFn(ctx, sorted, hints, matchers...)

pkg/querier/stats/stats.go

@@ -187,6 +187,22 @@ func (s *Stats) LoadQueueTime() time.Duration {
 	return time.Duration(atomic.LoadInt64((*int64)(&s.QueueTime)))
 }
 
+func (s *Stats) AddTotalSamples(c uint64) {
+	if s == nil {
+		return
+	}
+
+	atomic.AddUint64(&s.TotalSamples, c)
+}
+
+func (s *Stats) LoadTotalSamples() uint64 {
+	if s == nil {
+		return 0
+	}
+
+	return atomic.LoadUint64(&s.TotalSamples)
+}
+
 // Merge the provided Stats into this one.
 func (s *Stats) Merge(other *Stats) {
 	if s == nil || other == nil {
@@ -202,6 +218,7 @@ func (s *Stats) Merge(other *Stats) {
 	s.AddFetchedIndexBytes(other.LoadFetchedIndexBytes())
 	s.AddEstimatedSeriesCount(other.LoadEstimatedSeriesCount())
 	s.AddQueueTime(other.LoadQueueTime())
+	s.AddTotalSamples(other.LoadTotalSamples())
 }
 
 // Copy returns a copy of the stats. Use this rather than regular struct assignment

pkg/querier/stats/stats.proto

@@ -34,4 +34,7 @@ message Stats {
   uint64 estimated_series_count = 8;
   // The sum of durations that the query spent in the queue, before it was handled by querier.
   google.protobuf.Duration queue_time = 9 [(gogoproto.stdduration) = true, (gogoproto.nullable) = false];
+  // TotalSamples represents the total number of samples scanned while evaluating a query.
+  // This value is taken from the Prometheus engine github.com/prometheus/prometheus/util/stats/query_stats.go
+  uint64 total_samples = 10;
 }