Parameterize loadshed request values

AI assisted with development. Every line of code was either written by or carefully reviewed by me :)

Signed-off-by: Brett Wines <bwines@slack-corp.com>

Author: bgwines

go/pools/smartconnpool/waitlist.go

@@ -23,6 +23,7 @@ import (
 	"time"
 
 	"vitess.io/vitess/go/vt/servenv"
+	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/loadshed"
 )
 
@@ -34,7 +35,7 @@ type PoolConfig interface {
 type waiter[C Connection] struct {
 	// conn is a channel that will receive the connection when it's ready
 	conn    chan *Pooled[C]
-	request *loadshed.Request
+	request *loadshed.Request[*waiter[C]]
 	err     error
 }
 
@@ -98,7 +99,7 @@ func (wl *waitlist[C]) waitForConn(ctx context.Context, setting *Setting, closeC
 		}
 		return nil, ErrPoolWaiterCapReached
 	}
-	request, dropped := wl.snake.Enqueue(elem, valveID, priority)
+	request, dropped := wl.snake.Enqueue(elem, valveID, snakePriority(priority))
 	elem.request = request
 	wl.syncTimer()
 	wl.mu.Unlock()
@@ -149,6 +150,13 @@ func (wl *waitlist[C]) aboveWaiterCap(maxWaiters uint) bool {
 	return maxWaiters > 0 && wl.snake.Len() >= int(maxWaiters)
 }
 
+func snakePriority(priority float64) float64 {
+	if priority == loadshed.PriorityUndroppable {
+		return priority
+	}
+	return float64(sqlparser.MaxPriorityValue) - priority
+}
+
 func (wl *waitlist[C]) maybeStarvingCount() int {
 	return wl.snake.Len()
 }

go/pools/smartconnpool/waitlist_test.go

@@ -138,3 +138,10 @@ func TestWaitlistShedsQueuedRequests(t *testing.T) {
 		<-errs
 	}
 }
+
+func TestSnakePriority(t *testing.T) {
+	assert.Equal(t, float64(100), snakePriority(0))
+	assert.Equal(t, float64(70), snakePriority(30))
+	assert.Equal(t, float64(0), snakePriority(100))
+	assert.Equal(t, loadshed.PriorityUndroppable, snakePriority(loadshed.PriorityUndroppable))
+}

go/vt/vttablet/tabletserver/loadshed/codelq.go

@@ -144,7 +144,7 @@ type (
 	// CoDelQueue implements the CoDel (Controlled Delay) load-shedding
 	// algorithm. All methods are prefixed locked* and assume the caller holds
 	// the mutex, which is defined in the files for the higher-level structure
-	CoDelQueue struct {
+	CoDelQueue[T any] struct {
 		queue        *list.List
 		dropping     bool
 		dropNextNs   int64
@@ -153,16 +153,16 @@ type (
 		// droppable indexes the droppable queue entries by priority so the
 		// lowest-priority one is found in O(1) rather than an O(n) scan. Kept in
 		// lockstep with droppableLen: every insert/remove pairs with a ++/--.
-		droppable droppableIndex
+		droppable droppableIndex[T]
 
 		cfg               CoDelConfig
 		nowNs             func() int64
 		scheduleDropTimer func(delayNs int64)
 	}
 )
 
-func newCoDelQueue(cfg CoDelConfig, nowNs func() int64, scheduleDropTimer func(delayNs int64)) *CoDelQueue {
-	q := &CoDelQueue{
+func newCoDelQueue[T any](cfg CoDelConfig, nowNs func() int64, scheduleDropTimer func(delayNs int64)) *CoDelQueue[T] {
+	q := &CoDelQueue[T]{
 		queue:             list.New(),
 		count:             1,
 		cfg:               cfg,
@@ -173,15 +173,15 @@ func newCoDelQueue(cfg CoDelConfig, nowNs func() int64, scheduleDropTimer func(d
 	return q
 }
 
-func (q *CoDelQueue) lockedLen() int {
+func (q *CoDelQueue[T]) lockedLen() int {
 	return q.queue.Len()
 }
 
-func (q *CoDelQueue) lockedIsHealthy() bool {
+func (q *CoDelQueue[T]) lockedIsHealthy() bool {
 	return !q.dropping
 }
 
-func (q *CoDelQueue) lockedEnqueue(req *Request) {
+func (q *CoDelQueue[T]) lockedEnqueue(req *Request[T]) {
 	now := q.nowNs()
 
 	req.codelqEnqueuedAtNs = now
@@ -204,16 +204,16 @@ func (q *CoDelQueue) lockedEnqueue(req *Request) {
 }
 
 // lockedFirstWaiting returns the first waiting request in the queue.
-func (q *CoDelQueue) lockedFirstWaiting() *Request {
+func (q *CoDelQueue[T]) lockedFirstWaiting() *Request[T] {
 	first := q.queue.Front()
 	if first == nil {
 		return nil
 	}
-	return first.Value.(*Request)
+	return first.Value.(*Request[T])
 }
 
 // lockedRemove removes a specific request from the queue.
-func (q *CoDelQueue) lockedRemove(r *Request) {
+func (q *CoDelQueue[T]) lockedRemove(r *Request[T]) {
 	if r.codelqElem == nil {
 		return
 	}
@@ -230,7 +230,7 @@ func (q *CoDelQueue) lockedRemove(r *Request) {
 	}
 }
 
-func (q *CoDelQueue) lockedDequeue(r *Request) {
+func (q *CoDelQueue[T]) lockedDequeue(r *Request[T]) {
 	// CoDel health check, measured at dequeue: if this request's queue-wait
 	// (now - enqueue) was under target, the system is healthy — leave the
 	// dropping state. Separate from the droppableLen==0 clear below.
@@ -243,7 +243,7 @@ func (q *CoDelQueue) lockedDequeue(r *Request) {
 // lockedFindLowestPriorityDroppable finds the lowest-priority droppable
 // element in the queue — the oldest one at the lowest priority present — or nil
 // if none exists. O(1) via the droppable priority index (see droppableIndex).
-func (q *CoDelQueue) lockedFindLowestPriorityDroppable() *list.Element {
+func (q *CoDelQueue[T]) lockedFindLowestPriorityDroppable() *list.Element {
 	req := q.droppable.min()
 	if req == nil {
 		return nil
@@ -254,11 +254,11 @@ func (q *CoDelQueue) lockedFindLowestPriorityDroppable() *list.Element {
 // lockedRunTimer runs the CoDel drop logic. It is invoked both by the backstop
 // timer and synchronously from the dequeue path, so shedding is driven
 // as slots free rather than waiting for the (possibly late) timer to fire.
-func (q *CoDelQueue) lockedRunTimer(dropFn func() bool) {
+func (q *CoDelQueue[T]) lockedRunTimer(dropFn func() bool) {
 	q.lockedRunTimerLimited(dropFn, -1)
 }
 
-func (q *CoDelQueue) lockedRunTimerLimited(dropFn func() bool, maxDrops int) {
+func (q *CoDelQueue[T]) lockedRunTimerLimited(dropFn func() bool, maxDrops int) {
 	now := q.nowNs()
 
 	// Paced work: only advance the drop/ease control law and re-arm when a drop
@@ -282,11 +282,11 @@ func (q *CoDelQueue) lockedRunTimerLimited(dropFn func() bool, maxDrops int) {
 // `now`, so calling it is idempotent and safe outside the timer — the dequeue
 // path invokes it to shed stale requests in real time rather than
 // waiting on the possibly-late backstop timer. It does NOT arm/disarm the timer.
-func (q *CoDelQueue) lockedAdvance(now int64, dropFn func() bool) {
+func (q *CoDelQueue[T]) lockedAdvance(now int64, dropFn func() bool) {
 	q.lockedAdvanceLimited(now, dropFn, -1)
 }
 
-func (q *CoDelQueue) lockedAdvanceLimited(now int64, dropFn func() bool, maxDrops int) {
+func (q *CoDelQueue[T]) lockedAdvanceLimited(now int64, dropFn func() bool, maxDrops int) {
 	drops := 0
 	// Step the control law per interval while a drop is due AND there is still
 	// work to do: either a droppable backlog to shed, or an elevated count that
@@ -325,7 +325,7 @@ func (q *CoDelQueue) lockedAdvanceLimited(now int64, dropFn func() bool, maxDrop
 // lockedEaseCount returns the next drop count during easing:
 // count -= floor(log_base(count) / base), floored at 1. A larger base yields a
 // smaller step (gentler ease-out); base defaults to 3 when unset or <= 1.
-func (q *CoDelQueue) lockedEaseCount() int {
+func (q *CoDelQueue[T]) lockedEaseCount() int {
 	base := 3.0
 	if q.cfg.EasingLogBase != nil {
 		base = q.cfg.EasingLogBase()
@@ -341,15 +341,15 @@ func (q *CoDelQueue) lockedEaseCount() int {
 // inverse proportion to count^exponent, exploiting the non-linear
 // relationship between drop rate and throughput to achieve linear change
 // in throughput.
-func (q *CoDelQueue) lockedControlLaw(t int64) int64 {
+func (q *CoDelQueue[T]) lockedControlLaw(t int64) int64 {
 	return t + q.lockedCurrentInterval()
 }
 
 // lockedCurrentInterval returns the current interval for the control law.
 // The interval is compressed whenever count > 1 — both in the dropping state
 // and during easing (!dropping, count > 1), so that the ease-out timer fires
 // at progressively longer intervals as the count decays toward 1.
-func (q *CoDelQueue) lockedCurrentInterval() int64 {
+func (q *CoDelQueue[T]) lockedCurrentInterval() int64 {
 	interval := q.cfg.IntervalNs()
 	if q.count <= 1 {
 		return interval
@@ -360,7 +360,7 @@ func (q *CoDelQueue) lockedCurrentInterval() int64 {
 	return max(result, 1)
 }
 
-func (q *CoDelQueue) lockedArmDropTimer() {
+func (q *CoDelQueue[T]) lockedArmDropTimer() {
 	// Mark the episode active for this armed interval; the next timer fire
 	// re-evaluates health.
 	q.dropping = q.droppableLen > 0

go/vt/vttablet/tabletserver/loadshed/codelq_test.go

@@ -24,6 +24,9 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+type testRequest = Request[struct{}]
+type testCoDelQueue = CoDelQueue[struct{}]
+
 func defaultTestConfig() CoDelConfig {
 	return CoDelConfig{
 		IntervalNs:     func() int64 { return int64(1e9) },
@@ -74,14 +77,14 @@ func (r *testDropTimerRecorder) reset() {
 	r.scheduled = false
 }
 
-func newTestQueue(cfg CoDelConfig, clock *testClock) (*CoDelQueue, *testDropTimerRecorder) {
+func newTestQueue(cfg CoDelConfig, clock *testClock) (*testCoDelQueue, *testDropTimerRecorder) {
 	rec := &testDropTimerRecorder{}
-	q := newCoDelQueue(cfg, clock.nowFunc, rec.schedule)
+	q := newCoDelQueue[struct{}](cfg, clock.nowFunc, rec.schedule)
 	return q, rec
 }
 
-func testEnqueue(q *CoDelQueue, priority float64) *Request {
-	req := newRequest(priority)
+func testEnqueue(q *testCoDelQueue, priority float64) *testRequest {
+	req := newRequest[struct{}](priority)
 	q.lockedEnqueue(req)
 	return req
 }
@@ -135,7 +138,7 @@ func TestCoDelQueue_Enqueue_UndroppableNoSchedule(t *testing.T) {
 }
 
 // testDequeue removes the oldest waiting request.
-func testDequeue(q *CoDelQueue) *Request {
+func testDequeue(q *testCoDelQueue) *testRequest {
 	req := q.lockedFirstWaiting()
 	if req == nil {
 		return nil
@@ -226,7 +229,7 @@ func TestCoDelQueue_FindLowestPriorityDroppable_Basic(t *testing.T) {
 
 	elem := q.lockedFindLowestPriorityDroppable()
 	require.NotNil(t, elem)
-	dropped := elem.Value.(*Request)
+	dropped := elem.Value.(*testRequest)
 	q.lockedRemove(dropped)
 	assert.Equal(t, float64(1), dropped.priority)
 	assert.Equal(t, 2, q.lockedLen())
@@ -243,7 +246,7 @@ func TestCoDelQueue_FindLowestPriorityDroppable_ZeroInstantPick(t *testing.T) {
 
 	elem := q.lockedFindLowestPriorityDroppable()
 	require.NotNil(t, elem)
-	assert.Same(t, r2, elem.Value.(*Request))
+	assert.Same(t, r2, elem.Value.(*testRequest))
 }
 
 func TestCoDelQueue_DropSkipsUndroppable(t *testing.T) {
@@ -255,7 +258,7 @@ func TestCoDelQueue_DropSkipsUndroppable(t *testing.T) {
 
 	elem := q.lockedFindLowestPriorityDroppable()
 	require.NotNil(t, elem)
-	assert.Same(t, droppable, elem.Value.(*Request))
+	assert.Same(t, droppable, elem.Value.(*testRequest))
 	q.lockedRemove(droppable)
 	assert.Equal(t, 1, q.lockedLen())
 }
@@ -280,7 +283,7 @@ func TestCoDelQueue_DropUndroppableVsInf(t *testing.T) {
 
 	elem := q.lockedFindLowestPriorityDroppable()
 	require.NotNil(t, elem)
-	assert.Same(t, inf, elem.Value.(*Request))
+	assert.Same(t, inf, elem.Value.(*testRequest))
 }
 
 func TestCoDelQueue_DropAllInf_NoPanic(t *testing.T) {
@@ -351,7 +354,7 @@ func TestCoDelQueue_RunScheduledDrop_EntersDropping(t *testing.T) {
 		if elem == nil {
 			return false
 		}
-		q.lockedRemove(elem.Value.(*Request))
+		q.lockedRemove(elem.Value.(*testRequest))
 		return true
 	}
 	rec.reset()
@@ -373,7 +376,7 @@ func TestCoDelQueue_RunScheduledDrop_NothingDroppable(t *testing.T) {
 		if elem == nil {
 			return false
 		}
-		q.lockedRemove(elem.Value.(*Request))
+		q.lockedRemove(elem.Value.(*testRequest))
 		return true
 	}
 	q.lockedRunTimer(dropFn)
@@ -667,7 +670,7 @@ func TestCoDelQueue_Easing_DroppableLen_ReentersDroppingWithCurrentCount(t *test
 		if elem == nil {
 			return false
 		}
-		q.lockedRemove(elem.Value.(*Request))
+		q.lockedRemove(elem.Value.(*testRequest))
 		return true
 	}
 
@@ -794,7 +797,7 @@ func TestCoDelQueue_SlowMoving_Drops(t *testing.T) {
 		if elem == nil {
 			return false
 		}
-		q.lockedRemove(elem.Value.(*Request))
+		q.lockedRemove(elem.Value.(*testRequest))
 		return true
 	}
 	q.lockedRunTimer(dropFn)

go/vt/vttablet/tabletserver/loadshed/dequeue_shed_test.go

@@ -24,13 +24,13 @@ import (
 )
 
 // dropAll is the standard test dropFn: sheds the lowest-priority droppable head.
-func dropAllFn(q *CoDelQueue) func() bool {
+func dropAllFn(q *testCoDelQueue) func() bool {
 	return func() bool {
 		elem := q.lockedFindLowestPriorityDroppable()
 		if elem == nil {
 			return false
 		}
-		q.lockedRemove(elem.Value.(*Request))
+		q.lockedRemove(elem.Value.(*testRequest))
 		return true
 	}
 }
@@ -117,7 +117,7 @@ func TestValved_DisabledDropAdvancesCoDelWithoutDropping(t *testing.T) {
 	sq.codelq.cfg.IntervalNs = func() int64 { return 10_000_000 }
 
 	const backlog = 5
-	reqs := make([]*Request, backlog)
+	reqs := make([]*testRequest, backlog)
 	for i := range reqs {
 		reqs[i] = sq.lockedEnqueue(string(rune('a'+i)), 0)
 	}
@@ -142,7 +142,7 @@ func TestValved_EnablementSnapshottedOncePerBatch(t *testing.T) {
 	sq.codelq.cfg.IntervalNs = func() int64 { return 10_000_000 }
 
 	const backlog = 10
-	reqs := make([]*Request, backlog)
+	reqs := make([]*testRequest, backlog)
 	for i := range reqs {
 		reqs[i] = sq.lockedEnqueue(string(rune('a'+i)), 0)
 	}

go/vt/vttablet/tabletserver/loadshed/request.go

@@ -24,13 +24,13 @@ import (
 type (
 	// Request represents an entry in the CoDel queue. Named a 'request' since
 	// it may be dropped or dequeued.
-	Request struct {
+	Request[T any] struct {
 		priority           float64
 		codelqEnqueuedAtNs int64
 		codelqElem         *list.Element
 		valveID            string
 		queued             bool
-		value              any
+		value              T
 
 		// bucketElem locates this request in the droppableIndex while it is a
 		// droppable queue entry: it is the request's node in its priority
@@ -48,12 +48,12 @@ type (
 // schemas).
 var PriorityUndroppable = math.Inf(-1)
 
-func newRequest(priority float64) *Request {
-	return &Request{
+func newRequest[T any](priority float64) *Request[T] {
+	return &Request[T]{
 		priority: priority,
 	}
 }
 
-func (r *Request) isDroppable() bool {
+func (r *Request[T]) isDroppable() bool {
 	return r.priority != PriorityUndroppable
 }