QFE: pass "stats" parameter to queriers, fix prometheus stats merging

CHANGELOG.md

@@ -23,6 +23,7 @@ We use *breaking :warning:* to mark changes that are not backward compatible (re
 - [#7679](https://github.com/thanos-io/thanos/pull/7679) Query: respect store.limit.* flags when evaluating queries
 - [#7821](https://github.com/thanos-io/thanos/pull/7679) Query/Receive: Fix coroutine leak introduced in https://github.com/thanos-io/thanos/pull/7796.
 - [#7843](https://github.com/thanos-io/thanos/pull/7843) Query Frontend: fix slow query logging for non-query endpoints.
+- [#7852](https://github.com/thanos-io/thanos/pull/7852) Query Frontend: pass "stats" parameter forward to queriers and fix Prometheus stats merging.
 
 ### Added
 - [#7763](https://github.com/thanos-io/thanos/pull/7763) Ruler: use native histograms for client latency metrics.

internal/cortex/querier/queryrange/query_range.go

@@ -8,7 +8,7 @@ import (
 	"context"
 	stdjson "encoding/json"
 	"fmt"
-	io "io"
+	"io"
 	"math"
 	"net/http"
 	"net/url"
@@ -329,6 +329,7 @@ func (prometheusCodec) DecodeRequest(_ context.Context, r *http.Request, forward
 	result.Query = r.FormValue("query")
 	result.Stats = r.FormValue("stats")
 	result.Path = r.URL.Path
+	result.Stats = r.FormValue("stats")
 
 	// Include the specified headers from http request in prometheusRequest.
 	for _, header := range forwardHeaders {
@@ -706,6 +707,8 @@ func (s *PrometheusInstantQueryData) MarshalJSON() ([]byte, error) {
 func StatsMerge(resps []Response) *PrometheusResponseStats {
 	output := map[int64]*PrometheusResponseQueryableSamplesStatsPerStep{}
 	hasStats := false
+	peakSamples := int32(0)
+	totalSamples := int64(0)
 	for _, resp := range resps {
 		stats := resp.GetStats()
 		if stats == nil {
@@ -720,6 +723,11 @@ func StatsMerge(resps []Response) *PrometheusResponseStats {
 		for _, s := range stats.Samples.TotalQueryableSamplesPerStep {
 			output[s.GetTimestampMs()] = s
 		}
+
+		if stats.Samples.PeakSamples > peakSamples {
+			peakSamples = stats.Samples.PeakSamples
+		}
+		totalSamples += stats.Samples.TotalQueryableSamples
 	}
 
 	if !hasStats {
@@ -733,10 +741,12 @@ func StatsMerge(resps []Response) *PrometheusResponseStats {
 
 	sort.Slice(keys, func(i, j int) bool { return keys[i] < keys[j] })
 
-	result := &PrometheusResponseStats{Samples: &PrometheusResponseSamplesStats{}}
+	result := &PrometheusResponseStats{Samples: &PrometheusResponseSamplesStats{
+		PeakSamples:           peakSamples,
+		TotalQueryableSamples: totalSamples,
+	}}
 	for _, key := range keys {
 		result.Samples.TotalQueryableSamplesPerStep = append(result.Samples.TotalQueryableSamplesPerStep, output[key])
-		result.Samples.TotalQueryableSamples += output[key].Value
 	}
 
 	return result

internal/cortex/querier/queryrange/query_range_test.go

@@ -6,7 +6,7 @@ package queryrange
 import (
 	"bytes"
 	"context"
-	io "io"
+	"io"
 	"net/http"
 	"strconv"
 	"testing"
@@ -36,6 +36,20 @@ func TestRequest(t *testing.T) {
 			url:      query,
 			expected: &parsedRequestWithHeaders,
 		},
+		{
+			url: "/api/v1/query_range?end=60&query=sum%28container_memory_rss%29+by+%28namespace%29&start=0&stats=all&step=14",
+			expected: &PrometheusRequest{
+				Path:  "/api/v1/query_range",
+				Start: 0,
+				End:   60_000,
+				Step:  14_000,
+				Query: "sum(container_memory_rss) by (namespace)",
+				Stats: "all",
+				Headers: []*PrometheusRequestHeader{
+					{Name: "Test-Header", Values: []string{"test"}},
+				},
+			},
+		},
 		{
 			url:         "api/v1/query_range?start=foo&stats=all",
 			expectedErr: httpgrpc.Errorf(http.StatusBadRequest, "invalid parameter \"start\"; cannot parse \"foo\" to a valid timestamp"),
@@ -129,7 +143,7 @@ func TestResponseWithStats(t *testing.T) {
 		expected *PrometheusResponse
 	}{
 		{
-			body: `{"status":"success","data":{"resultType":"matrix","result":[{"metric":{"foo":"bar"},"values":[[1536673680,"137"],[1536673780,"137"]]}],"stats":{"samples":{"totalQueryableSamples":10,"totalQueryableSamplesPerStep":[[1536673680,5],[1536673780,5]]}},"analysis":null}}`,
+			body: `{"status":"success","data":{"resultType":"matrix","result":[{"metric":{"foo":"bar"},"values":[[1536673680,"137"],[1536673780,"137"]]}],"stats":{"samples":{"peakSamples":0,"totalQueryableSamples":10,"totalQueryableSamplesPerStep":[[1536673680,5],[1536673780,5]]}},"analysis":null}}`,
 			expected: &PrometheusResponse{
 				Status: "success",
 				Data: PrometheusData{
@@ -158,7 +172,7 @@ func TestResponseWithStats(t *testing.T) {
 			},
 		},
 		{
-			body: `{"status":"success","data":{"resultType":"matrix","result":[{"metric":{"foo":"bar"},"values":[[1536673680,"137"],[1536673780,"137"]]}],"stats":{"samples":{"totalQueryableSamples":10,"totalQueryableSamplesPerStep":[[1536673680,5],[1536673780,5]]}},"analysis":{"name":"[noArgFunction]","executionTime":"1s","children":null}}}`,
+			body: `{"status":"success","data":{"resultType":"matrix","result":[{"metric":{"foo":"bar"},"values":[[1536673680,"137"],[1536673780,"137"]]}],"stats":{"samples":{"peakSamples":0,"totalQueryableSamples":10,"totalQueryableSamplesPerStep":[[1536673680,5],[1536673780,5]]}},"analysis":{"name":"[noArgFunction]","executionTime":"1s","children":null}}}`,
 			expected: &PrometheusResponse{
 				Status: "success",
 				Data: PrometheusData{
@@ -211,11 +225,27 @@ func TestResponseWithStats(t *testing.T) {
 			}
 			resp2, err := PrometheusCodec.EncodeResponse(context.Background(), resp)
 			require.NoError(t, err)
-			assert.Equal(t, response, resp2)
+			assert.Equal(t, prettyPrintJsonBody(t, response.Body), prettyPrintJsonBody(t, resp2.Body))
 		})
 	}
 }
 
+func prettyPrintJsonBody(t *testing.T, body io.ReadCloser) string {
+	t.Helper()
+
+	bodyContent, err := io.ReadAll(body)
+	require.NoError(t, err)
+
+	var jsonData interface{}
+	err = json.Unmarshal(bodyContent, &jsonData)
+	require.NoError(t, err)
+
+	prettyBytes, err := json.MarshalIndent(jsonData, "", "  ")
+	require.NoError(t, err)
+
+	return string(prettyBytes)
+}
+
 func TestMergeAPIResponses(t *testing.T) {
 	for _, tc := range []struct {
 		name     string
@@ -659,7 +689,7 @@ func TestMergeAPIResponses(t *testing.T) {
 						},
 					},
 					Stats: &PrometheusResponseStats{Samples: &PrometheusResponseSamplesStats{
-						TotalQueryableSamples: 25,
+						TotalQueryableSamples: 30,
 						TotalQueryableSamplesPerStep: []*PrometheusResponseQueryableSamplesStatsPerStep{
 							{Value: 5, TimestampMs: 1000},
 							{Value: 5, TimestampMs: 2000},
@@ -696,7 +726,7 @@ func TestMergeAPIResponses(t *testing.T) {
 						},
 					},
 					Stats: &PrometheusResponseStats{Samples: &PrometheusResponseSamplesStats{
-						TotalQueryableSamples: 28,
+						TotalQueryableSamples: 36,
 						TotalQueryableSamplesPerStep: []*PrometheusResponseQueryableSamplesStatsPerStep{
 							{Value: 1, TimestampMs: 1000},
 							{Value: 2, TimestampMs: 2000},
@@ -734,7 +764,7 @@ func TestMergeAPIResponses(t *testing.T) {
 						},
 					},
 					Stats: &PrometheusResponseStats{Samples: &PrometheusResponseSamplesStats{
-						TotalQueryableSamples: 15,
+						TotalQueryableSamples: 26,
 						TotalQueryableSamplesPerStep: []*PrometheusResponseQueryableSamplesStatsPerStep{
 							{Value: 1, TimestampMs: 1000},
 							{Value: 2, TimestampMs: 2000},
@@ -769,7 +799,7 @@ func TestMergeAPIResponses(t *testing.T) {
 						},
 					},
 					Stats: &PrometheusResponseStats{Samples: &PrometheusResponseSamplesStats{
-						TotalQueryableSamples: 14,
+						TotalQueryableSamples: 40,
 						TotalQueryableSamplesPerStep: []*PrometheusResponseQueryableSamplesStatsPerStep{
 							{Value: 2, TimestampMs: 2000},
 							{Value: 3, TimestampMs: 3000},