Return server-side total bytes processed statistics as a header through query frontend (grafana#9645)

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

* update chanlog.md

* lint error fix in changelog.md

* Add bytes processed stats to existing Server-Timing header

* Fix failing integration tests

Author: madhu-reddy-peram

CHANGELOG.md

@@ -44,6 +44,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

integration/query_frontend_test.go

@@ -383,7 +383,7 @@ func runQueryFrontendTest(t *testing.T, cfg queryFrontendTestConfig) {
 		if userID == 0 && cfg.queryStatsEnabled {
 			res, _, err := c.QueryRaw("{instance=~\"hello.*\"}")
 			require.NoError(t, err)
-			require.Regexp(t, "querier_wall_time;dur=[0-9.]*, response_time;dur=[0-9.]*$", res.Header.Values("Server-Timing")[0])
+			require.Regexp(t, "querier_wall_time;dur=[0-9.]*, response_time;dur=[0-9.]*, bytes_processed=[0-9.]*$", res.Header.Values("Server-Timing")[0])
 		}
 
 		// Beyond the range of -querier.query-ingesters-within should return nothing. No need to repeat it for each user.

pkg/frontend/transport/handler.go

@@ -484,6 +484,7 @@ func writeServiceTimingHeader(queryResponseTime time.Duration, headers http.Head
 		parts := make([]string, 0)
 		parts = append(parts, statsValue("querier_wall_time", stats.LoadWallTime()))
 		parts = append(parts, statsValue("response_time", queryResponseTime))
+		parts = append(parts, statsBytesProcessedValue("bytes_processed", stats.LoadFetchedChunkBytes()+stats.LoadFetchedIndexBytes()))
 		headers.Set(ServiceTimingHeaderName, strings.Join(parts, ", "))
 	}
 }
@@ -493,6 +494,10 @@ func statsValue(name string, d time.Duration) string {
 	return name + ";dur=" + durationInMs
 }
 
+func statsBytesProcessedValue(name string, value uint64) string {
+	return name + "=" + strconv.FormatUint(value, 10)
+}
+
 func httpRequestActivity(request *http.Request, userAgent string, requestParams url.Values) string {
 	tenantID := "(unknown)"
 	if tenantIDs, err := tenant.TenantIDs(request.Context()); err == nil {

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 ServiceTimingHeader",
+			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.Contains(t, headers.Get(ServiceTimingHeaderName), "bytes_processed=0")
+			},
+		},
 	} {
 		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