Add reserved querier logic

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

Author: justinjung04

CHANGELOG.md

@@ -42,7 +42,7 @@
 * [FEATURE] Store Gateway: Add `-blocks-storage.bucket-store.max-inflight-requests` for store gateways to reject further requests upon reaching the limit. #5553
 * [FEATURE] Store Gateway: Add `cortex_bucket_store_block_load_duration_seconds` histogram to track time to load blocks. #5580
 * [FEATURE] AlertManager: Add `cortex_alertmanager_dispatcher_aggregation_groups` and `cortex_alertmanager_dispatcher_alert_processing_duration_seconds` metrics for dispatcher. #5592
-* [FEATURE] Query Frontend/Scheduler: Add high-priority query support. #5605
+* [FEATURE] Query Frontend/Scheduler: Add query priority support. #5605
 * [ENHANCEMENT] Distributor/Ingester: Add span on push path #5319
 * [ENHANCEMENT] Support object storage backends for runtime configuration file. #5292
 * [ENHANCEMENT] Query Frontend: Reject subquery with too small step size. #5323

pkg/frontend/v1/frontend.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"reflect"
 	"time"
 
 	"github.com/go-kit/log"
@@ -79,6 +80,9 @@ type Frontend struct {
 	requestQueue *queue.RequestQueue
 	activeUsers  *util.ActiveUsersCleanupService
 
+	queryPriority         validation.QueryPriority
+	compiledQueryPriority validation.QueryPriority
+
 	// Subservices manager.
 	subservices        *services.Manager
 	subservicesWatcher *services.FailureWatcher
@@ -207,8 +211,15 @@ func (f *Frontend) RoundTripGRPC(ctx context.Context, req *httpgrpc.HTTPRequest,
 		}
 
 		queryPriority := f.limits.QueryPriority(userID)
+
 		if reqParams != nil && queryPriority.Enabled {
-			request.priority = util_query.GetPriority(reqParams, ts, queryPriority)
+			queryPriorityChanged := !reflect.DeepEqual(f.queryPriority, queryPriority)
+			if queryPriorityChanged {
+				f.queryPriority = queryPriority
+				f.compiledQueryPriority = queryPriority
+			}
+
+			request.priority = util_query.GetPriority(reqParams, ts, &f.compiledQueryPriority, queryPriorityChanged)
 		}
 
 		if err := f.queueRequest(ctx, &request); err != nil {

pkg/frontend/v2/frontend.go

@@ -7,6 +7,7 @@ import (
 	"math/rand"
 	"net/http"
 	"net/url"
+	"reflect"
 	"sync"
 	"time"
 
@@ -30,6 +31,7 @@ import (
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
 	util_query "github.com/cortexproject/cortex/pkg/util/query"
 	"github.com/cortexproject/cortex/pkg/util/services"
+	"github.com/cortexproject/cortex/pkg/util/validation"
 )
 
 // Config for a Frontend.
@@ -74,6 +76,9 @@ type Frontend struct {
 
 	lastQueryID atomic.Uint64
 
+	queryPriority         validation.QueryPriority
+	compiledQueryPriority validation.QueryPriority
+
 	// frontend workers will read from this channel, and send request to scheduler.
 	requestsCh chan *frontendRequest
 
@@ -211,8 +216,15 @@ func (f *Frontend) RoundTripGRPC(ctx context.Context, req *httpgrpc.HTTPRequest,
 		}
 
 		queryPriority := f.limits.QueryPriority(userID)
+
 		if reqParams != nil && queryPriority.Enabled {
-			freq.priority = util_query.GetPriority(reqParams, ts, queryPriority)
+			queryPriorityChanged := !reflect.DeepEqual(f.queryPriority, queryPriority)
+			if queryPriorityChanged {
+				f.queryPriority = queryPriority
+				f.compiledQueryPriority = queryPriority
+			}
+
+			freq.priority = util_query.GetPriority(reqParams, ts, &f.compiledQueryPriority, queryPriorityChanged)
 		}
 
 		f.requests.put(freq)

pkg/scheduler/queue/queue.go

@@ -98,7 +98,8 @@ func (q *RequestQueue) EnqueueRequest(userID string, req Request, maxQueriers fl
 	}
 
 	shardSize := util.DynamicShardSize(maxQueriers, len(q.queues.queriers))
-	queue := q.queues.getOrAddQueue(userID, shardSize)
+	priorityList, priorityEnabled := q.getPriorityList(userID)
+	queue := q.queues.getOrAddQueue(userID, shardSize, priorityList, priorityEnabled)
 	maxOutstandingRequests := q.queues.limits.MaxOutstandingPerTenant(userID)
 
 	if queue == nil {
@@ -123,6 +124,24 @@ func (q *RequestQueue) EnqueueRequest(userID string, req Request, maxQueriers fl
 	return nil
 }
 
+func (q *RequestQueue) getPriorityList(userID string) ([]int64, bool) {
+	var priorityList []int64
+
+	queryPriority := q.queues.limits.QueryPriority(userID)
+
+	if queryPriority.Enabled {
+		for _, priority := range queryPriority.Priorities {
+			reservedQuerierShardSize := util.DynamicShardSize(priority.ReservedQueriers, len(q.queues.queriers))
+
+			for i := 0; i < reservedQuerierShardSize; i++ {
+				priorityList = append(priorityList, priority.Priority)
+			}
+		}
+	}
+
+	return priorityList, queryPriority.Enabled
+}
+
 // GetNextRequestForQuerier find next user queue and takes the next request off of it. Will block if there are no requests.
 // By passing user index from previous call of this method, querier guarantees that it iterates over all users fairly.
 // If querier finds that request from the user is already expired, it can get a request for the same user by using UserIndex.ReuseLastUser.
@@ -156,7 +175,8 @@ FindQueue:
 
 		// Pick next request from the queue.
 		for {
-			request := queue.dequeueRequest(q.queues.getMinPriority(userID, querierID))
+			minPriority, checkMinPriority := q.queues.getMinPriority(userID, querierID)
+			request := queue.dequeueRequest(minPriority, checkMinPriority)
 			if request == nil {
 				// the queue does not contain request with the min priority, break to wait for more requests
 				break

pkg/scheduler/queue/queue_test.go

@@ -3,7 +3,6 @@ package queue
 import (
 	"context"
 	"fmt"
-	"github.com/cortexproject/cortex/pkg/util/validation"
 	"strconv"
 	"sync"
 	"testing"
@@ -14,6 +13,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/cortexproject/cortex/pkg/util/services"
+	"github.com/cortexproject/cortex/pkg/util/validation"
 )
 
 func BenchmarkGetNextRequest(b *testing.B) {

pkg/scheduler/queue/user_queues.go

@@ -2,6 +2,7 @@ package queue
 
 import (
 	"math/rand"
+	"reflect"
 	"sort"
 	"time"
 
@@ -63,8 +64,11 @@ type userQueue struct {
 
 	// If not nil, only these queriers can handle user requests. If nil, all queriers can.
 	// We set this to nil if number of available queriers <= maxQueriers.
-	queriers    map[string]struct{}
-	maxQueriers int
+	queriers         map[string]struct{}
+	reservedQueriers map[string]int64
+	priorityList     []int64
+	priorityEnabled  bool
+	maxQueriers      int
 
 	// Seed for shuffle sharding of queriers. This seed is based on userID only and is therefore consistent
 	// between different frontends.
@@ -110,7 +114,7 @@ func (q *queues) deleteQueue(userID string) {
 // MaxQueriers is used to compute which queriers should handle requests for this user.
 // If maxQueriers is <= 0, all queriers can handle this user's requests.
 // If maxQueriers has changed since the last call, queriers for this are recomputed.
-func (q *queues) getOrAddQueue(userID string, maxQueriers int) userRequestQueue {
+func (q *queues) getOrAddQueue(userID string, maxQueriers int, priorityList []int64, priorityEnabled bool) userRequestQueue {
 	// Empty user is not allowed, as that would break our users list ("" is used for free spot).
 	if userID == "" {
 		return nil
@@ -121,48 +125,12 @@ func (q *queues) getOrAddQueue(userID string, maxQueriers int) userRequestQueue
 	}
 
 	uq := q.userQueues[userID]
-	queryPriority := q.limits.QueryPriority(userID)
 
 	if uq == nil {
-		queueSize := q.limits.MaxOutstandingPerTenant(userID)
-		// 0 is the default value of the flag. If the old flag is set
-		// then we use its value for compatibility reason.
-		if q.maxUserQueueSize != 0 {
-			queueSize = q.maxUserQueueSize
-		}
-		uq = &userQueue{
-			seed:  util.ShuffleShardSeed(userID, ""),
-			index: -1,
-		}
-
-		if queryPriority.Enabled {
-			uq.queue = NewPriorityRequestQueue(util.NewPriorityQueue(nil))
-		} else {
-			uq.queue = NewFIFORequestQueue(make(chan Request, queueSize))
-		}
-
-		q.userQueues[userID] = uq
-
-		// Add user to the list of users... find first free spot, and put it there.
-		for ix, u := range q.users {
-			if u == "" {
-				uq.index = ix
-				q.users[ix] = userID
-				break
-			}
-		}
-
-		// ... or add to the end.
-		if uq.index < 0 {
-			uq.index = len(q.users)
-			q.users = append(q.users, userID)
-		}
+		uq = q.createUserQueue(userID)
 	}
 
-	if uq.maxQueriers != maxQueriers {
-		uq.maxQueriers = maxQueriers
-		uq.queriers = shuffleQueriersForUser(uq.seed, maxQueriers, q.sortedQueriers, nil)
-	}
+	q.updateUserQueuesAttributes(uq, userID, maxQueriers, priorityList, priorityEnabled)
 
 	return uq.queue
 }
@@ -196,21 +164,98 @@ func (q *queues) getNextQueueForQuerier(lastUserIndex int, querierID string) (us
 			}
 		}
 
-		// TODO: justinjung04, reserved queriers
-		//if priority, isReserved := uq.reservedQueriers[querierID]; isReserved {
-		//	return uq.queues[priority], u, uid
-		//}
-
 		return uq.queue, u, uid
 	}
 	return nil, "", uid
 }
 
-func (q *queues) getMinPriority(userID string, querierID string) int64 {
-	// TODO: justinjung04 reserved querier
-	// check list of queriers and QueryPriority config
-	// from QueryPriority config, establish map of
-	return 0
+func (q *queues) createUserQueue(userID string) *userQueue {
+	uq := &userQueue{
+		seed:  util.ShuffleShardSeed(userID, ""),
+		index: -1,
+	}
+
+	uq.queue = q.createUserRequestQueue(userID)
+	q.userQueues[userID] = uq
+
+	// Add user to the list of users... find first free spot, and put it there.
+	for ix, u := range q.users {
+		if u == "" {
+			uq.index = ix
+			q.users[ix] = userID
+			break
+		}
+	}
+
+	// ... or add to the end.
+	if uq.index < 0 {
+		uq.index = len(q.users)
+		q.users = append(q.users, userID)
+	}
+
+	return uq
+}
+
+func (q *queues) createUserRequestQueue(userID string) userRequestQueue {
+	if q.limits.QueryPriority(userID).Enabled {
+		return NewPriorityRequestQueue(util.NewPriorityQueue(nil))
+	}
+
+	return NewFIFORequestQueue(make(chan Request, q.getQueueSize(userID)))
+}
+
+func (q *queues) getQueueSize(userID string) int {
+	queueSize := q.limits.MaxOutstandingPerTenant(userID)
+
+	// 0 is the default value of the flag. If the old flag is set
+	// then we use its value for compatibility reason.
+	if q.maxUserQueueSize != 0 {
+		queueSize = q.maxUserQueueSize
+	}
+
+	return queueSize
+}
+
+func (q *queues) updateUserQueuesAttributes(uq *userQueue, userID string, maxQueriers int, priorityList []int64, priorityEnabled bool) {
+	if uq.maxQueriers != maxQueriers {
+		uq.maxQueriers = maxQueriers
+		uq.queriers = shuffleQueriersForUser(uq.seed, maxQueriers, q.sortedQueriers, nil)
+	}
+
+	// if query priority is newly enabled/disabled, transfer the requests to the new queue
+	if uq.priorityEnabled != priorityEnabled {
+		tmpQueue := q.createUserRequestQueue(userID)
+		uq.queue.closeQueue()
+
+		for uq.queue.length() > 0 {
+			tmpQueue.enqueueRequest(uq.queue.dequeueRequest(0, false))
+		}
+		uq.queue = tmpQueue
+	}
+
+	if priorityEnabled && !reflect.DeepEqual(uq.priorityList, priorityList) {
+		reservedQueriers := make(map[string]int64)
+
+		i := 0
+		for querierID := range uq.queriers {
+			reservedQueriers[querierID] = priorityList[i]
+			i++
+			if i == len(priorityList) {
+				break
+			}
+		}
+
+		uq.reservedQueriers = reservedQueriers
+		uq.priorityList = priorityList
+		uq.priorityEnabled = priorityEnabled
+	}
+}
+
+func (q *queues) getMinPriority(userID string, querierID string) (int64, bool) {
+	if priority, ok := q.userQueues[userID].reservedQueriers[querierID]; ok {
+		return priority, true
+	}
+	return 0, false
 }
 
 func (q *queues) addQuerierConnection(querierID string) {
@@ -349,7 +394,7 @@ type MockLimits struct {
 	queryPriority      validation.QueryPriority
 }
 
-func (l MockLimits) MaxQueriersPerUser(user string) float64 {
+func (l MockLimits) MaxQueriersPerUser(_ string) float64 {
 	return l.maxQueriersPerUser
 }
 

pkg/scheduler/queue/user_queues_test.go

@@ -156,7 +156,7 @@ func TestQueuesConsistency(t *testing.T) {
 			conns := map[string]int{}
 
 			for i := 0; i < 10000; i++ {
-				queue := uq.getOrAddQueue(generateTenant(r), 3)
+				queue := uq.getOrAddQueue(generateTenant(r), 3, []int64{}, true)
 				switch r.Int() % 6 {
 				case 0:
 					assert.NotNil(t, queue)
@@ -352,6 +352,10 @@ func TestQueues_ForgetDelay_ShouldCorrectlyHandleQuerierReconnectingBeforeForget
 	}
 }
 
+func TestQueuesUpdatedConfig(t *testing.T) {
+	// TODO: justinjung04
+}
+
 func generateTenant(r *rand.Rand) string {
 	return fmt.Sprint("tenant-", r.Int()%5)
 }
@@ -361,10 +365,10 @@ func generateQuerier(r *rand.Rand) string {
 }
 
 func getOrAdd(t *testing.T, uq *queues, tenant string, maxQueriers int) userRequestQueue {
-	q := uq.getOrAddQueue(tenant, maxQueriers)
+	q := uq.getOrAddQueue(tenant, maxQueriers, []int64{}, true)
 	assert.NotNil(t, q)
 	assert.NoError(t, isConsistent(uq))
-	assert.Equal(t, q, uq.getOrAddQueue(tenant, maxQueriers))
+	assert.Equal(t, q, uq.getOrAddQueue(tenant, maxQueriers, []int64{}, true))
 	return q
 }