Bind to existing stats tag

Signed-off-by: Pedro Tanaka <pedro.tanaka@shopify.com>
thanos-io · Nov 4, 2024 · 1265ac9 · 1265ac9
1 parent b69c9bb
commit 1265ac9
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 12 deletions.
diff --git a/cmd/thanos/query_frontend.go b/cmd/thanos/query_frontend.go
@@ -146,7 +146,7 @@ func registerQueryFrontend(app *extkingpin.App) {
 	cmd.Flag("query-frontend.log-queries-longer-than", "Log queries that are slower than the specified duration. "+
 		"Set to 0 to disable. Set to < 0 to enable on all queries.").Default("0").DurationVar(&cfg.CortexHandlerConfig.LogQueriesLongerThan)
 
-	cmd.Flag("query-frontend.force-query-stats", "Will always pass \"stats\" param to upstream queriers and collect query statistics reporting them as logs.").Default("false").BoolVar(&cfg.ForceQueryStats)
+	cmd.Flag("query-frontend.force-query-stats", "Enables query statistics for all queries and will export statistics as logs and service headers.").Default("false").BoolVar(&cfg.CortexHandlerConfig.QueryStatsEnabled)
 
 	cmd.Flag("query-frontend.org-id-header", "Deprecation Warning - This flag will be soon deprecated in favor of query-frontend.tenant-header"+
 		" and both flags cannot be used at the same time. "+

diff --git a/internal/cortex/frontend/transport/handler.go b/internal/cortex/frontend/transport/handler.go
@@ -107,9 +107,11 @@ func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	// Initialise the stats in the context and make sure it's propagated
 	// down the request chain.
-	var ctx context.Context
-	stats, ctx = querier_stats.ContextWithEmptyStats(r.Context())
-	r = r.WithContext(ctx)
+	if f.cfg.QueryStatsEnabled {
+		var ctx context.Context
+		stats, ctx = querier_stats.ContextWithEmptyStats(r.Context())
+		r = r.WithContext(ctx)
+	}
 
 	defer func() {
 		_ = r.Body.Close()
@@ -211,7 +213,7 @@ func (f *Handler) reportSlowQuery(
 		"trace_id", thanosTraceID,
 	}, formatQueryString(queryString)...)
 
-	logMessage = addQueryRangeToLogMessage(queryString, logMessage)
+	logMessage = addQueryRangeToLogMessage(logMessage, queryString)
 	logMessage = f.addStatsToLogMessage(logMessage, stats)
 
 	level.Info(util_log.WithContext(r.Context(), f.log)).Log(logMessage...)
@@ -247,6 +249,8 @@ func (f *Handler) reportQueryStats(r *http.Request, queryString url.Values, quer
 		"fetched_series_count", numSeries,
 		"fetched_chunks_bytes", numBytes,
 	}, formatQueryString(queryString)...)
+	f.addStatsToLogMessage(logMessage, stats)
+	addQueryRangeToLogMessage(logMessage, queryString)
 
 	level.Info(util_log.WithContext(r.Context(), f.log)).Log(logMessage...)
 }
@@ -281,7 +285,7 @@ func (f *Handler) addStatsToLogMessage(message []interface{}, stats *querier_sta
 	return message
 }
 
-func addQueryRangeToLogMessage(queryString url.Values, logMessage []interface{}) []interface{} {
+func addQueryRangeToLogMessage(logMessage []interface{}, queryString url.Values) []interface{} {
 	queryRange := extractQueryRange(queryString)
 	if queryRange != time.Duration(0) {
 		logMessage = append(logMessage, "query_range_hours", int(queryRange.Hours()))

diff --git a/pkg/queryfrontend/config.go b/pkg/queryfrontend/config.go
@@ -211,7 +211,6 @@ type Config struct {
 	DefaultTenant          string
 	TenantCertField        string
 	EnableXFunctions       bool
-	ForceQueryStats        bool
 }
 
 // QueryRangeConfig holds the config for query range tripperware.

diff --git a/pkg/queryfrontend/roundtrip.go b/pkg/queryfrontend/roundtrip.go
@@ -62,7 +62,7 @@ func NewTripperware(config Config, reg prometheus.Registerer, logger log.Logger)
 		queryRangeLimits,
 		queryRangeCodec,
 		config.NumShards,
-		config.ForceQueryStats,
+		config.CortexHandlerConfig.QueryStatsEnabled,
 		prometheus.WrapRegistererWith(prometheus.Labels{"tripperware": "query_range"}, reg), logger, config.ForwardHeaders)
 	if err != nil {
 		return nil, err
@@ -79,7 +79,7 @@ func NewTripperware(config Config, reg prometheus.Registerer, logger log.Logger)
 		queryInstantCodec,
 		prometheus.WrapRegistererWith(prometheus.Labels{"tripperware": "query_instant"}, reg),
 		config.ForwardHeaders,
-		config.ForceQueryStats,
+		config.CortexHandlerConfig.QueryStatsEnabled,
 	)
 	return func(next http.RoundTripper) http.RoundTripper {
 		tripper := newRoundTripper(

diff --git a/pkg/queryfrontend/roundtrip_test.go b/pkg/queryfrontend/roundtrip_test.go
@@ -13,17 +13,17 @@ import (
 	"testing"
 	"time"
 
-	"go.uber.org/atomic"
-
+	"github.com/efficientgo/core/testutil"
 	"github.com/go-kit/log"
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/promql/parser"
 	"github.com/weaveworks/common/user"
+	"go.uber.org/atomic"
 
-	"github.com/efficientgo/core/testutil"
 	cortexcache "github.com/thanos-io/thanos/internal/cortex/chunk/cache"
 	"github.com/thanos-io/thanos/internal/cortex/cortexpb"
+	"github.com/thanos-io/thanos/internal/cortex/frontend/transport"
 	"github.com/thanos-io/thanos/internal/cortex/querier/queryrange"
 	cortexvalidation "github.com/thanos-io/thanos/internal/cortex/util/validation"
 	"github.com/thanos-io/thanos/pkg/store/labelpb"
@@ -186,6 +186,7 @@ func TestRoundTripRetryMiddleware(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			tpw, err := NewTripperware(
 				Config{
+					CortexHandlerConfig: &transport.HandlerConfig{},
 					QueryRangeConfig: QueryRangeConfig{
 						MaxRetries:             tc.maxRetries,
 						Limits:                 defaultLimits,