merge PR 20052 relay-abort-index

src/relay/client-request-aborts.ts

@@ -1,40 +1,69 @@
-export class ClientRequestAborts {
-  private readonly controllers = new Map<string, AbortController>()
+/** Opaque handle returned by `create`, so a release needs no string parsing to find its owner. */
+export type ClientRequestAbortHandle = {
+  readonly clientId: number
+  readonly requestId: number
+}
 
-  create(clientId: number, requestId: number): { key: string; controller: AbortController } {
-    const key = this.key(clientId, requestId)
+export class ClientRequestAborts {
+  // Why indexed by client instead of one flat map under composite `${clientId}:${requestId}` keys:
+  // abortClient runs on every closeClient and every setWrite, and against a flat map it had to scan
+  // every entry to find one client's. A scan with an early break cannot fix that -- the matching
+  // keys are scattered through the map, so any correct loop still visits every entry, which made a
+  // full churn of N clients cost K*N*(N+1)/2 visits. Only an index makes a teardown proportional to
+  // what that client actually owns.
+  private readonly byClient = new Map<number, Map<number, AbortController>>()
+
+  create(
+    clientId: number,
+    requestId: number
+  ): { key: ClientRequestAbortHandle; controller: AbortController } {
     const controller = new AbortController()
-    this.controllers.set(key, controller)
-    return { key, controller }
+    let requests = this.byClient.get(clientId)
+    if (!requests) {
+      requests = new Map<number, AbortController>()
+      this.byClient.set(clientId, requests)
+    }
+    requests.set(requestId, controller)
+    return { key: { clientId, requestId }, controller }
   }
 
   get(clientId: number, requestId: number): AbortController | undefined {
-    return this.controllers.get(this.key(clientId, requestId))
+    return this.byClient.get(clientId)?.get(requestId)
   }
 
-  delete(key: string): void {
-    this.controllers.delete(key)
+  delete(key: ClientRequestAbortHandle): void {
+    const requests = this.byClient.get(key.clientId)
+    if (!requests) {
+      return
+    }
+    requests.delete(key.requestId)
+    // Why drop the empty bucket: otherwise a churned client leaves an entry behind for the life of
+    // the relay, which is the retention the index exists to avoid.
+    if (requests.size === 0) {
+      this.byClient.delete(key.clientId)
+    }
   }
 
   abortClient(clientId: number): void {
-    const prefix = `${clientId}:`
-    for (const [key, controller] of this.controllers) {
-      if (!key.startsWith(prefix)) {
-        continue
-      }
+    const requests = this.byClient.get(clientId)
+    if (!requests) {
+      return
+    }
+    // Unlink before aborting: an abort listener that reaches back in must not see a half-emptied
+    // bucket, and the whole bucket is going regardless.
+    this.byClient.delete(clientId)
+    for (const controller of requests.values()) {
       controller.abort()
-      this.controllers.delete(key)
     }
   }
 
   abortAll(): void {
-    for (const [, controller] of this.controllers) {
-      controller.abort()
+    const buckets = Array.from(this.byClient.values())
+    this.byClient.clear()
+    for (const requests of buckets) {
+      for (const controller of requests.values()) {
+        controller.abort()
+      }
     }
-    this.controllers.clear()
-  }
-
-  private key(clientId: number, requestId: number): string {
-    return `${clientId}:${requestId}`
   }
 }

src/relay/dispatcher-capacity-signals.ts

@@ -55,7 +55,7 @@ export abstract class RelayDispatcherCapacitySignals extends RelayDispatcherClie
   }
 
   get legacyRetentionBelowLowWater(): boolean {
-    return this.publicationLedger.belowLowWater(this.activeClientKeys())
+    return this.publicationLedger.belowLowWater(() => this.activeClientKeys())
   }
 
   /**
@@ -138,7 +138,7 @@ export abstract class RelayDispatcherCapacitySignals extends RelayDispatcherClie
       this.deferredLegacyCapacity ||= !force
       return
     }
-    if (!force && !this.publicationLedger.belowLowWater(this.activeClientKeys())) {
+    if (!force && !this.publicationLedger.belowLowWater(() => this.activeClientKeys())) {
       return
     }
     for (const listener of this.legacyCapacityListeners) {

src/relay/dispatcher-per-connection-state-baseline.test.ts

@@ -0,0 +1,117 @@
+import { describe, expect, it, vi, afterEach } from 'vitest'
+import { RelayDispatcher } from './dispatcher'
+
+// Why a census rather than a duration: "the dispatcher releases per-connection state" is a
+// statement about what is still *held* after churn, so it is measured by counting containers,
+// not by timing a teardown. Every number here is exact and load-independent.
+
+type Probed = {
+  attachClient: (w: (b: Buffer) => void) => number
+  detachClient: (id: number) => void
+  onClientCapacity: (id: number, listener: () => void) => (() => void) | null
+  clients: Map<number, unknown>
+  requestHandlers: Map<string, unknown>
+  notificationHandlers: Map<string, unknown>
+  requestAborts: {
+    byClient: Map<number, Map<number, AbortController>>
+    create: (clientId: number, requestId: number) => unknown
+  }
+  publicationLedger: { clientBytes: Map<string, number>; aggregateBytes: number }
+  pendingRelayRequests: Map<number, unknown>
+  clientDetachListeners: Set<unknown>
+  disposeListeners: Set<unknown>
+  legacyCapacityListeners: Set<unknown>
+  clientCapacityListeners: Map<number, unknown>
+  ptyDataPublicationAdmission: unknown
+  keepaliveTimer: unknown
+  activeClients: () => unknown[]
+  tryPublishToClients: (clients: unknown[], msg: unknown, lane: string) => boolean
+  dispose: () => void
+}
+
+function countAbortControllers(d: Probed): number {
+  let total = 0
+  for (const bucket of d.requestAborts.byClient.values()) {
+    total += bucket.size
+  }
+  return total
+}
+
+function census(d: Probed): Record<string, number | string> {
+  return {
+    clients: d.clients.size,
+    requestHandlers: d.requestHandlers.size,
+    notificationHandlers: d.notificationHandlers.size,
+    requestAbortControllers: countAbortControllers(d),
+    ledgerClientBytes: d.publicationLedger.clientBytes.size,
+    ledgerAggregateBytes: d.publicationLedger.aggregateBytes,
+    pendingRelayRequests: d.pendingRelayRequests.size,
+    clientDetachListeners: d.clientDetachListeners.size,
+    disposeListeners: d.disposeListeners.size,
+    legacyCapacityListeners: d.legacyCapacityListeners.size,
+    clientCapacityListeners: d.clientCapacityListeners.size,
+    ptyDataPublicationAdmission: d.ptyDataPublicationAdmission === null ? 'null' : 'set',
+    keepaliveTimer: d.keepaliveTimer === null ? 'null' : 'armed'
+  }
+}
+
+function newDispatcher(): Probed {
+  return new RelayDispatcher(() => {}) as unknown as Probed
+}
+
+const CLIENTS_PER_CYCLE = 100
+
+describe('relay dispatcher per-connection state', () => {
+  afterEach(() => vi.useRealTimers())
+
+  it('returns every per-connection container to baseline across repeated churn', () => {
+    vi.useFakeTimers()
+    const d = newDispatcher()
+    const baseline = census(d)
+
+    for (let cycle = 0; cycle < 3; cycle++) {
+      const ids: number[] = []
+      for (let i = 0; i < CLIENTS_PER_CYCLE; i++) {
+        ids.push(d.attachClient(() => {}))
+      }
+      for (const id of ids) {
+        d.onClientCapacity(id, () => {})
+        d.requestAborts.create(id, 1)
+      }
+
+      // The census must be able to find things: these two are the containers that stay 0 unless
+      // deliberately loaded, so assert they actually moved before trusting that they came back.
+      expect(census(d).clients).toBe(CLIENTS_PER_CYCLE + 1)
+      expect(census(d).clientCapacityListeners).toBe(CLIENTS_PER_CYCLE)
+      expect(census(d).requestAbortControllers).toBe(CLIENTS_PER_CYCLE)
+
+      d.tryPublishToClients(
+        d.activeClients(),
+        { jsonrpc: '2.0', method: 'pty.data', params: { d: 'x'.repeat(256) } },
+        'bulk'
+      )
+      for (const id of ids) {
+        d.detachClient(id)
+      }
+      expect(census(d)).toEqual(baseline)
+    }
+    d.dispose()
+  })
+
+  // Why this is separate: the ledger is the one container with no per-client teardown. Every other
+  // container is reclaimed by closeClient; a ledger entry is reclaimed only by its own lease's
+  // release(). Normal closes settle every queued and in-flight entry, so this never fires in
+  // practice -- but nothing in the close path would reclaim an entry that did survive.
+  it('does not reclaim a publication-ledger entry on client close', () => {
+    vi.useFakeTimers()
+    const d = newDispatcher()
+    const id = d.attachClient(() => {})
+    d.publicationLedger.clientBytes.set('stranded-key', 4096)
+
+    d.detachClient(id)
+
+    expect(d.clients.size).toBe(1)
+    expect(d.publicationLedger.clientBytes.has('stranded-key')).toBe(true)
+    d.dispose()
+  })
+})

src/relay/legacy-relay-publication-ledger.ts

@@ -83,11 +83,18 @@ export class LegacyRelayPublicationLedger {
     })
   }
 
-  belowLowWater(clientKeys?: readonly string[]): boolean {
+  // Why the thunk overload: the aggregate ceiling below decides on its own most of the time, and it
+  // decides *first*. A caller passing an eager array has already built one string per client before
+  // learning the keys were never going to be read -- and it pays that most in the loaded case,
+  // because that is exactly when the aggregate check short-circuits.
+  belowLowWater(clientKeys?: readonly string[] | (() => readonly string[])): boolean {
     if (this.aggregateBytes > this.relayLowBytes) {
       return false
     }
-    const keys = clientKeys ?? Array.from(this.clientBytes.keys())
+    const keys =
+      typeof clientKeys === 'function'
+        ? clientKeys()
+        : (clientKeys ?? Array.from(this.clientBytes.keys()))
     return keys.every((clientKey) => (this.clientBytes.get(clientKey) ?? 0) <= this.clientLowBytes)
   }
 

src/relay/relay-hot-path-operation-counts.test.ts

@@ -0,0 +1,157 @@
+import { describe, expect, it, vi, afterEach } from 'vitest'
+import { RelayDispatcher } from './dispatcher'
+import { ClientRequestAborts } from './client-request-aborts'
+
+// Why operation counts and not milliseconds: each assertion below is about how many entries a hot
+// path visits, which is the property. A duration is only a proxy for it, and a proxy needs a
+// threshold calibrated against observed runtimes -- which makes the test about the observation.
+// These counts are exact and identical under any machine load.
+
+/** Counts entries yielded by a real Map's iterators without changing the code under test. */
+class CountingMap<K, V> extends Map<K, V> {
+  visits = 0
+  getCalls = 0
+
+  private countingIterator<T>(inner: IterableIterator<T>): IterableIterator<T> {
+    const bump = (): void => {
+      this.visits++
+    }
+    return {
+      next(): IteratorResult<T> {
+        const r = inner.next()
+        if (!r.done) {
+          bump()
+        }
+        return r
+      },
+      [Symbol.iterator]() {
+        return this
+      }
+    } as IterableIterator<T>
+  }
+
+  override [Symbol.iterator](): MapIterator<[K, V]> {
+    return this.countingIterator(super[Symbol.iterator]()) as MapIterator<[K, V]>
+  }
+
+  override values(): MapIterator<V> {
+    return this.countingIterator(super.values()) as MapIterator<V>
+  }
+
+  override get(key: K): V | undefined {
+    this.getCalls++
+    return super.get(key)
+  }
+}
+
+type ProbedDispatcher = {
+  attachClient: (w: (b: Buffer) => void) => number
+  clients: Map<number, unknown>
+  publicationLedger: { clientBytes: Map<string, number> }
+  notifyLegacyCapacityIfLow: () => void
+  activeClients: () => unknown[]
+  tryPublishToClients: (clients: unknown[], msg: unknown, lane: string) => boolean
+  dispose: () => void
+}
+
+function dispatcherWithClients(clientCount: number): {
+  d: ProbedDispatcher
+  clients: CountingMap<number, unknown>
+  ledger: CountingMap<string, number>
+} {
+  const d = new RelayDispatcher(() => {}) as unknown as ProbedDispatcher
+  for (let i = 1; i < clientCount; i++) {
+    d.attachClient(() => {})
+  }
+  const clients = new CountingMap<number, unknown>()
+  for (const [k, v] of d.clients) {
+    clients.set(k, v)
+  }
+  d.clients = clients
+  const ledger = new CountingMap<string, number>()
+  d.publicationLedger.clientBytes = ledger
+  clients.visits = 0
+  ledger.getCalls = 0
+  return { d, clients, ledger }
+}
+
+describe('relay hot-path operation counts', () => {
+  afterEach(() => vi.useRealTimers())
+
+  // Why this is the guard and not a duration: abortClient runs on every closeClient and every
+  // setWrite. Under the flat composite-key map it replaced, one client's teardown enumerated every
+  // controller in the relay, so a full churn of N clients holding K requests cost K*N*(N+1)/2 visits
+  // -- measured at 50 -> 5,100, 100 -> 20,200, 200 -> 80,400, 400 -> 320,800, exactly 4x per
+  // doubling. Teardown must now visit only what the client owns, and must not enumerate the client
+  // index at all: enumerating it *is* the old scan.
+  it('abortClient visits only the target client, and never enumerates the client index', () => {
+    const clientCount = 40
+    const inFlightPerClient = 4
+    const aborts = new ClientRequestAborts()
+    for (let c = 1; c <= clientCount; c++) {
+      for (let r = 1; r <= inFlightPerClient; r++) {
+        aborts.create(c, r)
+      }
+    }
+    const byClient = (aborts as unknown as { byClient: Map<number, Map<number, AbortController>> })
+      .byClient
+
+    // The census must be able to find things: prove the maps really hold 160 controllers across 40
+    // buckets before asserting that a teardown only touches 4 of them.
+    expect(byClient.size).toBe(clientCount)
+    let totalControllers = 0
+    for (const bucket of byClient.values()) {
+      totalControllers += bucket.size
+    }
+    expect(totalControllers).toBe(clientCount * inFlightPerClient)
+
+    const index = new CountingMap<number, Map<number, AbortController>>()
+    for (const [k, v] of byClient) {
+      index.set(k, v)
+    }
+    const targetBucket = new CountingMap<number, AbortController>()
+    for (const [k, v] of byClient.get(1)!) {
+      targetBucket.set(k, v)
+    }
+    index.set(1, targetBucket)
+    ;(aborts as unknown as { byClient: Map<number, unknown> }).byClient = index
+    index.visits = 0
+    targetBucket.visits = 0
+
+    aborts.abortClient(1)
+
+    expect(targetBucket.visits).toBe(inFlightPerClient)
+    expect(index.visits).toBe(0)
+    expect(index.has(1)).toBe(false)
+  })
+
+  it('notifyLegacyCapacity costs exactly one ledger lookup per active client', () => {
+    vi.useFakeTimers()
+    for (const clientCount of [50, 100, 200, 400]) {
+      const { d, clients, ledger } = dispatcherWithClients(clientCount)
+
+      d.notifyLegacyCapacityIfLow()
+
+      expect(clients.visits, `clients enumerated at n=${clientCount}`).toBe(clientCount)
+      expect(ledger.getCalls, `ledger lookups at n=${clientCount}`).toBe(clientCount)
+      d.dispose()
+    }
+  })
+
+  it('one broadcast publication costs a fixed number of lookups per subscriber', () => {
+    vi.useFakeTimers()
+    for (const clientCount of [10, 20, 40]) {
+      const { d, clients, ledger } = dispatcherWithClients(clientCount)
+
+      d.tryPublishToClients(
+        d.activeClients(),
+        { jsonrpc: '2.0', method: 'pty.data', params: { d: 'x' } },
+        'bulk'
+      )
+
+      expect(clients.visits, `clients enumerated at n=${clientCount}`).toBe(clientCount * 2)
+      expect(ledger.getCalls, `ledger lookups at n=${clientCount}`).toBe(clientCount * 4)
+      d.dispose()
+    }
+  })
+})