add bytes fetched in query frontend

databricks · May 17, 2024 · dbdda52 · dbdda52
1 parent ca8da96
commit dbdda52
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 5 deletions.
diff --git a/internal/cortex/querier/queryrange/query_bytes_fetched.go b/internal/cortex/querier/queryrange/query_bytes_fetched.go
@@ -0,0 +1,44 @@
+package queryrange
+
+import (
+	"strconv"
+)
+
+const QueryBytesFetchedHeaderName = "M3-Fetched-Bytes-Estimate"
+
+func SumQueryBytesFetched(responses ...Response) int {
+	result := 0
+	for _, resp := range responses {
+		for _, hdr := range resp.GetHeaders() {
+			if hdr.GetName() == QueryBytesFetchedHeaderName {
+				for _, v := range hdr.Values {
+					n, err := strconv.Atoi(v)
+					if err != nil {
+						continue
+					}
+					result += n
+				}
+				break
+			}
+		}
+	}
+	return result
+}
+
+func QueryBytesFetchedPrometheusResponseHeaders(responses ...Response) []*PrometheusResponseHeader {
+	return []*PrometheusResponseHeader{{
+		Name:   QueryBytesFetchedHeaderName,
+		Values: []string{strconv.Itoa(SumQueryBytesFetched(responses...))},
+	}}
+}
+
+func QueryBytesFetchedHttpHeaderValue(response Response) []string {
+	result := []string{}
+	for _, hdr := range response.GetHeaders() {
+		if hdr.GetName() == QueryBytesFetchedHeaderName {
+			result = hdr.GetValues()
+			break
+		}
+	}
+	return result
+}
diff --git a/internal/cortex/querier/queryrange/query_range.go b/internal/cortex/querier/queryrange/query_range.go
@@ -278,12 +278,12 @@ func (prometheusCodec) MergeResponse(_ Request, responses ...Response) (Response
 			Analysis:   AnalyzesMerge(analyzes...),
 		},
 	}
-
+	response.Headers = QueryBytesFetchedPrometheusResponseHeaders(responses...)
 	if len(resultsCacheGenNumberHeaderValues) != 0 {
-		response.Headers = []*PrometheusResponseHeader{{
+		response.Headers = append(response.Headers, &PrometheusResponseHeader{
 			Name:   ResultsCacheGenNumberHeaderName,
 			Values: resultsCacheGenNumberHeaderValues,
-		}}
+		})
 	}
 
 	return &response, nil
@@ -449,7 +449,8 @@ func (prometheusCodec) EncodeResponse(ctx context.Context, res Response) (*http.
 
 	resp := http.Response{
 		Header: http.Header{
-			"Content-Type": []string{"application/json"},
+			"Content-Type":              []string{"application/json"},
+			QueryBytesFetchedHeaderName: QueryBytesFetchedHttpHeaderValue(res),
 		},
 		Body:          io.NopCloser(bytes.NewBuffer(b)),
 		StatusCode:    http.StatusOK,

diff --git a/pkg/queryfrontend/queryinstant_codec.go b/pkg/queryfrontend/queryinstant_codec.go
@@ -78,6 +78,7 @@ func (c queryInstantCodec) MergeResponse(req queryrange.Request, responses ...qu
 				Analysis: queryrange.AnalyzesMerge(analyzes...),
 				Stats:    queryrange.StatsMerge(responses),
 			},
+			Headers: queryrange.QueryBytesFetchedPrometheusResponseHeaders(responses...),
 		}
 	default:
 		v, err := vectorMerge(req, promResponses)
@@ -96,6 +97,7 @@ func (c queryInstantCodec) MergeResponse(req queryrange.Request, responses ...qu
 				Analysis: queryrange.AnalyzesMerge(analyzes...),
 				Stats:    queryrange.StatsMerge(responses),
 			},
+			Headers: queryrange.QueryBytesFetchedPrometheusResponseHeaders(responses...),
 		}
 	}
 
@@ -248,7 +250,8 @@ func (c queryInstantCodec) EncodeResponse(ctx context.Context, res queryrange.Re
 
 	resp := http.Response{
 		Header: http.Header{
-			"Content-Type": []string{"application/json"},
+			"Content-Type":                         []string{"application/json"},
+			queryrange.QueryBytesFetchedHeaderName: queryrange.QueryBytesFetchedHttpHeaderValue(res),
 		},
 		Body:          io.NopCloser(bytes.NewBuffer(b)),
 		StatusCode:    http.StatusOK,