From e2e695e8cdad15f3e9099ed628d72c9b35e533fe Mon Sep 17 00:00:00 2001
From: Travis Patterson <travis.patterson@grafana.com>
Date: Thu, 6 Jul 2023 12:29:11 -0600
Subject: [PATCH] Get basic statistics from series volume requests (#9832)

The `series_volume` endpoint returns stats about the call. This PR wires
things up such that basic stats are returned:
- Execution time
- Number of responses
---
 pkg/querier/queryrange/codec.go      | 10 ++++++++++
 pkg/querier/queryrange/roundtrip.go  |  1 +
 pkg/querier/queryrange/stats_test.go | 21 +++++++++++++++++++++
 3 files changed, 32 insertions(+)

diff --git a/pkg/querier/queryrange/codec.go b/pkg/querier/queryrange/codec.go
index 886b30e9d6..a04514b45a 100644
--- a/pkg/querier/queryrange/codec.go
+++ b/pkg/querier/queryrange/codec.go
@@ -1055,6 +1055,16 @@ func paramsFromRequest(req queryrangebase.Request) (logql.Params, error) {
 		return &paramsRangeWrapper{
 			LokiRequest: r,
 		}, nil
+	case *logproto.VolumeRequest:
+		return &paramsRangeWrapper{
+			LokiRequest: &LokiRequest{
+				Query:   r.GetQuery(),
+				Limit:   uint32(r.GetLimit()),
+				Step:    r.GetStep(),
+				StartTs: time.UnixMilli(r.GetStart()),
+				EndTs:   time.UnixMilli(r.GetEnd()),
+			},
+		}, nil
 	case *LokiInstantRequest:
 		return &paramsInstantWrapper{
 			LokiInstantRequest: r,
diff --git a/pkg/querier/queryrange/roundtrip.go b/pkg/querier/queryrange/roundtrip.go
index 0103c17542..179d466b5a 100644
--- a/pkg/querier/queryrange/roundtrip.go
+++ b/pkg/querier/queryrange/roundtrip.go
@@ -807,6 +807,7 @@ func volumeRangeTripperware(codec queryrangebase.Codec, nextTW queryrangebase.Tr
 			}
 
 			seriesVolumeMiddlewares := []queryrangebase.Middleware{
+				StatsCollectorMiddleware(),
 				NewSeriesVolumeMiddleware(),
 			}
 
diff --git a/pkg/querier/queryrange/stats_test.go b/pkg/querier/queryrange/stats_test.go
index a8e31df85d..c82e121a72 100644
--- a/pkg/querier/queryrange/stats_test.go
+++ b/pkg/querier/queryrange/stats_test.go
@@ -158,6 +158,27 @@ func Test_StatsHTTP(t *testing.T) {
 				require.Equal(t, streams, data.result)
 			},
 		},
+		{
+			"volume request",
+			http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				data := r.Context().Value(ctxKey).(*queryData)
+				data.recorded = true
+				data.params, _ = paramsFromRequest(&logproto.VolumeRequest{
+					Matchers: "foo",
+					Limit:    100,
+				})
+				data.statistics = &statsResult
+				data.result = streams
+				w.WriteHeader(http.StatusTeapot)
+			}),
+			func(t *testing.T, data *queryData) {
+				require.Equal(t, fmt.Sprintf("%d", http.StatusTeapot), data.status)
+				require.Equal(t, "foo", data.params.Query())
+				require.Equal(t, uint32(100), data.params.Limit())
+				require.Equal(t, statsResult, *data.statistics)
+				require.Equal(t, streams, data.result)
+			},
+		},
 	} {
 		t.Run(test.name, func(t *testing.T) {
 			statsHTTPMiddleware(metricRecorderFn(func(data *queryData) {