From 6a65d8406a0da7833c6eabbb2a4f2a149ebb5c2e Mon Sep 17 00:00:00 2001 From: Neil <4138956+nwparker@users.noreply.github.com> Date: Sun, 30 Aug 2026 23:05:37 -0700 Subject: [PATCH] perf(ssh): index source spans by ID (#17504) --- .../ipc/ssh-pty-source-ack-publication.ts | 8 +++- .../ssh-pty-source-obligation-ledger.test.ts | 45 +++++++++++++++++++ .../ipc/ssh-pty-source-obligation-ledger.ts | 18 ++++---- .../ipc/ssh-pty-source-obligation-state.ts | 31 +++++++++---- .../ssh-pty-source-obligation-transitions.ts | 9 ++-- 5 files changed, 88 insertions(+), 23 deletions(-) diff --git a/src/main/ipc/ssh-pty-source-ack-publication.ts b/src/main/ipc/ssh-pty-source-ack-publication.ts index db1d9581255..966e6d660f2 100644 --- a/src/main/ipc/ssh-pty-source-ack-publication.ts +++ b/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, + spanOwners: Map, onPublished: () => void ): SshPtySourceAckPublication { let settled = false diff --git a/src/main/ipc/ssh-pty-source-obligation-ledger.test.ts b/src/main/ipc/ssh-pty-source-obligation-ledger.test.ts index ec6ad29b41e..3f66c371e9a 100644 --- a/src/main/ipc/ssh-pty-source-obligation-ledger.test.ts +++ b/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 { + getCalls = 0 + + override get(key: string): unknown { + this.getCalls += 1 + return super.get(key) + } +} + function identity( deliveryToken = 'token-1', overrides: Partial = {} @@ -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 + tokens: Map + } + 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() diff --git a/src/main/ipc/ssh-pty-source-obligation-ledger.ts b/src/main/ipc/ssh-pty-source-obligation-ledger.ts index 976f7d2fc5c..1fbd64b9656 100644 --- a/src/main/ipc/ssh-pty-source-obligation-ledger.ts +++ b/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() private readonly closedSnapshots = new Map() private readonly reservations = new Map() - private readonly spanOwners = new Map() + private readonly spanOwners = new Map() 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 { diff --git a/src/main/ipc/ssh-pty-source-obligation-state.ts b/src/main/ipc/ssh-pty-source-obligation-state.ts index a51694518a4..fa407c9fc4b 100644 --- a/src/main/ipc/ssh-pty-source-obligation-state.ts +++ b/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 + 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 + spanOwners: Map ): 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 + spanOwners: Map ): 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 + spanOwners: Map ): boolean { const last = token.spans.at(-1) if ( @@ -210,15 +213,25 @@ export function rollbackCommittedSourceSpan( } export function requireSourceSpan( - spanOwners: ReadonlyMap, + spanOwners: ReadonlyMap, 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, + 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, closedSnapshots: Map, reservations: Map, - spanOwners: Map, + spanOwners: Map, onTokenClosed: (identity: PtySourceDeliveryIdentity) => void ): void { token.state = 'closed' diff --git a/src/main/ipc/ssh-pty-source-obligation-transitions.ts b/src/main/ipc/ssh-pty-source-obligation-transitions.ts index 692a891be48..d5ec17389bf 100644 --- a/src/main/ipc/ssh-pty-source-obligation-transitions.ts +++ b/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, + spanOwners: ReadonlyMap, spanId: string, consumer: SshPtySourceConsumerId, next: SshPtySourceObligationState @@ -40,7 +41,7 @@ export function transitionOpenSourceObligation( } export function commitSourceObligationTransfer( - spanOwners: ReadonlyMap, + spanOwners: ReadonlyMap, spanId: string, consumer: SshPtySourceConsumerId ): boolean { @@ -58,7 +59,7 @@ export function commitSourceObligationTransfer( } export function cancelSourceObligationTransfer( - spanOwners: ReadonlyMap, + spanOwners: ReadonlyMap, spanId: string, consumer: SshPtySourceConsumerId, reason: string @@ -73,7 +74,7 @@ export function cancelSourceObligationTransfer( } export function rollbackSourceObligationTransfer( - spanOwners: ReadonlyMap, + spanOwners: ReadonlyMap, spanId: string, consumer: SshPtySourceConsumerId ): boolean {