router: fix session boost grace hold being skipped

pkg/kthena-router/datastore/fairness_queue.go

@@ -234,6 +234,7 @@ type RequestPriorityQueue struct {
 	podCounter     PodCounter            // Optional; counts backend pods for inflight scaling
 	inflightCount  atomic.Int64          // In-flight requests in session-boost mode
 	releaseCh      chan struct{}         // Signals a permit release in session-boost mode
+	releaseMu      sync.Mutex            // Makes a release's freed slot and its signal visible together
 	recheckCh      chan struct{}         // Signals refreshed backend state in session-boost mode
 }
 

pkg/kthena-router/datastore/session_boost_queue.go

@@ -161,11 +161,15 @@ func (pq *RequestPriorityQueue) admitSessionBoost(req *Request) bool {
 		releaseOnce := sync.Once{}
 		req.Release = func() {
 			releaseOnce.Do(func() {
+				// Free the slot and signal under releaseMu, so
+				// dequeueOrHoldPendingRelease sees both or neither.
+				pq.releaseMu.Lock()
 				pq.inflightCount.Add(-1)
 				select {
 				case pq.releaseCh <- struct{}{}:
 				default:
 				}
+				pq.releaseMu.Unlock()
 				pq.metricDecInflight(req.ModelName)
 			})
 		}
@@ -244,12 +248,24 @@ func (pq *RequestPriorityQueue) runSessionBoostMode(ctx context.Context) {
 // an arrival or a backend metrics refresh. With grace enabled, if a release is
 // also pending, it routes through the grace path so the freed slot is held for a
 // same-session follow-up; otherwise it re-checks admission directly.
+//
+// The release check and the inflight count are read together under releaseMu.
+// A slot freed after that read is left out of this pass, so it is still held for
+// a follow-up through the grace path.
 func (pq *RequestPriorityQueue) dequeueOrHoldPendingRelease(ctx context.Context, grace bool) {
-	if grace && pq.drainPendingRelease() {
+	if !grace {
+		pq.tryBackpressureDequeue(ctx)
+		return
+	}
+	pq.releaseMu.Lock()
+	released := pq.drainPendingRelease()
+	inflight := pq.inflightCount.Load()
+	pq.releaseMu.Unlock()
+	if released {
 		pq.waitGraceAndDequeue(ctx)
 		return
 	}
-	pq.tryBackpressureDequeue(ctx)
+	pq.tryBackpressureDequeueAbove(ctx, inflight)
 }
 
 // RecheckBackpressure asks the session-boost dequeue loop to re-evaluate held
@@ -350,6 +366,13 @@ func (pq *RequestPriorityQueue) drainCancelledLocked() int {
 // the queue is empty. This avoids the one-request-per-tick bottleneck during
 // initial ramp-up and whenever spare capacity exists.
 func (pq *RequestPriorityQueue) tryBackpressureDequeue(ctx context.Context) {
+	pq.tryBackpressureDequeueAbove(ctx, 0)
+}
+
+// tryBackpressureDequeueAbove is tryBackpressureDequeue with a lower bound for the
+// inflight count: the count the caller read, plus the requests admitted in this
+// pass. Slots freed after the caller's read are not used here.
+func (pq *RequestPriorityQueue) tryBackpressureDequeueAbove(ctx context.Context, minInflight int64) {
 	// In session-boost mode, the total inflight limit is InflightPerPod scaled by
 	// the number of backend pods serving the model.
 	perPod := pq.config.InflightPerPod
@@ -366,7 +389,7 @@ func (pq *RequestPriorityQueue) tryBackpressureDequeue(ctx context.Context) {
 	}
 
 	for {
-		currentInflight := pq.inflightCount.Load()
+		currentInflight := max(pq.inflightCount.Load(), minInflight)
 
 		if currentInflight >= maxInflight {
 			pq.mu.Lock()
@@ -405,6 +428,7 @@ func (pq *RequestPriorityQueue) tryBackpressureDequeue(ctx context.Context) {
 			// continue admitting the next queued request.
 			continue
 		}
+		minInflight++
 
 		klog.V(4).Infof("[SessionBoost] backpressure dequeue: user=%s model=%s sessionBoost=%v inflight=%d/%d",
 			req.UserID, req.ModelName, req.SessionBoost, pq.inflightCount.Load(), maxInflight)

pkg/kthena-router/datastore/session_boost_queue_test.go

@@ -574,22 +574,25 @@ func TestSessionBoostQueue_GracePeriod_BoostedArrives(t *testing.T) {
 }
 
 // TestSessionBoostQueue_DequeueOrHoldPendingRelease covers the handling of an
-// arrival or a recheck. When a release is pending at the same time, select may
-// pick either signal first, so this handling must not hand the freed slot to the
-// head of the queue without the grace wait. The context is cancelled up front:
-// a grace wait then returns without dispatching, while a direct dequeue still
-// admits, so whether the held request is admitted shows which path was taken.
+// arrival or a recheck. When a release is pending at the same time (select may
+// pick either signal first), or lands while the handling runs, this handling must
+// not hand the freed slot to the head of the queue without the grace wait. The
+// context is cancelled up front: a grace wait then returns without dispatching,
+// while a direct dequeue still admits, so whether the held request is admitted
+// shows which path was taken.
 func TestSessionBoostQueue_DequeueOrHoldPendingRelease(t *testing.T) {
 	tests := []struct {
-		name               string
-		gracePeriod        time.Duration
-		releasePending     bool
-		wantAdmitted       bool
-		wantReleasePending bool
+		name                  string
+		gracePeriod           time.Duration
+		releasePending        bool
+		releaseDuringHandling bool
+		wantAdmitted          bool
+		wantReleasePending    bool
 	}{
 		{name: "pending release holds the freed slot for grace", gracePeriod: time.Hour, releasePending: true, wantAdmitted: false, wantReleasePending: false},
 		{name: "no pending release admits directly", gracePeriod: time.Hour, releasePending: false, wantAdmitted: true, wantReleasePending: false},
 		{name: "grace disabled admits directly", gracePeriod: 0, releasePending: true, wantAdmitted: true, wantReleasePending: true},
+		{name: "release during handling holds the freed slot for grace", gracePeriod: time.Hour, releaseDuringHandling: true, wantAdmitted: false, wantReleasePending: true},
 	}
 
 	for _, tt := range tests {
@@ -619,11 +622,20 @@ func TestSessionBoostQueue_DequeueOrHoldPendingRelease(t *testing.T) {
 			}
 			<-q.notifyCh
 
-			// Free the slot. Without a pending release, its signal has already
-			// been handled.
-			req1.Release()
-			if !tt.releasePending {
-				<-q.releaseCh
+			if tt.releaseDuringHandling {
+				// podCounter runs inside the handling, after the release check,
+				// so freeing the slot from it lands the release in between.
+				q.podCounter = func() int {
+					req1.Release()
+					return 1
+				}
+			} else {
+				// Free the slot. Without a pending release, its signal has
+				// already been handled.
+				req1.Release()
+				if !tt.releasePending {
+					<-q.releaseCh
+				}
 			}
 
 			ctx, cancel := context.WithCancel(context.Background())