Query priority

CHANGELOG.md

@@ -5,6 +5,7 @@
 * [CHANGE] Store Gateway: Add a new fastcache based inmemory index cache. #5619
 * [CHANGE] Index Cache: Multi level cache backfilling operation becomes async. Added `-blocks-storage.bucket-store.index-cache.multilevel.max-async-concurrency` and `-blocks-storage.bucket-store.index-cache.multilevel.max-async-buffer-size` configs and metric `cortex_store_multilevel_index_cache_backfill_dropped_items_total` for number of dropped items. #5661
 * [FEATURE] Ingester: Add per-tenant new metric `cortex_ingester_tsdb_data_replay_duration_seconds`. #5477
+* [FEATURE] Query Frontend/Scheduler: Add query priority support. #5605
 * [ENHANCEMENT] Store Gateway: Added `-store-gateway.enabled-tenants` and `-store-gateway.disabled-tenants` to explicitly enable or disable store-gateway for specific tenants. #5638
 
 ## 1.16.0 2023-11-20

docs/configuration/config-file-reference.md

@@ -3041,6 +3041,18 @@ The `limits_config` configures default and per-tenant limits imposed by Cortex s
 # CLI flag: -frontend.max-outstanding-requests-per-tenant
 [max_outstanding_requests_per_tenant: <int> | default = 100]
 
+# Configuration for query priority.
+query_priority:
+  # Whether queries are assigned with priorities.
+  [enabled: <boolean> | default = false]
+
+  # Priority assigned to all queries by default. Must be a unique value. Use
+  # this as a baseline to make certain queries higher/lower priority.
+  [default_priority: <int> | default = 0]
+
+  # List of priority definitions.
+  [priorities: <list of PriorityDef> | default = []]
+
 # Duration to delay the evaluation of rules to ensure the underlying metrics
 # have been pushed to Cortex.
 # CLI flag: -ruler.evaluation-delay-duration
@@ -5032,6 +5044,33 @@ otel:
     [tls_insecure_skip_verify: <boolean> | default = false]
 ```
 
+### `PriorityDef`
+
+```yaml
+# Priority level. Must be a unique value.
+[priority: <int> | default = 0]
+
+# Number of reserved queriers to handle priorities higher or equal to the
+# priority level. Value between 0 and 1 will be used as a percentage.
+[reserved_queriers: <float> | default = 0]
+
+# List of query attributes to assign the priority.
+[query_attributes: <list of QueryAttribute> | default = []]
+```
+
+### `QueryAttribute`
+
+```yaml
+# Query string regex. If set to empty string, it will not match anything.
+[regex: <string> | default = ""]
+
+# Query start time. If set to 0, the start time won't be checked.
+[start_time: <int> | default = 0]
+
+# Query end time. If set to 0, the end time won't be checked.
+[end_time: <int> | default = 0]
+```
+
 ### `DisabledRuleGroup`
 
 ```yaml

pkg/frontend/config.go

@@ -59,7 +59,7 @@ func InitFrontend(cfg CombinedFrontendConfig, limits v1.Limits, grpcListenPort i
 			cfg.FrontendV2.Port = grpcListenPort
 		}
 
-		fr, err := v2.NewFrontend(cfg.FrontendV2, log, reg, retry)
+		fr, err := v2.NewFrontend(cfg.FrontendV2, limits, log, reg, retry)
 		return transport.AdaptGrpcRoundTripperToHTTPRoundTripper(fr), nil, fr, err
 
 	default:

pkg/frontend/transport/handler.go

@@ -200,6 +200,7 @@ func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		r.Body = io.NopCloser(&buf)
 	}
 
+	r.Header.Get("test")
 	startTime := time.Now()
 	resp, err := f.roundTripper.RoundTrip(r)
 	queryResponseTime := time.Since(startTime)
@@ -343,6 +344,9 @@ func (f *Handler) reportQueryStats(r *http.Request, userID string, queryString u
 	if ua := r.Header.Get("User-Agent"); len(ua) > 0 {
 		logMessage = append(logMessage, "user_agent", ua)
 	}
+	if queryPriority := r.Header.Get(util.QueryPriorityHeaderKey); len(queryPriority) > 0 {
+		logMessage = append(logMessage, "priority", queryPriority)
+	}
 
 	if error != nil {
 		s, ok := status.FromError(error)

pkg/frontend/v1/frontend.go

@@ -5,6 +5,7 @@ import (
 	"flag"
 	"fmt"
 	"net/http"
+	"strconv"
 	"time"
 
 	"github.com/go-kit/log"
@@ -51,14 +52,19 @@ type Limits interface {
 
 // MockLimits implements the Limits interface. Used in tests only.
 type MockLimits struct {
-	Queriers float64
+	Queriers      float64
+	queryPriority validation.QueryPriority
 	queue.MockLimits
 }
 
 func (l MockLimits) MaxQueriersPerUser(_ string) float64 {
 	return l.Queriers
 }
 
+func (l MockLimits) QueryPriority(_ string) validation.QueryPriority {
+	return l.queryPriority
+}
+
 // Frontend queues HTTP requests, dispatches them to backends, and handles retries
 // for requests which failed.
 type Frontend struct {
@@ -93,6 +99,15 @@ type request struct {
 	response chan *httpgrpc.HTTPResponse
 }
 
+func (r request) Priority() int64 {
+	priority, err := strconv.ParseInt(httpgrpcutil.GetHeader(*r.request, util.QueryPriorityHeaderKey), 10, 64)
+	if err != nil {
+		return 0
+	}
+
+	return priority
+}
+
 // New creates a new frontend. Frontend implements service, and must be started and stopped.
 func New(cfg Config, limits Limits, log log.Logger, registerer prometheus.Registerer, retry *transport.Retry) (*Frontend, error) {
 	f := &Frontend{
@@ -103,11 +118,11 @@ func New(cfg Config, limits Limits, log log.Logger, registerer prometheus.Regist
 		queueLength: promauto.With(registerer).NewGaugeVec(prometheus.GaugeOpts{
 			Name: "cortex_query_frontend_queue_length",
 			Help: "Number of queries in the queue.",
-		}, []string{"user"}),
+		}, []string{"user", "priority", "type"}),
 		discardedRequests: promauto.With(registerer).NewCounterVec(prometheus.CounterOpts{
 			Name: "cortex_query_frontend_discarded_requests_total",
 			Help: "Total number of query requests discarded.",
-		}, []string{"user"}),
+		}, []string{"user", "priority"}),
 		queueDuration: promauto.With(registerer).NewHistogram(prometheus.HistogramOpts{
 			Name:    "cortex_query_frontend_queue_duration_seconds",
 			Help:    "Time spend by requests queued.",
@@ -160,8 +175,12 @@ func (f *Frontend) stopping(_ error) error {
 }
 
 func (f *Frontend) cleanupInactiveUserMetrics(user string) {
-	f.queueLength.DeleteLabelValues(user)
-	f.discardedRequests.DeleteLabelValues(user)
+	f.queueLength.DeletePartialMatch(prometheus.Labels{
+		"user": user,
+	})
+	f.discardedRequests.DeletePartialMatch(prometheus.Labels{
+		"user": user,
+	})
 }
 
 // RoundTripGRPC round trips a proto (instead of a HTTP request).

pkg/frontend/v1/frontend_test.go

@@ -211,7 +211,7 @@ func TestFrontendMetricsCleanup(t *testing.T) {
 			require.NoError(t, testutil.GatherAndCompare(reg, strings.NewReader(`
 				# HELP cortex_query_frontend_queue_length Number of queries in the queue.
 				# TYPE cortex_query_frontend_queue_length gauge
-				cortex_query_frontend_queue_length{user="1"} 0
+				cortex_query_frontend_queue_length{priority="0",type="fifo",user="1"} 0
 			`), "cortex_query_frontend_queue_length"))
 
 			fr.cleanupInactiveUserMetrics("1")

pkg/frontend/v2/frontend.go

@@ -21,6 +21,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/frontend/transport"
 	"github.com/cortexproject/cortex/pkg/frontend/v2/frontendv2pb"
 	"github.com/cortexproject/cortex/pkg/querier/stats"
+	"github.com/cortexproject/cortex/pkg/scheduler"
 	"github.com/cortexproject/cortex/pkg/tenant"
 	"github.com/cortexproject/cortex/pkg/util/flagext"
 	"github.com/cortexproject/cortex/pkg/util/grpcclient"
@@ -64,10 +65,10 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 type Frontend struct {
 	services.Service
 
-	cfg Config
-	log log.Logger
-
-	retry *transport.Retry
+	cfg    Config
+	log    log.Logger
+	limits scheduler.Limits
+	retry  *transport.Retry
 
 	lastQueryID atomic.Uint64
 
@@ -112,7 +113,7 @@ type enqueueResult struct {
 }
 
 // NewFrontend creates a new frontend.
-func NewFrontend(cfg Config, log log.Logger, reg prometheus.Registerer, retry *transport.Retry) (*Frontend, error) {
+func NewFrontend(cfg Config, limits scheduler.Limits, log log.Logger, reg prometheus.Registerer, retry *transport.Retry) (*Frontend, error) {
 	requestsCh := make(chan *frontendRequest)
 
 	schedulerWorkers, err := newFrontendSchedulerWorkers(cfg, fmt.Sprintf("%s:%d", cfg.Addr, cfg.Port), requestsCh, log)
@@ -122,6 +123,7 @@ func NewFrontend(cfg Config, log log.Logger, reg prometheus.Registerer, retry *t
 
 	f := &Frontend{
 		cfg:              cfg,
+		limits:           limits,
 		log:              log,
 		requestsCh:       requestsCh,
 		schedulerWorkers: schedulerWorkers,

pkg/frontend/v2/frontend_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/frontend/transport"
 	"github.com/cortexproject/cortex/pkg/frontend/v2/frontendv2pb"
 	"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/util/flagext"
 	"github.com/cortexproject/cortex/pkg/util/services"
@@ -48,7 +49,7 @@ func setupFrontend(t *testing.T, schedulerReplyFunc func(f *Frontend, msg *sched
 
 	//logger := log.NewLogfmtLogger(os.Stdout)
 	logger := log.NewNopLogger()
-	f, err := NewFrontend(cfg, logger, nil, transport.NewRetry(maxRetries, nil))
+	f, err := NewFrontend(cfg, queue.MockLimits{}, logger, nil, transport.NewRetry(maxRetries, nil))
 	require.NoError(t, err)
 
 	frontendv2pb.RegisterFrontendForQuerierServer(server, f)

pkg/querier/tripperware/instantquery/instant_query.go

@@ -4,18 +4,17 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"github.com/cortexproject/cortex/pkg/util"
 	"io"
 	"net/http"
 	"net/url"
 	"sort"
-	"strconv"
 	"strings"
 	"time"
 
 	jsoniter "github.com/json-iterator/go"
 	"github.com/opentracing/opentracing-go"
 	otlog "github.com/opentracing/opentracing-go/log"
-	"github.com/pkg/errors"
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/model/timestamp"
@@ -26,7 +25,6 @@ import (
 	"github.com/cortexproject/cortex/pkg/cortexpb"
 	"github.com/cortexproject/cortex/pkg/querier/tripperware"
 	"github.com/cortexproject/cortex/pkg/querier/tripperware/queryrange"
-	"github.com/cortexproject/cortex/pkg/util"
 	"github.com/cortexproject/cortex/pkg/util/spanlogger"
 )
 
@@ -132,7 +130,7 @@ func (resp *PrometheusInstantQueryResponse) HTTPHeaders() map[string][]string {
 func (c instantQueryCodec) DecodeRequest(_ context.Context, r *http.Request, forwardHeaders []string) (tripperware.Request, error) {
 	result := PrometheusRequest{Headers: map[string][]string{}}
 	var err error
-	result.Time, err = parseTimeParam(r, "time", c.now().Unix())
+	result.Time, err = util.ParseTimeParam(r, "time", c.now().Unix())
 	if err != nil {
 		return nil, decorateWithParamName(err, "time")
 	}
@@ -630,15 +628,3 @@ func (s *PrometheusInstantQueryData) MarshalJSON() ([]byte, error) {
 		return s.Result.GetRawBytes(), nil
 	}
 }
-
-func parseTimeParam(r *http.Request, paramName string, defaultValue int64) (int64, error) {
-	val := r.FormValue(paramName)
-	if val == "" {
-		val = strconv.FormatInt(defaultValue, 10)
-	}
-	result, err := util.ParseTime(val)
-	if err != nil {
-		return 0, errors.Wrapf(err, "Invalid time value for '%s'", paramName)
-	}
-	return result, nil
-}

pkg/querier/tripperware/limits.go

@@ -1,6 +1,10 @@
 package tripperware
 
-import "time"
+import (
+	"time"
+
+	"github.com/cortexproject/cortex/pkg/util/validation"
+)
 
 // Limits allows us to specify per-tenant runtime limits on the behavior of
 // the query handling code.
@@ -21,4 +25,7 @@ type Limits interface {
 
 	// QueryVerticalShardSize returns the maximum number of queriers that can handle requests for this user.
 	QueryVerticalShardSize(userID string) int
+
+	// QueryPriority returns the query priority config for the tenant, including different priorities and their attributes.
+	QueryPriority(userID string) validation.QueryPriority
 }

pkg/querier/tripperware/priority.go

@@ -0,0 +1,94 @@
+package tripperware
+
+import (
+	"net/http"
+	"strings"
+	"time"
+
+	"github.com/prometheus/prometheus/promql/parser"
+
+	"github.com/cortexproject/cortex/pkg/util"
+	"github.com/cortexproject/cortex/pkg/util/validation"
+)
+
+func GetPriority(r *http.Request, userID string, limits Limits, now time.Time) (int64, error) {
+	isQuery := strings.HasSuffix(r.URL.Path, "/query")
+	isQueryRange := strings.HasSuffix(r.URL.Path, "/query_range")
+	queryPriority := limits.QueryPriority(userID)
+	query := r.FormValue("query")
+
+	if (!isQuery && !isQueryRange) || !queryPriority.Enabled || query == "" {
+		return 0, nil
+	}
+
+	expr, err := parser.ParseExpr(query)
+	if err != nil {
+		return 0, err
+	}
+
+	var startTime, endTime int64
+	if isQuery {
+		if t, err := util.ParseTimeParam(r, "time", now.Unix()); err == nil {
+			startTime = t
+			endTime = t
+		}
+	} else if isQueryRange {
+		if st, err := util.ParseTime(r.FormValue("start")); err == nil {
+			if et, err := util.ParseTime(r.FormValue("end")); err == nil {
+				startTime = st
+				endTime = et
+			}
+		}
+	}
+
+	es := &parser.EvalStmt{
+		Expr:          expr,
+		Start:         util.TimeFromMillis(startTime),
+		End:           util.TimeFromMillis(endTime),
+		LookbackDelta: limits.MaxQueryLookback(userID), // this is available from querier flag.
+	}
+
+	minTime, maxTime := FindMinMaxTime(es)
+
+	for _, priority := range queryPriority.Priorities {
+		for _, attribute := range priority.QueryAttributes {
+			if attribute.Regex == "" || (attribute.CompiledRegex != nil && !attribute.CompiledRegex.MatchString(query)) {
+				continue
+			}
+
+			if isWithinTimeAttributes(attribute, now, minTime, maxTime) {
+				return priority.Priority, nil
+			}
+		}
+	}
+
+	return queryPriority.DefaultPriority, nil
+}
+
+func isWithinTimeAttributes(attribute validation.QueryAttribute, now time.Time, startTime, endTime int64) bool {
+	if attribute.StartTime == 0 && attribute.EndTime == 0 {
+		return true
+	}
+
+	if attribute.StartTime != 0 {
+		startTimeThreshold := now.Add(-1 * time.Duration(attribute.StartTime).Abs()).Truncate(time.Second).Unix()
+		if startTime < startTimeThreshold {
+			return false
+		}
+	}
+
+	if attribute.EndTime != 0 {
+		endTimeThreshold := now.Add(-1 * time.Duration(attribute.EndTime).Abs()).Add(1 * time.Second).Truncate(time.Second).Unix()
+		if endTime > endTimeThreshold {
+			return false
+		}
+	}
+
+	return true
+}
+
+func FindMinMaxTime(s *parser.EvalStmt) (int64, int64) {
+	// Placeholder until Prometheus is updated to >=0.48.0
+	// which includes https://github.com/prometheus/prometheus/commit/9e3df532d8294d4fe3284bde7bc96db336a33552
+	return s.Start.Unix(), s.End.Unix()
+}