populate stats

Author: yuchen-db

internal/cortex/frontend/transport/handler.go

@@ -21,6 +21,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/thanos-io/thanos/internal/cortex/frontend/transport/utils"
+	"github.com/thanos-io/thanos/internal/cortex/querier/queryrange"
 	querier_stats "github.com/thanos-io/thanos/internal/cortex/querier/stats"
 	"github.com/thanos-io/thanos/internal/cortex/tenant"
 	"github.com/thanos-io/thanos/internal/cortex/util"
@@ -212,6 +213,8 @@ func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		f.reportSlowQuery(r, hs, queryString, queryResponseTime)
 	}
 	if f.cfg.QueryStatsEnabled {
+		stats.FetchedChunkBytes = queryrange.GetQueryBytesFetchedFromHeader(resp.Header)
+		stats.FetchedSeriesCount = queryrange.GetQuerySeriesFetchedFromHeader(resp.Header)
 		f.reportQueryStats(r, queryString, queryResponseTime, stats)
 	}
 }

internal/cortex/querier/queryrange/query_bytes_fetched.go

@@ -4,12 +4,14 @@
 package queryrange
 
 import (
+	"net/http"
 	"strconv"
 )
 
 // QueryBytesFetchedHeaderName is the http header name of number of bytes fetched by a query from m3readcoord.
 // This name is compatible with M3 and rule manager code
 const QueryBytesFetchedHeaderName = "M3-Fetched-Bytes-Estimate"
+const QuerySeriesFetchedHeaderName = "M3-Fetched-Series-Estimate"
 
 func sumQueryBytesFetched(responses ...Response) uint64 {
 	var result uint64
@@ -52,3 +54,25 @@ func QueryBytesFetchedHttpHeaderValue(response Response) []string {
 	}
 	return result
 }
+
+func getHeaderValue(hdr http.Header, key string) uint64 {
+	if val, ok := hdr[key]; ok {
+		if len(val) != 1 {
+			return 0
+		}
+		n, err := strconv.ParseUint(val[0], 10, 64)
+		if err != nil {
+			return 0
+		}
+		return n
+	}
+	return 0
+}
+
+func GetQueryBytesFetchedFromHeader(hdr http.Header) uint64 {
+	return getHeaderValue(hdr, QueryBytesFetchedHeaderName)
+}
+
+func GetQuerySeriesFetchedFromHeader(hdr http.Header) uint64 {
+	return getHeaderValue(hdr, QuerySeriesFetchedHeaderName)
+}

internal/cortex/querier/queryrange/query_range.go

@@ -483,6 +483,7 @@ func (prometheusCodec) EncodeResponse(ctx context.Context, res Response) (*http.
 	} else if res.(*PrometheusResponse).Data.SeriesStatsCounter != nil {
 		// Pantheon code path
 		httpHeader[QueryBytesFetchedHeaderName] = []string{strconv.FormatInt(res.(*PrometheusResponse).Data.SeriesStatsCounter.Bytes, 10)}
+		httpHeader[QuerySeriesFetchedHeaderName] = []string{strconv.FormatInt(res.(*PrometheusResponse).Data.SeriesStatsCounter.Series, 10)}
 	}
 	resp := http.Response{
 		Header:        httpHeader,