Preserve waitlist behavior with Snake

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

@@ -33,10 +33,14 @@ type PoolConfig interface {
 
 // waiter represents a client waiting for a connection in the waitlist
 type waiter[C Connection] struct {
+	// setting is the connection Setting that we'd like, or nil if we'd like a
+	// a connection with no Setting applied
+	setting *Setting
 	// conn is a channel that will receive the connection when it's ready
-	conn    chan *Pooled[C]
-	request *loadshed.Request[*waiter[C]]
-	err     error
+	conn chan *Pooled[C]
+	err  error
+	// age is the amount of cycles this client has been on the waitlist
+	age uint32
 }
 
 type waitlist[C Connection] struct {
@@ -62,7 +66,7 @@ func (wl *waitlist[C]) waitForConn(ctx context.Context, setting *Setting, closeC
 	defer wl.nodes.Put(elem)
 
 	conn := elem.conn
-	*elem = waiter[C]{conn: conn}
+	*elem = waiter[C]{conn: conn, setting: setting}
 
 	// Fast path: reject early using an atomic read of the list length to avoid
 	// contending on the mutex under high query rates. This is racy — the count
@@ -99,9 +103,12 @@ func (wl *waitlist[C]) waitForConn(ctx context.Context, setting *Setting, closeC
 		}
 		return nil, ErrPoolWaiterCapReached
 	}
-	request, dropped := wl.snake.Enqueue(elem, valveID, snakePriority(priority))
-	elem.request = request
-	wl.syncTimer()
+	if priority != loadshed.PriorityUndroppable {
+		// Translate the Vitess proto priority (0 = most important) into Snake's
+		// convention (higher priority shed last).
+		priority = float64(sqlparser.MaxPriorityValue) - priority
+	}
+	request, dropped := wl.snake.Enqueue(elem, valveID, priority)
 	wl.mu.Unlock()
 	wl.reject(dropped)
 
@@ -110,8 +117,7 @@ func (wl *waitlist[C]) waitForConn(ctx context.Context, setting *Setting, closeC
 		// Pool was closed while we were waiting.
 		wl.mu.Lock()
 		// Try to find and remove ourselves from the list.
-		removed, dropped := wl.snake.Cancel(elem.request)
-		wl.syncTimer()
+		removed, dropped := wl.snake.Cancel(request)
 		wl.mu.Unlock()
 		wl.reject(dropped)
 
@@ -128,8 +134,7 @@ func (wl *waitlist[C]) waitForConn(ctx context.Context, setting *Setting, closeC
 		// prevent another goroutine from trying to hand us a connection later on.
 		wl.mu.Lock()
 		// Try to find and remove ourselves from the list.
-		removed, dropped := wl.snake.Cancel(elem.request)
-		wl.syncTimer()
+		removed, dropped := wl.snake.Cancel(request)
 		wl.mu.Unlock()
 		wl.reject(dropped)
 
@@ -150,15 +155,9 @@ 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()
+	// TODO: Remove the age/starvation code since Snake guarantees prompt grant-or-shed.
+	return 0
 }
 
 // tryReturnConn tries handing over a connection to one of the waiters in the pool.
@@ -172,9 +171,25 @@ func (wl *waitlist[D]) tryReturnConn(conn *Pooled[D]) bool {
 }
 
 func (wl *waitlist[D]) tryReturnConnSlow(conn *Pooled[D]) bool {
+	const maxAge = 8
+	connSetting := conn.Conn.Setting()
+
 	wl.mu.Lock()
-	waiter, ok, dropped := wl.snake.Dequeue()
-	wl.syncTimer()
+	// iterate through the waitlist looking for either waiters that have been
+	// here too long, or a waiter that is looking exactly for the same Setting
+	// as the one we have in our connection.
+	waiter, ok, dropped := wl.snake.DequeueMatching(func(waiter *waiter[D]) bool {
+		if waiter.age > maxAge || waiter.setting == connSetting {
+			return true
+		}
+		// this only ages the waiters that are being skipped over: we'll start
+		// aging the waiters in the back once they get to the front of the pool.
+		// the maxAge of 8 has been set empirically: smaller values cause clients
+		// with a specific setting to slightly starve, and aging all the clients
+		// in the list every time leads to unfairness when the system is at capacity
+		waiter.age++
+		return false
+	})
 	wl.mu.Unlock()
 	wl.reject(dropped)
 
@@ -200,16 +215,9 @@ func (wl *waitlist[C]) reject(waiters []*waiter[C]) {
 	}
 }
 
-func (wl *waitlist[C]) syncTimer() {
-	if delay, ok := wl.snake.LockedTimerUpdate(); ok {
-		time.AfterFunc(delay, wl.runDropTimer)
-	}
-}
-
 func (wl *waitlist[C]) runDropTimer() {
 	wl.mu.Lock()
 	dropped := wl.snake.LockedDropTimerFired()
-	wl.syncTimer()
 	wl.mu.Unlock()
 	wl.reject(dropped)
 }
@@ -228,6 +236,7 @@ func (wl *waitlist[C]) init(poolName string, config PoolConfig) {
 
 	wl.snake = loadshed.NewSnake[*waiter[C]](loadshed.SnakeConfig{
 		LoadsheddingAllowed: enabled,
+		DropTimerFired:      wl.runDropTimer,
 		CoDel: loadshed.CoDelConfig{
 			IntervalNs:     func() int64 { return interval().Nanoseconds() },
 			TargetNs:       func() int64 { return target().Nanoseconds() },

go/pools/smartconnpool/waitlist_test.go

@@ -139,9 +139,25 @@ func TestWaitlistShedsQueuedRequests(t *testing.T) {
 	}
 }
 
-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))
+func TestWaitlistPreservesSettingAffinityAndAging(t *testing.T) {
+	wl := waitlist[*TestConn]{}
+	wl.init("", nil)
+
+	foo := &waiter[*TestConn]{setting: sFoo, conn: make(chan *Pooled[*TestConn], 1)}
+	wl.snake.Enqueue(foo, "", loadshed.PriorityUndroppable)
+	bar := &waiter[*TestConn]{setting: sBar, conn: make(chan *Pooled[*TestConn], 1)}
+	wl.snake.Enqueue(bar, "", loadshed.PriorityUndroppable)
+	conn := &Pooled[*TestConn]{Conn: &TestConn{setting: sBar}}
+
+	require.True(t, wl.tryReturnConn(conn))
+	assert.Same(t, conn, <-bar.conn)
+	assert.Equal(t, uint32(1), foo.age)
+	assert.Equal(t, 0, wl.maybeStarvingCount())
+
+	foo.age = 9
+	bar = &waiter[*TestConn]{setting: sBar, conn: make(chan *Pooled[*TestConn], 1)}
+	wl.snake.Enqueue(bar, "", loadshed.PriorityUndroppable)
+
+	require.True(t, wl.tryReturnConn(conn))
+	assert.Same(t, conn, <-foo.conn)
 }

go/vt/vttablet/tabletserver/loadshed/snake.go

@@ -44,6 +44,7 @@ type (
 		dropTimerExpectedNs int64
 		cfg                 SnakeConfig
 		clockFunc           func() int64
+		length              atomic.Int64
 
 		shedCount atomic.Int64
 		// shedByPriority breaks shedCount down by the shed request's priority label
@@ -111,6 +112,7 @@ func (s *Snake[T]) Enqueue(value T, valveID string, priority float64) (*Request[
 
 	req := s.q.lockedEnqueue(valveID, priority)
 	req.value = value
+	s.length.Add(1)
 	if valveID != "" {
 		s.lockedObserveValveDepth(valveID)
 	}
@@ -142,6 +144,7 @@ func (s *Snake[T]) dequeue(match func(T) bool) (T, bool, []T) {
 	if req != nil {
 		s.q.lockedDequeue(req)
 		req.signal(grantSentinel)
+		s.length.Add(-1)
 		now := s.clockFunc()
 		s.lockedAccrueDropping(now)
 		s.sojourn.Add(now - req.codelqEnqueuedAtNs)
@@ -155,11 +158,16 @@ func (s *Snake[T]) dequeue(match func(T) bool) (T, bool, []T) {
 	return value, ok, s.droppedValues(pending)
 }
 
+func (s *Snake[T]) Len() int {
+	return int(s.length.Load())
+}
+
 func (s *Snake[T]) Cancel(req *Request[T]) (bool, []T) {
 	if req.signaledValue != nil {
 		return false, nil
 	}
 	s.q.lockedCancel(req)
+	s.length.Add(-1)
 	var zero T
 	req.value = zero
 	dropped := s.q.lockedTakePendingDrops()
@@ -207,6 +215,7 @@ func (s *Snake[T]) droppedValues(requests []*Request[T]) []T {
 	}
 	values := make([]T, len(requests))
 	for i, req := range requests {
+		s.length.Add(-1)
 		s.shedCount.Add(1)
 		if s.shedByPriority != nil {
 			s.shedByPriority.Add([]string{shedPriorityLabel(req.priority)}, 1)

go/vt/vttablet/tabletserver/query_executor.go

@@ -827,6 +827,8 @@ func (qre *QueryExecutor) getConn() (*connpool.PooledConn, error) {
 	}(time.Now())
 
 	priority := float64(priorityFromOptions(qre.options, qre.tsv.config.TxThrottlerDefaultPriority))
+	// Queries against a configured schema (e.g. performance_schema health
+	// checks) are marked undroppable instead, so they are never shed.
 	if matchesUndroppableSchema(qre.plan.SchemaQualifiers, qre.tsv.Config().LoadshedOltpRead.UndroppableSchemasValue()) {
 		priority = loadshed.PriorityUndroppable
 	}