refactor(flowcontrol): single-ended SafeQueue backed by container/heap by LukeAVanDrie · Pull Request #1690 · llm-d/llm-d-router

LukeAVanDrie · 2026-06-18T21:14:58Z

What type of PR is this?

/kind cleanup

What this PR does / why we need it:

The flow-control SafeQueue was double-ended (PeekHead + PeekTail), backed by a hand-rolled max-min heap. @RishabhSaini identified a latent bug in that bespoke data structure (#1689). In practice nothing consumes PeekTail, the only reference was the flowQueueAccessor passthrough itself. Dispatch (EDF/SLO ordering) needs only the head,
and eviction uses an independent ordering, so a double-ended dispatch queue buys nothing.

This makes the queue single-ended and swaps the custom heap for the standard library:

Remove PeekTail from QueueInspectionMethods and every implementer, mock, and test.
Rename PeekHead -> Peek now that there is a single end.
Replace maxminheap.go with priorityqueue.go, a container/heap-backed priority queue (PriorityQueue). This deletes the custom sift logic and the latent bug class with it, while keeping O(log n) targeted removal via per-item index tracking.
Drop the internal handle map; Remove validates membership by index identity, preserving the existing error contract with less state.
Rename the registered queue MaxMinHeap -> PriorityQueue.

container/heap with index tracking is already the established pattern in flowcontrol/eviction/queue.go, so this aligns the dispatch queue with existing style. FCFS / ListQueue behavior is unchanged.

Which issue(s) this PR fixes:

N/A, but a more aggressive refactoring that resolves the issue surfaced in #1689.

Release note (write NONE if no user-facing change):

flowcontrol: the priority-ordered dispatch queue is now backed by the standard library container/heap and registered as `PriorityQueue` (previously `MaxMinHeap`). Any configuration that explicitly set the queue to `MaxMinHeap` must be updated to `PriorityQueue`.

The dispatch SafeQueue was double-ended, but nothing consumed PeekTail other than the flowQueueAccessor passthrough itself: dispatch needs only the head, and eviction uses an independent ordering. With PeekTail gone the queue is single-ended, so the bespoke max-min heap that gave O(1) access to both ends is no longer justified. - Remove the unused PeekTail from the SafeQueue contract and all implementers, mocks, and tests; rename PeekHead -> Peek now that there is a single end. - Replace the hand-rolled max-min heap (maxminheap.go) with a standard library container/heap priority queue (priorityqueue.go), eliminating the custom bubble-up/down logic and the latent bug class it carried. - Drop the side handle map: membership is validated by index identity (a *heapItem lives in exactly one queue), preserving the contract's Remove error semantics with less state. O(log n) targeted removal is retained via per-item index tracking. - Rename the registered queue MaxMinHeap -> PriorityQueue. Signed-off-by: Luke Van Drie <lukevandrie@google.com>

RishabhSaini

The stdlib swap eliminates the bug class entirely and improves maintainability. Code is correct, all invariants preserved.
Nit: the PR description says "N/A" for the issue link but references #1689. Consider linking it as "Fixes #1689" so it auto-closes.
/lgtm

shmuelk · 2026-06-21T10:17:18Z

+	pq.mu.Lock()
+	defer pq.mu.Unlock()


The statements that need locking are 189 through 193, perhaps the Lock and Unlock calls should wrap those statements

Partially applied. I hoisted the nil check and the type assertion out (they only inspect the caller's handle). But the IsInvalidated() check has to stay inside the lock: isInvalidated (and index) are written under this same mutex by Remove/Cleanup/Drain, so reading them lock-free would be a data race.

So the locked region is now: IsInvalidated → identity check → heap.Remove → invalidate.

shmuelk · 2026-06-21T10:19:55Z

+			hi.index = -1
+			pq.byteSize.Add(^hi.item.OriginalRequest().ByteSize() + 1) // Atomic subtraction.
+		} else {
+			itemsToKeep = append(itemsToKeep, hi)


As itemsToKeep is pre-allocated to be the size of the current heap, which will be more than sufficient even if no items are removed by the predicate. Why not have a counter and do simple assignments instead of calling append here?

Nice call; I switched to in-place compaction which drops the second slice allocation entirely. Added a small tail-nil loop so the removed *heapItems aren't retained by the vacated capacity.

shmuelk · 2026-06-21T10:23:22Z

+	pq.mu.Lock()
+	defer pq.mu.Unlock()
+
+	drainedItems := make([]flowcontrol.QueueItemAccessor, len(pq.heap.items))
+	for i, hi := range pq.heap.items {
+		drainedItems[i] = hi.item
+		hi.Invalidate()
+		hi.index = -1
+	}
+
+	pq.heap.items = make([]*heapItem, 0)
+	pq.byteSize.Store(0)


Suggested change

pq.mu.Lock()

defer pq.mu.Unlock()

drainedItems := make([]flowcontrol.QueueItemAccessor, len(pq.heap.items))

for i, hi := range pq.heap.items {

drainedItems[i] = hi.item

hi.Invalidate()

hi.index = -1

}

pq.heap.items = make([]*heapItem, 0)

pq.byteSize.Store(0)

pq.mu.Lock()

old := pq.heap.items

pq.heap.items = make([]*heapItem, 0)

pq.byteSize.Store(0)

pq.mu.Unlock()

drainedItems := make([]flowcontrol.QueueItemAccessor, len(old))

for i, hi := range old {

drainedItems[i] = hi.item

hi.Invalidate()

hi.index = -1

}

I'd like to push back on this one. Moving the loop outside the lock makes hi.Invalidate() / hi.index = -1 race with a concurrent Remove(handle). Even after the slice is swapped out, a caller can still hold a handle to one of the drained items and call Remove, which reads hi.isInvalidated and hi.index under the lock. Drain mutating those same fields lock-free is a data race.

If we ever want to shorten this O(n) section, the clean way would be making isInvalidated/index atomic so invalidation can go lock-free, but that's a broader change I'd rather not fold into this PR.

Also, Drain is only ran on the graceful shutdown path. I don't think we need to be concerned with optimizing this method.

@shmuelk

Addresses review comments from @shmuelk on llm-d#1690: - Correct the new-file license headers to "2026 The llm-d Authors". - Clarify the PriorityQueue doc: ordering is maintained by an internal container/heap (the queue does not itself implement heap.Interface). - Narrow the Add and Remove critical sections to the shared heap mutation. Precondition checks (nil, type assertion) and the atomic byteSize update move outside the lock; the isInvalidated and index reads stay inside it, since those fields are written under the lock by Remove/Cleanup/Drain and reading them lock-free would race. - Compact Cleanup survivors in place instead of allocating a second slice, clearing the vacated tail so removed items are not retained. Signed-off-by: Luke Van Drie <lukevandrie@google.com>

Signed-off-by: llm-d-router-release-notes[bot] <287676111+llm-d-router-release-notes[bot]@users.noreply.github.com> Co-authored-by: llm-d-router-release-notes[bot] <287676111+llm-d-router-release-notes[bot]@users.noreply.github.com>

LukeAVanDrie requested review from a team and shmuelk as code owners June 18, 2026 21:14

LukeAVanDrie requested review from elevran and vMaroon June 18, 2026 21:14

github-actions Bot added size/XXL Denotes a PR that changes 1000+ lines, ignoring generated files. kind/cleanup and removed kind/cleanup labels Jun 18, 2026

LukeAVanDrie mentioned this pull request Jun 18, 2026

fix(flowcontrol): call up() after down() in maxMinHeap.Remove #1689

Closed

2 tasks

RishabhSaini approved these changes Jun 19, 2026

View reviewed changes