Propagate whether query stats is enabled from query-frontend down the request path (#3595)

* Propagate whether query stats is enabled from query-frontend down the request path

Signed-off-by: Marco Pracucci <[email protected]>

* Update CHANGELOG.md

Signed-off-by: Marco Pracucci <[email protected]>

Co-authored-by: Peter Štibraný <[email protected]>

Co-authored-by: Peter Štibraný <[email protected]>

Author: pracucci

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## master / unreleased
 
+* [CHANGE] Querier: it's not required to set `-frontend.query-stats-enabled=true` in the querier anymore to enable query statistics logging in the query-frontend. The flag is now required to be configured only in the query-frontend and it will be propagated to the queriers. #3595
 * [CHANGE] Blocks storage: compactor is now required when running a Cortex cluster with the blocks storage, because it also keeps the bucket index updated. #3583
 * [CHANGE] Blocks storage: block deletion marks are now stored in a per-tenant global markers/ location too, other than within the block location. The compactor, at startup, will copy deletion marks from the block location to the global location. This migration is required only once, so you can safely disable it via `-compactor.block-deletion-marks-migration-enabled=false` once new compactor has successfully started once in your cluster. #3583
 * [ENHANCEMENT] Blocks storage: introduced a per-tenant bucket index, periodically updated by the compactor, used to avoid full bucket scanning done by queriers and store-gateways. The bucket index is updated by the compactor during blocks cleanup, on every `-compactor.cleanup-interval`. #3553 #3555 #3561 #3583

docs/configuration/config-file-reference.md

@@ -877,8 +877,7 @@ The `query_frontend_config` configures the Cortex query-frontend.
 [max_body_size: <int> | default = 10485760]
 
 # True to enable query statistics tracking. When enabled, a message with some
-# statistics is logged for every query. This configuration option must be set
-# both on query-frontend and querier.
+# statistics is logged for every query.
 # CLI flag: -frontend.query-stats-enabled
 [query_stats_enabled: <boolean> | default = false]
 

pkg/cortex/modules.go

@@ -306,7 +306,6 @@ func (t *Cortex) initQuerier() (serv services.Service, err error) {
 	}
 
 	t.Cfg.Worker.MaxConcurrentRequests = t.Cfg.Querier.MaxConcurrent
-	t.Cfg.Worker.QueryStatsEnabled = t.Cfg.Frontend.Handler.QueryStatsEnabled
 	return querier_worker.NewQuerierWorker(t.Cfg.Worker, httpgrpc_server.NewServer(internalQuerierRouter), util.Logger, prometheus.DefaultRegisterer)
 }
 

pkg/frontend/transport/handler.go

@@ -45,7 +45,7 @@ type HandlerConfig struct {
 func (cfg *HandlerConfig) RegisterFlags(f *flag.FlagSet) {
 	f.DurationVar(&cfg.LogQueriesLongerThan, "frontend.log-queries-longer-than", 0, "Log queries that are slower than the specified duration. Set to 0 to disable. Set to < 0 to enable on all queries.")
 	f.Int64Var(&cfg.MaxBodySize, "frontend.max-body-size", 10*1024*1024, "Max body size for downstream prometheus.")
-	f.BoolVar(&cfg.QueryStatsEnabled, "frontend.query-stats-enabled", false, "True to enable query statistics tracking. When enabled, a message with some statistics is logged for every query. This configuration option must be set both on query-frontend and querier.")
+	f.BoolVar(&cfg.QueryStatsEnabled, "frontend.query-stats-enabled", false, "True to enable query statistics tracking. When enabled, a message with some statistics is logged for every query.")
 }
 
 // Handler accepts queries and forwards them to RoundTripper. It can log slow queries,

pkg/frontend/v1/frontend.go

@@ -190,8 +190,9 @@ func (f *Frontend) Process(server frontendv1pb.Frontend_ProcessServer) error {
 		errs := make(chan error, 1)
 		go func() {
 			err = server.Send(&frontendv1pb.FrontendToClient{
-				Type:        frontendv1pb.HTTP_REQUEST,
-				HttpRequest: req.request,
+				Type:         frontendv1pb.HTTP_REQUEST,
+				HttpRequest:  req.request,
+				StatsEnabled: stats.IsEnabled(req.originalCtx),
 			})
 			if err != nil {
 				errs <- err

pkg/frontend/v1/frontendv1pb/frontend.pb.go

@@ -27,6 +27,10 @@ enum Type {
 message FrontendToClient {
   httpgrpc.HTTPRequest httpRequest = 1;
   Type type = 2;
+
+  // Whether query statistics tracking should be enabled. The response will include
+  // statistics only when this option is enabled.
+  bool statsEnabled = 3;
 }
 
 message ClientToFrontend {

pkg/frontend/v1/frontendv1pb/frontend.proto

@@ -73,9 +73,10 @@ type Frontend struct {
 }
 
 type frontendRequest struct {
-	queryID uint64
-	request *httpgrpc.HTTPRequest
-	userID  string
+	queryID      uint64
+	request      *httpgrpc.HTTPRequest
+	userID       string
+	statsEnabled bool
 
 	cancel context.CancelFunc
 
@@ -170,9 +171,10 @@ func (f *Frontend) RoundTripGRPC(ctx context.Context, req *httpgrpc.HTTPRequest)
 	defer cancel()
 
 	freq := &frontendRequest{
-		queryID: f.lastQueryID.Inc(),
-		request: req,
-		userID:  userID,
+		queryID:      f.lastQueryID.Inc(),
+		request:      req,
+		userID:       userID,
+		statsEnabled: stats.IsEnabled(ctx),
 
 		cancel: cancel,
 

pkg/frontend/v2/frontend.go

@@ -261,6 +261,7 @@ func (w *frontendSchedulerWorker) schedulerLoop(loop schedulerpb.SchedulerForFro
 				UserID:          req.userID,
 				HttpRequest:     req.request,
 				FrontendAddress: w.frontendAddr,
+				StatsEnabled:    req.statsEnabled,
 			})
 
 			if err != nil {

pkg/frontend/v2/frontend_scheduler_worker.go

@@ -29,6 +29,13 @@ func FromContext(ctx context.Context) *Stats {
 	return o.(*Stats)
 }
 
+// IsEnabled returns whether stats tracking is enabled in the context.
+func IsEnabled(ctx context.Context) bool {
+	// When query statistics are enabled, the stats object is already initialised
+	// within the context, so we can just check it.
+	return FromContext(ctx) != nil
+}
+
 // AddWallTime adds some time to the counter.
 func (s *Stats) AddWallTime(t time.Duration) {
 	if s == nil {

pkg/querier/stats/stats.go

@@ -16,14 +16,15 @@ func NewWallTimeMiddleware() WallTimeMiddleware {
 // Wrap implements middleware.Interface.
 func (m WallTimeMiddleware) Wrap(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		stats := FromContext(r.Context())
-		if stats == nil {
+		if !IsEnabled(r.Context()) {
 			next.ServeHTTP(w, r)
 			return
 		}
 
 		startTime := time.Now()
 		next.ServeHTTP(w, r)
+
+		stats := FromContext(r.Context())
 		stats.AddWallTime(time.Since(startTime))
 	})
 }

pkg/querier/stats/time_middleware.go

@@ -26,20 +26,18 @@ var (
 
 func newFrontendProcessor(cfg Config, handler RequestHandler, log log.Logger) processor {
 	return &frontendProcessor{
-		log:               log,
-		handler:           handler,
-		maxMessageSize:    cfg.GRPCClientConfig.GRPC.MaxSendMsgSize,
-		querierID:         cfg.QuerierID,
-		queryStatsEnabled: cfg.QueryStatsEnabled,
+		log:            log,
+		handler:        handler,
+		maxMessageSize: cfg.GRPCClientConfig.GRPC.MaxSendMsgSize,
+		querierID:      cfg.QuerierID,
 	}
 }
 
 // Handles incoming queries from frontend.
 type frontendProcessor struct {
-	handler           RequestHandler
-	maxMessageSize    int
-	querierID         string
-	queryStatsEnabled bool
+	handler        RequestHandler
+	maxMessageSize int
+	querierID      string
 
 	log log.Logger
 }
@@ -86,7 +84,7 @@ func (fp *frontendProcessor) process(c frontendv1pb.Frontend_ProcessClient) erro
 			// and cancel the query.  We don't actually handle queries in parallel
 			// here, as we're running in lock step with the server - each Recv is
 			// paired with a Send.
-			go fp.runRequest(ctx, request.HttpRequest, func(response *httpgrpc.HTTPResponse, stats *stats.Stats) error {
+			go fp.runRequest(ctx, request.HttpRequest, request.StatsEnabled, func(response *httpgrpc.HTTPResponse, stats *stats.Stats) error {
 				return c.Send(&frontendv1pb.ClientToFrontend{
 					HttpResponse: response,
 					Stats:        stats,
@@ -105,9 +103,9 @@ func (fp *frontendProcessor) process(c frontendv1pb.Frontend_ProcessClient) erro
 	}
 }
 
-func (fp *frontendProcessor) runRequest(ctx context.Context, request *httpgrpc.HTTPRequest, sendHTTPResponse func(response *httpgrpc.HTTPResponse, stats *stats.Stats) error) {
+func (fp *frontendProcessor) runRequest(ctx context.Context, request *httpgrpc.HTTPRequest, statsEnabled bool, sendHTTPResponse func(response *httpgrpc.HTTPResponse, stats *stats.Stats) error) {
 	var stats *querier_stats.Stats
-	if fp.queryStatsEnabled {
+	if statsEnabled {
 		stats, ctx = querier_stats.ContextWithEmptyStats(ctx)
 	}