perf(ssh): index source spans by ID (#17504)

src/main/ipc/ssh-pty-source-ack-publication.ts

@@ -1,10 +1,14 @@
 import type { SshPtySourceAckPublication } from './ssh-pty-source-obligation-contract'
-import { reclaimPublishedSourcePrefix, type TokenRecord } from './ssh-pty-source-obligation-state'
+import {
+  reclaimPublishedSourcePrefix,
+  type SpanRecord,
+  type TokenRecord
+} from './ssh-pty-source-obligation-state'
 
 export function createSshPtySourceAckPublication(
   token: TokenRecord,
   endSu: number,
-  spanOwners: Map<string, TokenRecord>,
+  spanOwners: Map<string, SpanRecord>,
   onPublished: () => void
 ): SshPtySourceAckPublication {
   let settled = false

src/main/ipc/ssh-pty-source-obligation-ledger.test.ts

@@ -5,6 +5,15 @@ import type {
 } from '../../shared/pty-source-credit-contract'
 import { SshPtySourceObligationLedger } from './ssh-pty-source-obligation-ledger'
 
+class CountingSpanOwnerMap extends Map<string, unknown> {
+  getCalls = 0
+
+  override get(key: string): unknown {
+    this.getCalls += 1
+    return super.get(key)
+  }
+}
+
 function identity(
   deliveryToken = 'token-1',
   overrides: Partial<PtySourceDeliveryIdentity> = {}
@@ -55,6 +64,42 @@ function commitSpan(
 }
 
 describe('SshPtySourceObligationLedger', () => {
+  it('looks up retained spans directly by ID as the ledger grows', () => {
+    const spanCount = 1_024
+    const ledger = new SshPtySourceObligationLedger()
+    const owner = identity()
+    ledger.open(owner)
+    for (let index = 0; index < spanCount; index += 1) {
+      commitSpan(ledger, owner, span(owner, `span-${index}`, index, 'x'))
+    }
+
+    const internals = ledger as unknown as {
+      spanOwners: Map<string, unknown>
+      tokens: Map<string, { spans: readonly { span: PtySourceSpan }[] }>
+    }
+    const countedSpanOwners = new CountingSpanOwnerMap(internals.spanOwners)
+    internals.spanOwners = countedSpanOwners
+    const tokenRecord = Array.from(internals.tokens.values())[0]
+    if (!tokenRecord) {
+      throw new Error('test token record missing')
+    }
+
+    let legacyVisits = 0
+    for (let index = 0; index < spanCount; index += 1) {
+      for (const candidate of tokenRecord.spans) {
+        legacyVisits += 1
+        if (candidate.span.spanId === `span-${index}`) {
+          break
+        }
+      }
+      expect(ledger.obligation(`span-${index}`, 'model').state).toBe('open')
+    }
+
+    expect(legacyVisits).toBe(524_800)
+    expect(countedSpanOwners.getCalls).toBe(spanCount)
+    expect(legacyVisits - countedSpanOwners.getCalls).toBe(523_776)
+  })
+
   it('rolls back an uncommitted admission without consuming its source coordinate', () => {
     const ledger = new SshPtySourceObligationLedger()
     const owner = identity()

src/main/ipc/ssh-pty-source-obligation-ledger.ts

@@ -27,10 +27,12 @@ import {
   createSourceToken,
   markSourceExitPublished,
   requireSourceSpan,
+  requireSourceReservation,
   rollbackCommittedSourceSpan,
   sealSourceToken,
   snapshotSourceToken,
   type ReservationRecord,
+  type SpanRecord,
   type TokenRecord
 } from './ssh-pty-source-obligation-state'
 import {
@@ -54,7 +56,7 @@ export class SshPtySourceObligationLedger {
   private readonly tokens = new Map<string, TokenRecord>()
   private readonly closedSnapshots = new Map<string, SshPtySourceTokenSnapshot>()
   private readonly reservations = new Map<string, ReservationRecord>()
-  private readonly spanOwners = new Map<string, TokenRecord>()
+  private readonly spanOwners = new Map<string, SpanRecord>()
   private nextReservationId = 1
 
   constructor(
@@ -113,10 +115,14 @@ export class SshPtySourceObligationLedger {
     if (token.state !== 'active' || token.receivedEndSu !== reservation.span.sourceStartSu) {
       throw new Error('SSH PTY source admission reservation became stale')
     }
-    const spanRecord = createSourceSpanRecord(reservation.span, reservation.requiredConsumers)
+    const spanRecord = createSourceSpanRecord(
+      token,
+      reservation.span,
+      reservation.requiredConsumers
+    )
     token.spans.push(spanRecord)
     token.receivedEndSu = reservation.span.sourceEndSu
-    this.spanOwners.set(reservation.span.spanId, token)
+    this.spanOwners.set(reservation.span.spanId, spanRecord)
     this.reservations.delete(reservation.reservationId)
   }
 
@@ -312,11 +318,7 @@ export class SshPtySourceObligationLedger {
   }
 
   private requireReservation(reservation: SshPtySourceAdmissionReservation): ReservationRecord {
-    const record = this.reservations.get(reservation.reservationId)
-    if (!record || record.reservation !== reservation) {
-      throw new Error('Unknown SSH PTY source admission reservation')
-    }
-    return record
+    return requireSourceReservation(this.reservations, reservation)
   }
 
   private requireToken(identity: PtySourceDeliveryIdentity): TokenRecord {

src/main/ipc/ssh-pty-source-obligation-state.ts

@@ -18,6 +18,7 @@ export const CLOSED_SOURCE_TOKEN_TOMBSTONE_LIMIT = 256
 export type SpanRecord = {
   span: PtySourceSpan
   obligations: Map<SshPtySourceConsumerId, SshPtySourceObligationState>
+  owner: TokenRecord
 }
 
 export type ReservationRecord = {
@@ -57,10 +58,12 @@ export function createSourceToken(
 }
 
 export function createSourceSpanRecord(
+  owner: TokenRecord,
   span: PtySourceSpan,
   consumers: readonly SshPtySourceConsumerId[]
 ): SpanRecord {
   return {
+    owner,
     span,
     obligations: new Map(
       consumers.map((consumer) => [consumer, Object.freeze({ state: 'open' as const })])
@@ -158,7 +161,7 @@ export function cancelOpenSourceObligations(token: TokenRecord, reason: string):
 
 export function reclaimPublishedSourcePrefix(
   token: TokenRecord,
-  spanOwners: Map<string, TokenRecord>
+  spanOwners: Map<string, SpanRecord>
 ): void {
   while (token.spans[0]?.span.sourceEndSu <= token.ackPublishedEndSu) {
     const record = token.spans.shift()!
@@ -168,7 +171,7 @@ export function reclaimPublishedSourcePrefix(
 
 export function releaseSourceTokenSpans(
   token: TokenRecord,
-  spanOwners: Map<string, TokenRecord>
+  spanOwners: Map<string, SpanRecord>
 ): void {
   for (const record of token.spans) {
     spanOwners.delete(record.span.spanId)
@@ -191,7 +194,7 @@ export function releaseSourceTokenReservations(
 export function rollbackCommittedSourceSpan(
   token: TokenRecord,
   reservation: SshPtySourceAdmissionReservation,
-  spanOwners: Map<string, TokenRecord>
+  spanOwners: Map<string, SpanRecord>
 ): boolean {
   const last = token.spans.at(-1)
   if (
@@ -210,15 +213,25 @@ export function rollbackCommittedSourceSpan(
 }
 
 export function requireSourceSpan(
-  spanOwners: ReadonlyMap<string, TokenRecord>,
+  spanOwners: ReadonlyMap<string, SpanRecord>,
   spanId: string
 ): { token: TokenRecord; span: SpanRecord } {
-  const token = spanOwners.get(spanId)
-  const span = token?.spans.find((candidate) => candidate.span.spanId === spanId)
-  if (!token || !span) {
+  const span = spanOwners.get(spanId)
+  if (!span || span.span.spanId !== spanId) {
     throw new Error('Unknown or reclaimed SSH PTY source span')
   }
-  return { token, span }
+  return { token: span.owner, span }
+}
+
+export function requireSourceReservation(
+  reservations: ReadonlyMap<string, ReservationRecord>,
+  reservation: SshPtySourceAdmissionReservation
+): ReservationRecord {
+  const record = reservations.get(reservation.reservationId)
+  if (!record || record.reservation !== reservation) {
+    throw new Error('Unknown SSH PTY source admission reservation')
+  }
+  return record
 }
 
 export function closeSourceGeneration(
@@ -271,7 +284,7 @@ export function closeSourceToken(
   tokens: Map<string, TokenRecord>,
   closedSnapshots: Map<string, SshPtySourceTokenSnapshot>,
   reservations: Map<string, ReservationRecord>,
-  spanOwners: Map<string, TokenRecord>,
+  spanOwners: Map<string, SpanRecord>,
   onTokenClosed: (identity: PtySourceDeliveryIdentity) => void
 ): void {
   token.state = 'closed'

src/main/ipc/ssh-pty-source-obligation-transitions.ts

@@ -6,6 +6,7 @@ import {
   advanceSourceTerminalEnd,
   cancelOpenSourceObligations,
   requireSourceSpan,
+  type SpanRecord,
   type TokenRecord
 } from './ssh-pty-source-obligation-state'
 
@@ -25,7 +26,7 @@ export function applySourceRecoveryCancellationProof(
 }
 
 export function transitionOpenSourceObligation(
-  spanOwners: ReadonlyMap<string, TokenRecord>,
+  spanOwners: ReadonlyMap<string, SpanRecord>,
   spanId: string,
   consumer: SshPtySourceConsumerId,
   next: SshPtySourceObligationState
@@ -40,7 +41,7 @@ export function transitionOpenSourceObligation(
 }
 
 export function commitSourceObligationTransfer(
-  spanOwners: ReadonlyMap<string, TokenRecord>,
+  spanOwners: ReadonlyMap<string, SpanRecord>,
   spanId: string,
   consumer: SshPtySourceConsumerId
 ): boolean {
@@ -58,7 +59,7 @@ export function commitSourceObligationTransfer(
 }
 
 export function cancelSourceObligationTransfer(
-  spanOwners: ReadonlyMap<string, TokenRecord>,
+  spanOwners: ReadonlyMap<string, SpanRecord>,
   spanId: string,
   consumer: SshPtySourceConsumerId,
   reason: string
@@ -73,7 +74,7 @@ export function cancelSourceObligationTransfer(
 }
 
 export function rollbackSourceObligationTransfer(
-  spanOwners: ReadonlyMap<string, TokenRecord>,
+  spanOwners: ReadonlyMap<string, SpanRecord>,
   spanId: string,
   consumer: SshPtySourceConsumerId
 ): boolean {