perf(flowcontrol): index active queues so dispatch scans O(active) flows

Every dispatch cycle, the fairness policy's Pick calls IterateQueues,
which snapshotted every registered queue in the band (empty ones
included) into a freshly allocated slice under the registry read lock.
Per-dispatch cost was O(registered flows) in CPU and allocation: 94-98%
of all allocation in the flow-control benchmarks, and superlinear
throughput decay under flow churn as registered flows accumulated
between GC cycles.

Each priority band now maintains a lock-free index of its non-empty
queues, updated on empty<->non-empty transitions detected inside the
queue's existing stats critical section (transitions are serialized per
queue by the queue mutex; the index is a sync.Map, so the update never
touches the registry mutex). IterateQueues ranges over the index
directly: no registry lock, no allocation, cost proportional to flows
that actually hold items.

This narrows IterateQueues to active (non-empty) queues, matching its
documented contract ("each active Flow"). Both production callers are
compatible: globalstrict.Pick already skips empty queues, and the
processor's cleanup/drain sweeps are no-ops on them. The view is
eventually consistent; callbacks must tolerate a visited queue reporting
Len() == 0, which they already did under the snapshot implementation.

Benchmarks (M4 Pro, -benchtime=3s unless noted):
- PerformanceMatrix L=1/P=1/F=5000/W=5000: 30.7k -> 277k d/s (9x),
  42.2KB -> 1.8KB per op
- TopologyChurn: 16.7k -> 260k d/s (15x), 83KB -> 1.8KB per op
- FullPath (registry churn): 2.3k -> 118k d/s (52x); per-op cost now
  flat with run length (was superlinear: 106us at 24k flows minted ->
  439us at 73k)

Signed-off-by: Luke Van Drie <lukevandrie@google.com>

Author: LukeAVanDrie

pkg/epp/flowcontrol/registry/managedqueue.go

@@ -63,6 +63,13 @@ type managedQueue struct {
 	// onStatsDelta is the callback used to propagate statistics changes up to the registry.
 	onStatsDelta propagateStatsDeltaFunc
 
+	// onActiveTransition is invoked when the queue transitions between empty and non-empty, so the
+	// owning priority band can maintain its index of active (non-empty) queues. Transitions are
+	// detected under `mu`, which serializes them per queue. The callback must be lock-free or take
+	// only leaf locks (never the registry mutex): it runs inside this queue's critical section.
+	// May be nil (e.g., in isolated unit tests).
+	onActiveTransition func(mq *managedQueue, active bool)
+
 	// --- State Protected by `mu` ---
 
 	// mu protects all mutating operations. It ensures that any changes to the underlying `queue` and the updates to the
@@ -92,14 +99,16 @@ func newManagedQueue(
 	key flowcontrol.FlowKey,
 	logger logr.Logger,
 	onStatsDelta propagateStatsDeltaFunc,
+	onActiveTransition func(mq *managedQueue, active bool),
 ) *managedQueue {
 	mqLogger := logger.WithName("managed-queue").WithValues("flowKey", key)
 	mq := &managedQueue{
-		queue:        queue,
-		policy:       policy,
-		key:          key,
-		onStatsDelta: onStatsDelta,
-		logger:       mqLogger,
+		queue:              queue,
+		policy:             policy,
+		key:                key,
+		onStatsDelta:       onStatsDelta,
+		onActiveTransition: onActiveTransition,
+		logger:             mqLogger,
 	}
 	mq.flowQueueAccessor = &flowQueueAccessor{mq: mq}
 	return mq
@@ -187,6 +196,17 @@ func (mq *managedQueue) propagateStatsDeltaLocked(lenDelta, byteSizeDelta int64)
 	}
 	mq.byteSize.Add(byteSizeDelta)
 
+	// Maintain the band's active-queue index on empty<->non-empty transitions. `mu` is held, so
+	// transitions are strictly alternating per queue; `newLen == lenDelta` implies the previous
+	// length was zero.
+	if mq.onActiveTransition != nil && lenDelta != 0 {
+		if lenDelta > 0 && newLen == lenDelta {
+			mq.onActiveTransition(mq, true)
+		} else if newLen == 0 {
+			mq.onActiveTransition(mq, false)
+		}
+	}
+
 	// Propagate the delta up to the registry. This propagation is lock-free and eventually consistent.
 	mq.onStatsDelta(mq.key.Priority, lenDelta, byteSizeDelta)
 }

pkg/epp/flowcontrol/registry/managedqueue_test.go

@@ -70,7 +70,7 @@ func newMqHarness(t *testing.T, queue contracts.SafeQueue, key flowcontrol.FlowK
 	propagator := &mockStatsPropagator{}
 	mockPolicy := &fwkfcmocks.MockOrderingPolicy{}
 
-	mq := newManagedQueue(queue, mockPolicy, key, logr.Discard(), propagator.propagate)
+	mq := newManagedQueue(queue, mockPolicy, key, logr.Discard(), propagator.propagate, nil)
 	require.NotNil(t, mq, "Test setup: newManagedQueue must return a valid instance")
 
 	return &mqTestHarness{

pkg/epp/flowcontrol/registry/registry_helpers.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"fmt"
 	"sort"
+	"sync"
 
 	"github.com/llm-d/llm-d-router/pkg/common/observability/logging"
 	"github.com/llm-d/llm-d-router/pkg/epp/flowcontrol/contracts"
@@ -50,6 +51,25 @@ type priorityBand struct {
 
 	// priorityBandAccessor is a preallocated flowcontrol.PriorityBandAccessor for this priorityBand
 	priorityBandAccessor *priorityBandAccessor
+
+	// activeQueues indexes the subset of `queues` that currently hold items, keyed by logical ID
+	// (values are *managedQueue). It is maintained by each queue's empty<->non-empty transitions
+	// (serialized per queue under the queue's own mutex) and read lock-free by IterateQueues, which
+	// keeps the dispatch hot path O(active flows) with zero allocation instead of O(registered
+	// flows) with a snapshot. The view is eventually consistent: a queue is always present here by
+	// the time an Add returns, but may linger briefly after draining; readers must tolerate
+	// observing an empty queue.
+	activeQueues sync.Map
+}
+
+// setQueueActivity is the onActiveTransition callback for this band's queues. It runs inside the
+// queue's critical section, so it must remain lock-free (sync.Map only, never the registry mutex).
+func (b *priorityBand) setQueueActivity(mq *managedQueue, active bool) {
+	if active {
+		b.activeQueues.Store(mq.key.ID, mq)
+	} else {
+		b.activeQueues.Delete(mq.key.ID)
+	}
 }
 
 // initPriorityBand constructs the runtime state for a single priority level and registers it within the registry.
@@ -161,7 +181,7 @@ func (fr *FlowRegistry) synchronizeFlow(
 
 	fr.logger.V(logging.TRACE).Info("Creating new queue for flow instance.", "flowKey", key)
 
-	mq := newManagedQueue(q, policy, key, fr.logger, fr.propagateStatsDelta)
+	mq := newManagedQueue(q, policy, key, fr.logger, fr.propagateStatsDelta, band.setQueueActivity)
 	band.queues[key.ID] = mq
 }
 
@@ -185,6 +205,7 @@ func (fr *FlowRegistry) deleteFlow(key flowcontrol.FlowKey) {
 			}
 		}
 		delete(band.queues, key.ID)
+		band.activeQueues.Delete(key.ID)
 	}
 }
 
@@ -244,22 +265,16 @@ func (a *priorityBandAccessor) Queue(id string) flowcontrol.FlowQueueAccessor {
 	return mq.FlowQueueAccessor()
 }
 
-// IterateQueues executes the given `callback` for each FlowQueueAccessor in this priority band.
+// IterateQueues executes the given `callback` for each active (non-empty) FlowQueueAccessor in
+// this priority band.
 //
-// To minimize lock contention, this implementation snapshots the queue accessors under a read lock and then executes
-// the callback on the snapshot, outside of the lock. This ensures that a potentially slow policy (the callback) does
-// not block other operations on the registry.
+// It ranges over the band's lock-free active-queue index, so it takes no registry lock and
+// performs no allocation, and its cost scales with the number of flows that currently hold items
+// rather than the number of registered flows. The view is eventually consistent: a queue drained
+// concurrently with iteration may still be visited, so callbacks must tolerate Len() == 0; a
+// queue is guaranteed to be visible once the Add that made it non-empty has returned.
 func (a *priorityBandAccessor) IterateQueues(callback func(queue flowcontrol.FlowQueueAccessor) bool) {
-	a.registry.mu.RLock()
-	accessors := make([]flowcontrol.FlowQueueAccessor, 0, len(a.band.queues))
-	for _, mq := range a.band.queues {
-		accessors = append(accessors, mq.FlowQueueAccessor())
-	}
-	a.registry.mu.RUnlock()
-
-	for _, accessor := range accessors {
-		if !callback(accessor) {
-			return
-		}
-	}
+	a.band.activeQueues.Range(func(_, v any) bool {
+		return callback(v.(*managedQueue).FlowQueueAccessor())
+	})
 }

pkg/epp/flowcontrol/registry/registry_helpers_test.go

@@ -264,20 +264,71 @@ func TestRegistry_PriorityBandAccessor(t *testing.T) {
 		t.Run("IterateQueues", func(t *testing.T) {
 			t.Parallel()
 
-			t.Run("ShouldVisitAllQueuesInBand", func(t *testing.T) {
+			t.Run("ShouldVisitAllActiveQueuesInBand", func(t *testing.T) {
 				t.Parallel()
+				h := newTestHarness(t) // Isolated harness: this subtest mutates queue contents.
+				accessor, err := h.registry.PriorityBandAccessor(highPriority)
+				require.NoError(t, err)
+				h.addItem(h.highPriorityKey1, 100)
+				h.addItem(h.highPriorityKey2, 100)
+
 				var iteratedKeys []flowcontrol.FlowKey
 				accessor.IterateQueues(func(queue flowcontrol.FlowQueueAccessor) bool {
 					iteratedKeys = append(iteratedKeys, queue.FlowKey())
 					return true
 				})
 				expectedKeys := []flowcontrol.FlowKey{h.highPriorityKey1, h.highPriorityKey2}
 				assert.ElementsMatch(t, expectedKeys, iteratedKeys,
-					"IterateQueues must visit every registered flow in the band exactly once")
+					"IterateQueues must visit every active (non-empty) flow in the band exactly once")
+			})
+
+			t.Run("ShouldSkipEmptyQueues", func(t *testing.T) {
+				t.Parallel()
+				h := newTestHarness(t) // Isolated harness: this subtest mutates queue contents.
+				accessor, err := h.registry.PriorityBandAccessor(highPriority)
+				require.NoError(t, err)
+				h.addItem(h.highPriorityKey1, 100)
+
+				var iteratedKeys []flowcontrol.FlowKey
+				accessor.IterateQueues(func(queue flowcontrol.FlowQueueAccessor) bool {
+					iteratedKeys = append(iteratedKeys, queue.FlowKey())
+					return true
+				})
+				assert.Equal(t, []flowcontrol.FlowKey{h.highPriorityKey1}, iteratedKeys,
+					"IterateQueues must visit only flows whose queues hold items; registered-but-empty flows are skipped")
+			})
+
+			t.Run("ShouldTrackEmptinessTransitions", func(t *testing.T) {
+				t.Parallel()
+				h := newTestHarness(t) // Isolated harness: this subtest mutates queue contents.
+				accessor, err := h.registry.PriorityBandAccessor(highPriority)
+				require.NoError(t, err)
+
+				countVisited := func() int {
+					var n int
+					accessor.IterateQueues(func(queue flowcontrol.FlowQueueAccessor) bool {
+						n++
+						return true
+					})
+					return n
+				}
+
+				item := h.addItem(h.highPriorityKey1, 100)
+				assert.Equal(t, 1, countVisited(), "A flow must become visible once its queue holds an item")
+				h.removeItem(h.highPriorityKey1, item)
+				assert.Equal(t, 0, countVisited(), "A flow must stop being visited once its queue drains")
+				h.addItem(h.highPriorityKey1, 100)
+				assert.Equal(t, 1, countVisited(), "A drained flow must become visible again on re-add")
 			})
 
 			t.Run("ShouldExitEarly_WhenCallbackReturnsFalse", func(t *testing.T) {
 				t.Parallel()
+				h := newTestHarness(t) // Isolated harness: this subtest mutates queue contents.
+				accessor, err := h.registry.PriorityBandAccessor(highPriority)
+				require.NoError(t, err)
+				h.addItem(h.highPriorityKey1, 100)
+				h.addItem(h.highPriorityKey2, 100)
+
 				var iterationCount int
 				accessor.IterateQueues(func(queue flowcontrol.FlowQueueAccessor) bool {
 					iterationCount++
@@ -307,12 +358,16 @@ func TestRegistry_PriorityBandAccessor(t *testing.T) {
 					}
 				}()
 
-				// Goroutine B: The Modifier (constantly writing)
+				// Goroutine B: The Modifier (constantly writing). Item add/remove cycles exercise the
+				// active-queue index transitions racing against iteration; flow create/delete cycles
+				// exercise registry topology changes.
 				go func() {
 					defer wg.Done()
 					for i := range 100 {
 						key := flowcontrol.FlowKey{ID: fmt.Sprintf("new-flow-%d", i), Priority: highPriority}
 						h.synchronizeFlow(key)
+						item := h.addItem(key, 100)
+						h.removeItem(key, item)
 						h.registry.mu.Lock()
 						h.registry.deleteFlow(key)
 						h.registry.mu.Unlock()

pkg/epp/framework/interface/flowcontrol/accessors.go

@@ -63,7 +63,9 @@ type PriorityBandAccessor interface {
 	// Returns nil if the ID is not found in this group.
 	Queue(id string) FlowQueueAccessor
 
-	// IterateQueues executes the given callback for each active Flow in this group.
+	// IterateQueues executes the given callback for each active Flow in this group. A Flow is active while its queue
+	// holds at least one item; implementations may skip empty queues entirely. The view may be eventually consistent
+	// under concurrent mutation, so callbacks must tolerate a visited queue reporting Len() == 0.
 	// Iteration stops if the callback returns false. The order of iteration is not guaranteed unless specified by
 	// the implementation.
 	IterateQueues(callback func(flow FlowQueueAccessor) (keepIterating bool))