Return server-side total bytes processed statistics as a header through query frontend

CHANGELOG.md

@@ -42,6 +42,7 @@
   * `-memberlist.max-concurrent-writes`
   * `-memberlist.acquire-writer-timeout`
 * [ENHANCEMENT] memberlist: Notifications can now be processed once per interval specified by `-memberlist.notify-interval` to reduce notify storm CPU activity in large clusters. #9594
+* [ENHANCEMENT] Return server-side total bytes processed statistics as a header through query frontend. #9645 
 * [BUGFIX] Fix issue where functions such as `rate()` over native histograms could return incorrect values if a float stale marker was present in the selected range. #9508
 * [BUGFIX] Fix issue where negation of native histograms (eg. `-some_native_histogram_series`) did nothing. #9508
 * [BUGFIX] Fix issue where `metric might not be a counter, name does not end in _total/_sum/_count/_bucket` annotation would be emitted even if `rate` or `increase` did not have enough samples to compute a result. #9508

pkg/frontend/transport/handler.go

@@ -39,10 +39,11 @@ import (
 
 const (
 	// StatusClientClosedRequest is the status code for when a client request cancellation of an http request
-	StatusClientClosedRequest = 499
-	ServiceTimingHeaderName   = "Server-Timing"
-	cacheControlHeader        = "Cache-Control"
-	cacheControlLogField      = "header_cache_control"
+	StatusClientClosedRequest  = 499
+	ServiceTimingHeaderName    = "Server-Timing"
+	cacheControlHeader         = "Cache-Control"
+	cacheControlLogField       = "header_cache_control"
+	ServiceTotalBytesProcessed = "X-Total-Bytes-Processed"
 )
 
 var (
@@ -246,6 +247,7 @@ func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	if f.cfg.QueryStatsEnabled {
 		writeServiceTimingHeader(queryResponseTime, hs, queryDetails.QuerierStats)
+		writeTotalBytesProcessedHeader(hs, queryDetails.QuerierStats)
 	}
 
 	w.WriteHeader(resp.StatusCode)
@@ -488,6 +490,13 @@ func writeServiceTimingHeader(queryResponseTime time.Duration, headers http.Head
 	}
 }
 
+func writeTotalBytesProcessedHeader(headers http.Header, stats *querier_stats.Stats) {
+	if stats != nil {
+		totalBytes := stats.LoadFetchedChunkBytes() + stats.LoadFetchedIndexBytes()
+		headers.Set(ServiceTotalBytesProcessed, strconv.FormatUint(totalBytes, 10))
+	}
+}
+
 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

@@ -94,6 +94,7 @@ func TestHandler_ServeHTTP(t *testing.T) {
 		expectedMetrics         int
 		expectedActivity        string
 		expectedReadConsistency string
+		assertHeaders           func(t *testing.T, headers http.Header)
 	}{
 		{
 			name: "handler with stats enabled, POST request with params",
@@ -284,6 +285,27 @@ func TestHandler_ServeHTTP(t *testing.T) {
 			expectedActivity:        "user:12345 UA: req:GET /api/v1/query query=some_metric&time=42",
 			expectedReadConsistency: "",
 		},
+		{
+			name: "handler with stats enabled, check ServiceTotalBytesProcessed header",
+			cfg:  HandlerConfig{QueryStatsEnabled: true, MaxBodySize: 1024},
+			request: func() *http.Request {
+				req := httptest.NewRequest(http.MethodPost, "/api/v1/query", strings.NewReader("query=some_metric&time=42"))
+				req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
+				return req
+			},
+			downstreamResponse: makeSuccessfulDownstreamResponse(),
+			expectedStatusCode: 200,
+			expectedParams: url.Values{
+				"query": []string{"some_metric"},
+				"time":  []string{"42"},
+			},
+			expectedMetrics:         5,
+			expectedActivity:        "user:12345 UA: req:POST /api/v1/query query=some_metric&time=42",
+			expectedReadConsistency: "",
+			assertHeaders: func(t *testing.T, headers http.Header) {
+				assert.Equal(t, "0", headers.Get(ServiceTotalBytesProcessed))
+			},
+		},
 	} {
 		t.Run(tt.name, func(t *testing.T) {
 			activityFile := filepath.Join(t.TempDir(), "activity-tracker")
@@ -370,6 +392,9 @@ func TestHandler_ServeHTTP(t *testing.T) {
 				if tt.expectedStatusCode >= 200 && tt.expectedStatusCode < 300 {
 					require.Equal(t, int64(len(responseData)), msg["response_size_bytes"])
 				}
+				if tt.assertHeaders != nil {
+					tt.assertHeaders(t, resp.Header())
+				}
 
 				// Check that the HTTP or Protobuf request parameters are logged.
 				paramsLogged := 0