Move priority, mint and maxt to query stats

Signed-off-by: Justin Jung <[email protected]>

Author: justinjung04

CHANGELOG.md

@@ -16,7 +16,7 @@
 * [ENHANCEMENT] Index Cache: Multi level cache adds config `max_backfill_items` to cap max items to backfill per async operation. #5686
 * [ENHANCEMENT] Query Frontend: Log number of split queries in `query stats` log. #5703
 * [ENHANCEMENT] Logging: Added new options for logging HTTP request headers: `-server.log-request-headers` enables logging HTTP request headers, `-server.log-request-headers-exclude-list` allows users to specify headers which should not be logged. #5744
-* [ENHANCEMENT] Query Frontend/Scheduler: Time check in query priority now considers overall data fetched time window (including range selectors, modifiers and lookback delta). #5758
+* [ENHANCEMENT] Query Frontend/Scheduler: Time check in query priority now considers overall data select time window (including range selectors, modifiers and lookback delta). #5758
 * [ENHANCEMENT] Querier: Added `querier.store-gateway-query-stats-enabled` to enable or disable store gateway query stats log. #5749
 * [BUGFIX] Distributor: Do not use label with empty values for sharding #5717
 * [BUGFIX] Query Frontend: queries with negative offset should check whether it is cacheable or not. #5719

docs/configuration/config-file-reference.md

@@ -5095,16 +5095,16 @@ otel:
 # Regex that the query string should match. If not set, it won't be checked.
 [regex: <string> | default = ""]
 
-# Overall data fetch time window (including range selectors, modifiers and
+# Overall data select time window (including range selectors, modifiers and
 # lookback delta) that the query should be within. If not set, it won't be
 # checked.
 time_window:
-  # Start of the data fetch time window (including range selectors, modifiers
+  # Start of the data select time window (including range selectors, modifiers
   # and lookback delta) that the query should be within. If set to 0, it won't
   # be checked.
   [start: <int> | default = 0]
 
-  # End of the data fetch time window (including range selectors, modifiers and
+  # End of the data select time window (including range selectors, modifiers and
   # lookback delta) that the query should be within. If set to 0, it won't be
   # checked.
   [end: <int> | default = 0]

pkg/frontend/transport/handler.go

@@ -295,6 +295,8 @@ func (f *Handler) reportQueryStats(r *http.Request, userID string, queryString u
 	numChunkBytes := stats.LoadFetchedChunkBytes()
 	numDataBytes := stats.LoadFetchedDataBytes()
 	splitQueries := stats.LoadSplitQueries()
+	dataSelectMaxTime := stats.LoadDataSelectMaxTime()
+	dataSelectMinTime := stats.LoadDataSelectMinTime()
 
 	// Track stats.
 	f.querySeconds.WithLabelValues(userID).Add(wallTime.Seconds())
@@ -339,20 +341,20 @@ func (f *Handler) reportQueryStats(r *http.Request, userID string, queryString u
 		logMessage = append(logMessage, "content_encoding", encoding)
 	}
 
+	if dataSelectMaxTime > 0 {
+		logMessage = append(logMessage, "data_select_max_time", util.FormatMillisToSeconds(dataSelectMaxTime))
+	}
+	if dataSelectMinTime > 0 {
+		logMessage = append(logMessage, "data_select_min_time", util.FormatMillisToSeconds(dataSelectMinTime))
+	}
 	if query := queryString.Get("query"); len(query) > 0 {
 		logMessage = append(logMessage, "query_length", len(query))
 	}
 	if ua := r.Header.Get("User-Agent"); len(ua) > 0 {
 		logMessage = append(logMessage, "user_agent", ua)
 	}
-	if queryPriority, ok := r.Context().Value(tripperware.QueryPriorityCtxKey).(int64); ok {
-		logMessage = append(logMessage, "priority", queryPriority)
-	}
-	if maxT, ok := r.Context().Value(tripperware.DataFetchedMaxTimeCtxKey).(int64); ok {
-		logMessage = append(logMessage, "data_fetched_max_time", util.FormatMillisToSeconds(maxT))
-	}
-	if minT, ok := r.Context().Value(tripperware.DataFetchedMinTimeCtxKey).(int64); ok {
-		logMessage = append(logMessage, "data_fetched_min_time", util.FormatMillisToSeconds(minT))
+	if priority, ok := stats.LoadPriority(); ok {
+		logMessage = append(logMessage, "priority", priority)
 	}
 
 	if error != nil {

pkg/frontend/transport/handler_test.go

@@ -21,7 +21,6 @@ import (
 	"github.com/weaveworks/common/user"
 
 	querier_stats "github.com/cortexproject/cortex/pkg/querier/stats"
-	"github.com/cortexproject/cortex/pkg/querier/tripperware"
 )
 
 type roundTripperFunc func(*http.Request) (*http.Response, error)
@@ -351,29 +350,25 @@ func TestReportQueryStatsFormat(t *testing.T) {
 		},
 		"should include query priority": {
 			queryString: url.Values(map[string][]string{"query": {"up"}}),
-			priority:    99,
+			queryStats: &querier_stats.QueryStats{
+				Priority:         99,
+				PriorityAssigned: true,
+			},
 			expectedLog: `level=info msg="query stats" component=query-frontend method=GET path=/prometheus/api/v1/query response_time=1s query_wall_time_seconds=0 fetched_series_count=0 fetched_chunks_count=0 fetched_samples_count=0 fetched_chunks_bytes=0 fetched_data_bytes=0 split_queries=0 status_code=200 response_size=1000 query_length=2 priority=99 param_query=up`,
 		},
 		"should include data fetch min and max time": {
 			queryString: url.Values(map[string][]string{"query": {"up"}}),
-			minT:        1704067200000,
-			maxT:        1704153600000,
-			expectedLog: `level=info msg="query stats" component=query-frontend method=GET path=/prometheus/api/v1/query response_time=1s query_wall_time_seconds=0 fetched_series_count=0 fetched_chunks_count=0 fetched_samples_count=0 fetched_chunks_bytes=0 fetched_data_bytes=0 split_queries=0 status_code=200 response_size=1000 query_length=2 maxt=1704153600 mint=1704067200 param_query=up`,
+			queryStats: &querier_stats.QueryStats{
+				DataSelectMaxTime: 1704153600000,
+				DataSelectMinTime: 1704067200000,
+			},
+			expectedLog: `level=info msg="query stats" component=query-frontend method=GET path=/prometheus/api/v1/query response_time=1s query_wall_time_seconds=0 fetched_series_count=0 fetched_chunks_count=0 fetched_samples_count=0 fetched_chunks_bytes=0 fetched_data_bytes=0 split_queries=0 status_code=200 response_size=1000 data_select_max_time=1704153600 data_select_min_time=1704067200 query_length=2 param_query=up`,
 		},
 	}
 
 	for testName, testData := range tests {
 		t.Run(testName, func(t *testing.T) {
 			req.Header = testData.header
-			if testData.priority > 0 {
-				*req = *req.WithContext(context.WithValue(req.Context(), tripperware.QueryPriorityCtxKey, testData.priority))
-			}
-			if testData.minT > 0 {
-				*req = *req.WithContext(context.WithValue(req.Context(), tripperware.DataFetchedMinTimeCtxKey, testData.minT))
-			}
-			if testData.maxT > 0 {
-				*req = *req.WithContext(context.WithValue(req.Context(), tripperware.DataFetchedMaxTimeCtxKey, testData.maxT))
-			}
 			handler.reportQueryStats(req, userID, testData.queryString, responseTime, testData.queryStats, testData.responseErr, statusCode, resp)
 			data, err := io.ReadAll(outputBuf)
 			require.NoError(t, err)

pkg/frontend/v1/frontend.go

@@ -18,7 +18,6 @@ import (
 	"github.com/cortexproject/cortex/pkg/frontend/transport"
 	"github.com/cortexproject/cortex/pkg/frontend/v1/frontendv1pb"
 	"github.com/cortexproject/cortex/pkg/querier/stats"
-	"github.com/cortexproject/cortex/pkg/querier/tripperware"
 	"github.com/cortexproject/cortex/pkg/scheduler/queue"
 	"github.com/cortexproject/cortex/pkg/tenant"
 	"github.com/cortexproject/cortex/pkg/util"
@@ -100,7 +99,7 @@ type request struct {
 }
 
 func (r request) Priority() int64 {
-	priority, ok := r.originalCtx.Value(tripperware.QueryPriorityCtxKey).(int64)
+	priority, ok := stats.FromContext(r.originalCtx).LoadPriority()
 	if !ok {
 		return 0
 	}

pkg/querier/stats/stats.go

@@ -16,7 +16,11 @@ var ctxKey = contextKey(0)
 
 type QueryStats struct {
 	Stats
-	m sync.Mutex
+	PriorityAssigned  bool
+	Priority          int64
+	DataSelectMaxTime int64
+	DataSelectMinTime int64
+	m                 sync.Mutex
 }
 
 // ContextWithEmptyStats returns a context with empty stats.
@@ -196,6 +200,58 @@ func (s *QueryStats) LoadSplitQueries() uint64 {
 	return atomic.LoadUint64(&s.SplitQueries)
 }
 
+func (s *QueryStats) SetPriority(priority int64) {
+	if s == nil {
+		return
+	}
+
+	if !s.PriorityAssigned {
+		s.PriorityAssigned = true
+	}
+
+	atomic.StoreInt64(&s.Priority, priority)
+}
+
+func (s *QueryStats) LoadPriority() (int64, bool) {
+	if s == nil {
+		return 0, false
+	}
+
+	return atomic.LoadInt64(&s.Priority), s.PriorityAssigned
+}
+
+func (s *QueryStats) SetDataSelectMaxTime(dataSelectMaxTime int64) {
+	if s == nil {
+		return
+	}
+
+	atomic.StoreInt64(&s.DataSelectMaxTime, dataSelectMaxTime)
+}
+
+func (s *QueryStats) LoadDataSelectMaxTime() int64 {
+	if s == nil {
+		return 0
+	}
+
+	return atomic.LoadInt64(&s.DataSelectMaxTime)
+}
+
+func (s *QueryStats) SetDataSelectMinTime(dataSelectMinTime int64) {
+	if s == nil {
+		return
+	}
+
+	atomic.StoreInt64(&s.DataSelectMinTime, dataSelectMinTime)
+}
+
+func (s *QueryStats) LoadDataSelectMinTime() int64 {
+	if s == nil {
+		return 0
+	}
+
+	return atomic.LoadInt64(&s.DataSelectMinTime)
+}
+
 // Merge the provided Stats into this one.
 func (s *QueryStats) Merge(other *QueryStats) {
 	if s == nil || other == nil {

pkg/querier/tripperware/priority.go

@@ -1,28 +1,14 @@
 package tripperware
 
 import (
-	"errors"
 	"time"
 
 	"github.com/cortexproject/cortex/pkg/util/validation"
 )
 
-var (
-	errQueryPriorityDisabled = errors.New("query priority disabled")
-	errEmptyQueryString      = errors.New("empty query string")
-)
-
-func GetPriority(query string, minTime, maxTime int64, now time.Time, queryPriority validation.QueryPriority) (int64, error) {
-	if !queryPriority.Enabled {
-		return 0, errQueryPriorityDisabled
-	}
-
-	if query == "" {
-		return 0, errEmptyQueryString
-	}
-
-	if len(queryPriority.Priorities) == 0 {
-		return queryPriority.DefaultPriority, nil
+func GetPriority(query string, minTime, maxTime int64, now time.Time, queryPriority validation.QueryPriority) int64 {
+	if !queryPriority.Enabled || query == "" || len(queryPriority.Priorities) == 0 {
+		return queryPriority.DefaultPriority
 	}
 
 	for _, priority := range queryPriority.Priorities {
@@ -34,12 +20,12 @@ func GetPriority(query string, minTime, maxTime int64, now time.Time, queryPrior
 			}
 
 			if isWithinTimeAttributes(attribute.TimeWindow, now, minTime, maxTime) {
-				return priority.Priority, nil
+				return priority.Priority
 			}
 		}
 	}
 
-	return queryPriority.DefaultPriority, nil
+	return queryPriority.DefaultPriority
 }
 
 func isWithinTimeAttributes(timeWindow validation.TimeWindow, now time.Time, startTime, endTime int64) bool {

pkg/querier/tripperware/priority_test.go

@@ -29,31 +29,23 @@ func Test_GetPriorityShouldReturnDefaultPriorityIfNotEnabledOrInvalidQueryString
 
 	type testCase struct {
 		query                string
-		err                  error
 		queryPriorityEnabled bool
 	}
 
 	tests := map[string]testCase{
 		"should miss if query priority not enabled": {
 			query: "up",
-			err:   errQueryPriorityDisabled,
 		},
 		"should miss if query string empty": {
 			query:                "",
-			err:                  errEmptyQueryString,
 			queryPriorityEnabled: true,
 		},
 	}
 
 	for testName, testData := range tests {
 		t.Run(testName, func(t *testing.T) {
 			limits.queryPriority.Enabled = testData.queryPriorityEnabled
-			priority, err := GetPriority(testData.query, 0, 0, now, limits.queryPriority)
-			if err != nil {
-				assert.Equal(t, testData.err, err)
-			} else {
-				assert.NoError(t, err)
-			}
+			priority := GetPriority(testData.query, 0, 0, now, limits.queryPriority)
 			assert.Equal(t, int64(0), priority)
 		})
 	}
@@ -111,8 +103,7 @@ func Test_GetPriorityShouldConsiderRegex(t *testing.T) {
 		t.Run(testName, func(t *testing.T) {
 			limits.queryPriority.Priorities[0].QueryAttributes[0].Regex = testData.regex
 			limits.queryPriority.Priorities[0].QueryAttributes[0].CompiledRegex = regexp.MustCompile(testData.regex)
-			priority, err := GetPriority(testData.query, 0, 0, now, limits.queryPriority)
-			assert.NoError(t, err)
+			priority := GetPriority(testData.query, 0, 0, now, limits.queryPriority)
 			assert.Equal(t, int64(testData.expectedPriority), priority)
 		})
 	}
@@ -180,8 +171,7 @@ func Test_GetPriorityShouldConsiderStartAndEndTime(t *testing.T) {
 
 	for testName, testData := range tests {
 		t.Run(testName, func(t *testing.T) {
-			priority, err := GetPriority("sum(up)", testData.start.UnixMilli(), testData.end.UnixMilli(), now, limits.queryPriority)
-			assert.NoError(t, err)
+			priority := GetPriority("sum(up)", testData.start.UnixMilli(), testData.end.UnixMilli(), now, limits.queryPriority)
 			assert.Equal(t, int64(testData.expectedPriority), priority)
 		})
 	}
@@ -225,8 +215,7 @@ func Test_GetPriorityShouldNotConsiderStartAndEndTimeIfEmpty(t *testing.T) {
 
 	for testName, testData := range tests {
 		t.Run(testName, func(t *testing.T) {
-			priority, err := GetPriority("sum(up)", testData.start.Unix(), testData.end.Unix(), now, limits.queryPriority)
-			assert.NoError(t, err)
+			priority := GetPriority("sum(up)", testData.start.Unix(), testData.end.Unix(), now, limits.queryPriority)
 			assert.Equal(t, int64(1), priority)
 		})
 	}

pkg/querier/tripperware/roundtrip.go

@@ -32,6 +32,7 @@ import (
 	"github.com/weaveworks/common/httpgrpc"
 	"github.com/weaveworks/common/user"
 
+	"github.com/cortexproject/cortex/pkg/querier/stats"
 	"github.com/cortexproject/cortex/pkg/tenant"
 	"github.com/cortexproject/cortex/pkg/util"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
@@ -165,17 +166,14 @@ func NewQueryTripperware(
 						return next.RoundTrip(r)
 					}
 
+					reqStats := stats.FromContext(r.Context())
 					minTime, maxTime := util.FindMinMaxTime(r, expr, lookbackDelta, now)
-					*r = *r.WithContext(context.WithValue(r.Context(), DataFetchedMinTimeCtxKey, minTime))
-					*r = *r.WithContext(context.WithValue(r.Context(), DataFetchedMaxTimeCtxKey, maxTime))
+					reqStats.SetDataSelectMaxTime(maxTime)
+					reqStats.SetDataSelectMinTime(minTime)
 
 					if limits != nil && limits.QueryPriority(userStr).Enabled {
-						priority, err := GetPriority(query, minTime, maxTime, now, limits.QueryPriority(userStr))
-						if err != nil {
-							level.Debug(log).Log("msg", "failed to get query priority for user", "user", userStr, "err", err.Error())
-						} else {
-							*r = *r.WithContext(context.WithValue(r.Context(), QueryPriorityCtxKey, priority))
-						}
+						priority := GetPriority(query, minTime, maxTime, now, limits.QueryPriority(userStr))
+						reqStats.SetPriority(priority)
 					}
 				}
 

pkg/querier/tripperware/util.go

@@ -10,14 +10,6 @@ import (
 	"github.com/cortexproject/cortex/pkg/util/validation"
 )
 
-type contextKey string
-
-const (
-	DataFetchedMinTimeCtxKey = contextKey("dataFetchedMinTime")
-	DataFetchedMaxTimeCtxKey = contextKey("dataFetchedMaxTime")
-	QueryPriorityCtxKey      = contextKey("queryPriority")
-)
-
 // RequestResponse contains a request response and the respective request that was used.
 type RequestResponse struct {
 	Request  Request

pkg/scheduler/scheduler.go

@@ -21,7 +21,7 @@ import (
 	"google.golang.org/grpc"
 
 	"github.com/cortexproject/cortex/pkg/frontend/v2/frontendv2pb"
-	"github.com/cortexproject/cortex/pkg/querier/tripperware"
+	"github.com/cortexproject/cortex/pkg/querier/stats"
 	"github.com/cortexproject/cortex/pkg/scheduler/queue"
 	"github.com/cortexproject/cortex/pkg/scheduler/schedulerpb"
 	"github.com/cortexproject/cortex/pkg/tenant"
@@ -167,7 +167,7 @@ type schedulerRequest struct {
 }
 
 func (s schedulerRequest) Priority() int64 {
-	priority, ok := s.ctx.Value(tripperware.QueryPriorityCtxKey).(int64)
+	priority, ok := stats.FromContext(s.ctx).LoadPriority()
 	if !ok {
 		return 0
 	}

pkg/util/validation/limits.go

@@ -64,13 +64,13 @@ type PriorityDef struct {
 
 type QueryAttribute struct {
 	Regex         string     `yaml:"regex" json:"regex" doc:"nocli|description=Regex that the query string should match. If not set, it won't be checked."`
-	TimeWindow    TimeWindow `yaml:"time_window" json:"time_window" doc:"nocli|description=Overall data fetch time window (including range selectors, modifiers and lookback delta) that the query should be within. If not set, it won't be checked."`
+	TimeWindow    TimeWindow `yaml:"time_window" json:"time_window" doc:"nocli|description=Overall data select time window (including range selectors, modifiers and lookback delta) that the query should be within. If not set, it won't be checked."`
 	CompiledRegex *regexp.Regexp
 }
 
 type TimeWindow struct {
-	Start model.Duration `yaml:"start" json:"start" doc:"nocli|description=Start of the data fetch time window (including range selectors, modifiers and lookback delta) that the query should be within. If set to 0, it won't be checked.|default=0"`
-	End   model.Duration `yaml:"end" json:"end" doc:"nocli|description=End of the data fetch time window (including range selectors, modifiers and lookback delta) that the query should be within. If set to 0, it won't be checked.|default=0"`
+	Start model.Duration `yaml:"start" json:"start" doc:"nocli|description=Start of the data select time window (including range selectors, modifiers and lookback delta) that the query should be within. If set to 0, it won't be checked.|default=0"`
+	End   model.Duration `yaml:"end" json:"end" doc:"nocli|description=End of the data select time window (including range selectors, modifiers and lookback delta) that the query should be within. If set to 0, it won't be checked.|default=0"`
 }
 
 // Limits describe all the limits for users; can be used to describe global default