From 96216601bc977bb9011dd7600cc7ac80bb3d736b Mon Sep 17 00:00:00 2001
From: OrcaWin <alpha-eng@stably.ai>
Date: Tue, 8 Sep 2026 19:40:25 -0700
Subject: [PATCH] perf: retire exhausted buckets from host-balanced listings
 (#19493)

* perf: retire exhausted buckets from host-balanced listings

* test(listing): guard the bucket-retirement invariant and starvation fix

Document why in-place compaction keeps bucket[round] defined and preserves
first-appearance host order, and cover cap-exact filling plus small-host
representation so a drifting retirement predicate cannot emit blank rows.

---------

Co-authored-by: m4air <m4air@m4airs-MacBook-Air.local>
Co-authored-by: Neil <4138956+nwparker@users.noreply.github.com>
---
 src/shared/host-balanced-listing-page.ts      | 33 +++++----
 .../host-balanced-listing-scaling.test.ts     | 67 +++++++++++++++++++
 2 files changed, 87 insertions(+), 13 deletions(-)
 create mode 100644 src/shared/host-balanced-listing-scaling.test.ts

diff --git a/src/shared/host-balanced-listing-page.ts b/src/shared/host-balanced-listing-page.ts
index f771d98fdd4..baff6e18a37 100644
--- a/src/shared/host-balanced-listing-page.ts
+++ b/src/shared/host-balanced-listing-page.ts
@@ -28,22 +28,29 @@ export function selectHostBalancedPage<TRow>(
     }
   })
   const buckets = [...indicesByHost.values()]
-  const cursors = buckets.map(() => 0)
   const chosen: number[] = []
-  while (chosen.length < limit) {
-    let advanced = false
-    for (let bucket = 0; bucket < buckets.length && chosen.length < limit; bucket += 1) {
-      const cursor = cursors[bucket] ?? 0
-      const index = buckets[bucket]?.[cursor]
-      if (index !== undefined) {
-        chosen.push(index)
-        cursors[bucket] = cursor + 1
-        advanced = true
+  // Compacting in place keeps `buckets[0..activeCount)` as exactly the buckets longer than
+  // `round`, in their original order, so `bucket[round]` is always defined and the round robin
+  // still visits hosts in first-appearance order. Retiring them keeps a long host bucket from
+  // re-scanning every exhausted one. Writes land at or before the slot just read, never ahead.
+  let activeCount = buckets.length
+  let round = 0
+  while (chosen.length < limit && activeCount > 0) {
+    let nextActiveCount = 0
+    for (
+      let bucketIndex = 0;
+      bucketIndex < activeCount && chosen.length < limit;
+      bucketIndex += 1
+    ) {
+      const bucket = buckets[bucketIndex]
+      chosen.push(bucket[round])
+      if (bucket.length > round + 1) {
+        buckets[nextActiveCount] = bucket
+        nextActiveCount += 1
       }
     }
-    if (!advanced) {
-      break
-    }
+    activeCount = nextActiveCount
+    round += 1
   }
   return chosen.sort((left, right) => left - right).map((index) => rows[index] as TRow)
 }
diff --git a/src/shared/host-balanced-listing-scaling.test.ts b/src/shared/host-balanced-listing-scaling.test.ts
new file mode 100644
index 00000000000..aa7e40c2eb6
--- /dev/null
+++ b/src/shared/host-balanced-listing-scaling.test.ts
@@ -0,0 +1,67 @@
+import { expect, it, vi } from 'vitest'
+import { selectHostBalancedPage } from './host-balanced-listing-page'
+
+it('retires exhausted host buckets from subsequent listing rounds', () => {
+  const rows = [
+    ...Array.from({ length: 1000 }, (_, index) => ({ host: `host-${index}`, id: index })),
+    ...Array.from({ length: 2000 }, (_, index) => ({ host: 'large-host', id: 1000 + index }))
+  ]
+  const original = Map.prototype.values
+  let reads = 0
+  const spy = vi
+    .spyOn(Map.prototype, 'values')
+    .mockImplementation(function (this: Map<string, number[]>) {
+      const iterator = original.call(this)
+      if (!this.has('large-host')) {
+        return iterator
+      }
+      return iterator.map(
+        (bucket: number[]) =>
+          new Proxy(bucket, {
+            get(target, key, receiver) {
+              if (typeof key === 'string' && /^\d+$/.test(key)) {
+                reads += 1
+              }
+              return Reflect.get(target, key, receiver)
+            }
+          })
+      )
+    })
+  let result: typeof rows
+  try {
+    result = selectHostBalancedPage(rows, 2000, (row) => row.host)
+  } finally {
+    spy.mockRestore()
+  }
+  expect(result!.map((row) => row.id)).toEqual(Array.from({ length: 2000 }, (_, index) => index))
+  expect(reads).toBeLessThanOrEqual(2000)
+})
+
+// Retirement is only safe while every retained bucket outlives the round it is read at. If that
+// predicate drifts, `bucket[round]` yields undefined and the page grows blank rows.
+it('fills every cap exactly, without undefined or duplicate rows', () => {
+  const rows = [
+    ...Array.from({ length: 3 }, (_, index) => ({ host: 'a', id: index })),
+    { host: 'b', id: 3 },
+    ...Array.from({ length: 9 }, (_, index) => ({ host: 'c', id: 4 + index }))
+  ]
+  for (let limit = 0; limit < rows.length; limit += 1) {
+    const page = selectHostBalancedPage(rows, limit, (row) => row.host)
+    const pageIds = page.map((row) => row.id)
+    expect(page).toHaveLength(limit)
+    expect(page.every((row) => row !== undefined)).toBe(true)
+    expect(new Set(pageIds).size).toBe(limit)
+    // the page must stay a subsequence of the caller's original order
+    expect(pageIds).toEqual([...pageIds].sort((left, right) => left - right))
+  }
+})
+
+it('keeps starved hosts represented once a dominant bucket is retired', () => {
+  const rows = [
+    ...Array.from({ length: 2000 }, (_, index) => ({ host: 'big', id: index })),
+    ...Array.from({ length: 40 }, (_, index) => ({ host: `small-${index}`, id: 2000 + index }))
+  ]
+  const page = selectHostBalancedPage(rows, 100, (row) => row.host)
+  expect(new Set(page.map((row) => row.host)).size).toBe(41)
+  expect(page.filter((row) => row.host === 'big')).toHaveLength(60)
+})