Create user request queue test

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

Author: justinjung04

pkg/scheduler/queue/queue.go

@@ -156,7 +156,12 @@ FindQueue:
 
 		// Pick next request from the queue.
 		for {
-			request := queue.dequeueRequest()
+			request := queue.dequeueRequest(q.queues.getMinPriority(userID, querierID))
+			if request == nil {
+				// the queue does not contain request with the min priority, break to wait for more requests
+				break
+			}
+
 			if queue.length() == 0 {
 				q.queues.deleteQueue(userID)
 			}

pkg/scheduler/queue/user_queues.go

@@ -59,13 +59,12 @@ type queues struct {
 }
 
 type userQueue struct {
-	queue requestQueue
+	queue userRequestQueue
 
 	// 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{}
-	reservedQueriers map[string]int64
-	maxQueriers      int
+	queriers    map[string]struct{}
+	maxQueriers int
 
 	// Seed for shuffle sharding of queriers. This seed is based on userID only and is therefore consistent
 	// between different frontends.
@@ -111,7 +110,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) requestQueue {
+func (q *queues) getOrAddQueue(userID string, maxQueriers int) userRequestQueue {
 	// Empty user is not allowed, as that would break our users list ("" is used for free spot).
 	if userID == "" {
 		return nil
@@ -139,7 +138,7 @@ func (q *queues) getOrAddQueue(userID string, maxQueriers int) requestQueue {
 		if queryPriority.Enabled {
 			uq.queue = NewPriorityRequestQueue(util.NewPriorityQueue(nil))
 		} else {
-			uq.queue = NewFIFOQueue(make(chan Request, queueSize))
+			uq.queue = NewFIFORequestQueue(make(chan Request, queueSize))
 		}
 
 		q.userQueues[userID] = uq
@@ -163,7 +162,6 @@ func (q *queues) getOrAddQueue(userID string, maxQueriers int) requestQueue {
 	if uq.maxQueriers != maxQueriers {
 		uq.maxQueriers = maxQueriers
 		uq.queriers = shuffleQueriersForUser(uq.seed, maxQueriers, q.sortedQueriers, nil)
-		uq.reservedQueriers = getReservedQueriers(uq.queriers, q.limits.QueryPriority(userID).Priorities)
 	}
 
 	return uq.queue
@@ -172,7 +170,7 @@ func (q *queues) getOrAddQueue(userID string, maxQueriers int) requestQueue {
 // Finds next queue for the querier. To support fair scheduling between users, client is expected
 // to pass last user index returned by this function as argument. Is there was no previous
 // last user index, use -1.
-func (q *queues) getNextQueueForQuerier(lastUserIndex int, querierID string) (requestQueue, string, int) {
+func (q *queues) getNextQueueForQuerier(lastUserIndex int, querierID string) (userRequestQueue, string, int) {
 	uid := lastUserIndex
 
 	for iters := 0; iters < len(q.users); iters++ {
@@ -198,7 +196,7 @@ func (q *queues) getNextQueueForQuerier(lastUserIndex int, querierID string) (re
 			}
 		}
 
-		//TODO: justinjung04, reserved queriers
+		// TODO: justinjung04, reserved queriers
 		//if priority, isReserved := uq.reservedQueriers[querierID]; isReserved {
 		//	return uq.queues[priority], u, uid
 		//}
@@ -208,6 +206,13 @@ func (q *queues) getNextQueueForQuerier(lastUserIndex int, querierID string) (re
 	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) addQuerierConnection(querierID string) {
 	info := q.queriers[querierID]
 	if info != nil {
@@ -320,16 +325,6 @@ func (q *queues) recomputeUserQueriers() {
 // In that case *all* queriers should be used.
 // Scratchpad is used for shuffling, to avoid new allocations. If nil, new slice is allocated.
 func shuffleQueriersForUser(userSeed int64, queriersToSelect int, allSortedQueriers []string, scratchpad []string) map[string]struct{} {
-	//numOfReservedQueriers := getNumOfReservedQueriers(queriersToSelect, len(allSortedQueriers), numOfReservedQueriersFloat)
-	//reservedQueriers := make(map[string]struct{}, numOfReservedQueriers)
-
-	//if queriersToSelect == 0 || len(allSortedQueriers) <= queriersToSelect {
-	//	for i := 0; i < numOfReservedQueriers; i++ {
-	//		reservedQueriers[allSortedQueriers[i]] = struct{}{}
-	//	}
-	//	return nil, reservedQueriers
-	//}
-
 	queriers := make(map[string]struct{}, queriersToSelect)
 	rnd := rand.New(rand.NewSource(userSeed))
 
@@ -340,41 +335,13 @@ func shuffleQueriersForUser(userSeed int64, queriersToSelect int, allSortedQueri
 	for i := 0; i < queriersToSelect; i++ {
 		r := rnd.Intn(last + 1)
 		queriers[scratchpad[r]] = struct{}{}
-		//if i < numOfReservedQueriers {
-		//	reservedQueriers[scratchpad[r]] = struct{}{}
-		//}
-		// move selected item to the end, it won't be selected anymore.
 		scratchpad[r], scratchpad[last] = scratchpad[last], scratchpad[r]
 		last--
 	}
 
 	return queriers
 }
 
-func getReservedQueriers(queriers map[string]struct{}, priorities []validation.PriorityDef) map[string]int64 {
-	// TODO: justinjung04
-	return map[string]int64{}
-}
-
-// TODO: justinjung04
-//func getNumOfReservedQueriers(queriersToSelect int, totalNumOfQueriers int, reservedQueriers float64) int {
-//	numOfReservedQueriers := int(reservedQueriers)
-//
-//	if reservedQueriers < 1 && reservedQueriers > 0 {
-//		if queriersToSelect == 0 || queriersToSelect > totalNumOfQueriers {
-//			queriersToSelect = totalNumOfQueriers
-//		}
-//
-//		numOfReservedQueriers = int(math.Ceil(float64(queriersToSelect) * reservedQueriers))
-//	}
-//
-//	if numOfReservedQueriers > totalNumOfQueriers {
-//		return totalNumOfQueriers
-//	}
-//
-//	return numOfReservedQueriers
-//}
-
 // MockLimits implements the Limits interface. Used in tests only.
 type MockLimits struct {
 	MaxOutstanding     int

pkg/scheduler/queue/user_queues_test.go

@@ -360,16 +360,16 @@ func generateQuerier(r *rand.Rand) string {
 	return fmt.Sprint("querier-", r.Int()%5)
 }
 
-func getOrAdd(t *testing.T, uq *queues, tenant string, maxQueriers int) requestQueue {
+func getOrAdd(t *testing.T, uq *queues, tenant string, maxQueriers int) userRequestQueue {
 	q := uq.getOrAddQueue(tenant, maxQueriers)
 	assert.NotNil(t, q)
 	assert.NoError(t, isConsistent(uq))
 	assert.Equal(t, q, uq.getOrAddQueue(tenant, maxQueriers))
 	return q
 }
 
-func confirmOrderForQuerier(t *testing.T, uq *queues, querier string, lastUserIndex int, qs ...requestQueue) int {
-	var n requestQueue
+func confirmOrderForQuerier(t *testing.T, uq *queues, querier string, lastUserIndex int, qs ...userRequestQueue) int {
+	var n userRequestQueue
 	for _, q := range qs {
 		n, _, lastUserIndex = uq.getNextQueueForQuerier(lastUserIndex, querier)
 		assert.Equal(t, q, n)

pkg/scheduler/queue/request_queue.go renamed to pkg/scheduler/queue/user_request_queue.go

@@ -2,9 +2,9 @@ package queue
 
 import "github.com/cortexproject/cortex/pkg/util"
 
-type requestQueue interface {
+type userRequestQueue interface {
 	enqueueRequest(Request)
-	dequeueRequest() Request
+	dequeueRequest(minPriority int64) Request
 	length() int
 	closeQueue()
 }
@@ -13,15 +13,15 @@ type FIFORequestQueue struct {
 	queue chan Request
 }
 
-func NewFIFOQueue(queue chan Request) *FIFORequestQueue {
+func NewFIFORequestQueue(queue chan Request) *FIFORequestQueue {
 	return &FIFORequestQueue{queue: queue}
 }
 
 func (f *FIFORequestQueue) enqueueRequest(r Request) {
 	f.queue <- r
 }
 
-func (f *FIFORequestQueue) dequeueRequest() Request {
+func (f *FIFORequestQueue) dequeueRequest(_ int64) Request {
 	return <-f.queue
 }
 
@@ -45,7 +45,10 @@ func (f *PriorityRequestQueue) enqueueRequest(r Request) {
 	f.queue.Enqueue(r)
 }
 
-func (f *PriorityRequestQueue) dequeueRequest() Request {
+func (f *PriorityRequestQueue) dequeueRequest(minPriority int64) Request {
+	if f.queue.Peek().Priority() < minPriority {
+		return nil
+	}
 	return f.queue.Dequeue()
 }
 

pkg/scheduler/queue/user_request_queue_test.go

@@ -0,0 +1,57 @@
+package queue
+
+import (
+	"github.com/cortexproject/cortex/pkg/util"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestFIFORequestQueue(t *testing.T) {
+	queue := NewFIFORequestQueue(make(chan Request, 2))
+	request1 := MockRequest{
+		id:       "request 1",
+		priority: 1,
+	}
+	request2 := MockRequest{
+		id:       "request 2",
+		priority: 2,
+	}
+
+	queue.enqueueRequest(request1)
+	queue.enqueueRequest(request2)
+	assert.Equal(t, 2, queue.length())
+	assert.Equal(t, request1, queue.dequeueRequest(0))
+	assert.Equal(t, 1, queue.length())
+	assert.Equal(t, request2, queue.dequeueRequest(0))
+	assert.Equal(t, 0, queue.length())
+	queue.closeQueue()
+	assert.Panics(t, func() { queue.enqueueRequest(request1) })
+}
+
+func TestPriorityRequestQueue(t *testing.T) {
+	queue := NewPriorityRequestQueue(util.NewPriorityQueue(nil))
+	request1 := MockRequest{
+		id:       "request 1",
+		priority: 1,
+	}
+	request2 := MockRequest{
+		id:       "request 2",
+		priority: 2,
+	}
+
+	queue.enqueueRequest(request1)
+	queue.enqueueRequest(request2)
+	assert.Equal(t, 2, queue.length())
+	assert.Equal(t, request2, queue.dequeueRequest(0))
+	assert.Equal(t, 1, queue.length())
+	assert.Equal(t, request1, queue.dequeueRequest(0))
+	assert.Equal(t, 0, queue.length())
+
+	queue.enqueueRequest(request1)
+	queue.enqueueRequest(request2)
+	assert.Equal(t, 2, queue.length())
+	assert.Equal(t, request2, queue.dequeueRequest(2))
+	
+	queue.closeQueue()
+	assert.Panics(t, func() { queue.enqueueRequest(request1) })
+}

pkg/util/priority_queue.go

@@ -93,7 +93,7 @@ func (pq *PriorityQueue) Enqueue(op PriorityOp) {
 	}
 }
 
-// Dequeue will return the op with the highest priority; block if queue is
+// Dequeue will remove and return the op with the highest priority; block if queue is
 // empty; returns nil if queue is closed.
 func (pq *PriorityQueue) Dequeue() PriorityOp {
 	pq.lock.Lock()
@@ -114,3 +114,9 @@ func (pq *PriorityQueue) Dequeue() PriorityOp {
 	}
 	return op
 }
+
+// Peek will return the op with the highest priority without removing it from the queue
+func (pq *PriorityQueue) Peek() PriorityOp {
+	op := pq.queue[0]
+	return op
+}

pkg/util/priority_queue_test.go

@@ -2,7 +2,6 @@ package util
 
 import (
 	"runtime"
-	"strconv"
 	"testing"
 	"time"
 
@@ -15,10 +14,6 @@ func (i simpleItem) Priority() int64 {
 	return int64(i)
 }
 
-func (i simpleItem) Key() string {
-	return strconv.FormatInt(int64(i), 10)
-}
-
 func TestPriorityQueueBasic(t *testing.T) {
 	queue := NewPriorityQueue(nil)
 	assert.Equal(t, 0, queue.Length(), "Expected length = 0")
@@ -39,6 +34,7 @@ func TestPriorityQueuePriorities(t *testing.T) {
 	queue.Enqueue(simpleItem(1))
 	queue.Enqueue(simpleItem(2))
 
+	assert.Equal(t, simpleItem(2), queue.Peek().(simpleItem), "Expected to peek simpleItem(2)")
 	assert.Equal(t, simpleItem(2), queue.Dequeue().(simpleItem), "Expected to dequeue simpleItem(2)")
 	assert.Equal(t, simpleItem(1), queue.Dequeue().(simpleItem), "Expected to dequeue simpleItem(1)")
 
@@ -51,6 +47,7 @@ func TestPriorityQueuePriorities2(t *testing.T) {
 	queue.Enqueue(simpleItem(2))
 	queue.Enqueue(simpleItem(1))
 
+	assert.Equal(t, simpleItem(2), queue.Peek().(simpleItem), "Expected to peek simpleItem(2)")
 	assert.Equal(t, simpleItem(2), queue.Dequeue().(simpleItem), "Expected to dequeue simpleItem(2)")
 	assert.Equal(t, simpleItem(1), queue.Dequeue().(simpleItem), "Expected to dequeue simpleItem(1)")