From ae2eeff55d40f36eca7df4a6e0cb531f3d11cf00 Mon Sep 17 00:00:00 2001 From: Neil <4138956+nwparker@users.noreply.github.com> Date: Mon, 31 Aug 2026 02:41:15 -0700 Subject: [PATCH] perf(relay): index PTY source-credit send spans (#17490) * perf(relay): index PTY source-credit send spans * perf(relay): maintain PTY source-credit retention totals * test(relay): pin PTY send-cursor rebase across ACK reclaim Cover the Math.max clamp branch in reclaimCreditedSpans where reclaim removes spans at or past the send cursor, and widen the seeded fuzz case to 20 spans per seed so the cursor actually traverses spans; assert the cursor never overshoots the span containing sentEndSu. * refactor(relay): drop dead retained-total helpers and pin retention counters The incremental PtySourceCreditRetention counters replaced the recompute-from-records helpers; delete the now-unreferenced exports and recompute the totals from the live records inside the ledger tests so the counters have an independent oracle. * test(relay): bound send-span reads instead of pinning the read pattern Address review feedback on the send-span cursor coverage: - replace the exact indexed-read pin and the tautological naive-visit assertion with a linear bound that still fails on the old Array.find path - drop the per-run bench console.log - assert retention totals immediately after rotate(), the only path that removes and re-adds a record in one call Also count the replacement delivery in retention as it enters the delivery map so the "in deliveries <=> counted" invariant never has a hole. --- src/relay/pty-source-credit-ledger.test.ts | 205 ++++++++++++++++++++- src/relay/pty-source-credit-ledger.ts | 47 ++--- src/relay/pty-source-credit-record.ts | 37 ++-- src/relay/pty-source-credit-retention.ts | 69 +++++++ src/relay/pty-source-credit-settlement.ts | 5 + 5 files changed, 313 insertions(+), 50 deletions(-) create mode 100644 src/relay/pty-source-credit-retention.ts diff --git a/src/relay/pty-source-credit-ledger.test.ts b/src/relay/pty-source-credit-ledger.test.ts index 37ad35a8f1e..7f3759e9591 100644 --- a/src/relay/pty-source-credit-ledger.test.ts +++ b/src/relay/pty-source-credit-ledger.test.ts @@ -1,7 +1,11 @@ import { describe, expect, it } from 'vitest' -import type { PtySourceDeliveryIdentity } from '../shared/pty-source-credit-contract' +import { + ptySourceDeliveryKey, + type PtySourceDeliveryIdentity, + type PtySourceSpan +} from '../shared/pty-source-credit-contract' import { RelayPtySourceCreditLedger } from './pty-source-credit-ledger' -import { CLOSED_DELIVERY_TOMBSTONE_LIMIT } from './pty-source-credit-record' +import { CLOSED_DELIVERY_TOMBSTONE_LIMIT, type DeliveryRecord } from './pty-source-credit-record' function identity( deliveryToken = 'token-1', @@ -51,19 +55,213 @@ function drainOne( return reservation } +type CursorRecord = { + spans: PtySourceSpan[] + sendSpanIndex: number +} + +function getCursorRecord( + ledger: RelayPtySourceCreditLedger, + owner: PtySourceDeliveryIdentity +): CursorRecord { + const internals = ledger as unknown as { deliveries: Map } + const record = internals.deliveries.get(ptySourceDeliveryKey(owner)) + if (!record) { + throw new Error('test delivery record missing') + } + return record +} + +// Why: retention totals are maintained incrementally, so recomputing from the live records is +// the only independent check that no mutation path skipped a counter update. +function expectRetentionMatchesRecords(ledger: RelayPtySourceCreditLedger): void { + const internals = ledger as unknown as { deliveries: Map } + const expected = { sourceSu: 0, dataBytes: 0, spans: 0 } + for (const record of internals.deliveries.values()) { + expected.sourceSu += record.receivedEndSu - record.creditedEndSu + expected.dataBytes += record.retainedDataBytes + expected.spans += record.spans.length + } + expect(ledger.retentionSnapshot()).toEqual(expected) +} + describe('RelayPtySourceCreditLedger', () => { + it('keeps send-span lookup near-linear across a retained-frame burst', () => { + const spanCount = 1_024 + const ledger = new RelayPtySourceCreditLedger({ + maxRetainedSourceSu: spanCount * 2, + maxAggregateRetainedSourceSu: spanCount * 2, + maxRetainedDataBytes: spanCount * 1_024, + maxAggregateRetainedDataBytes: spanCount * 1_024, + maxRetainedSpans: spanCount, + maxAggregateRetainedSpans: spanCount + }) + const owner = identity() + ledger.open(owner, spanCount * 2) + for (let index = 0; index < spanCount; index += 1) { + append(ledger, owner, 'x', `span-${index}`) + } + + const record = getCursorRecord(ledger, owner) + let indexedReads = 0 + const spans = record.spans + record.spans = new Proxy(spans, { + get(target, property, receiver) { + if (typeof property === 'string' && /^\d+$/.test(property)) { + indexedReads += 1 + } + return Reflect.get(target, property, receiver) + } + }) + + let sends = 0 + while (true) { + const reservation = ledger.reserveNextSend(owner, 1) + if (!reservation) { + break + } + ledger.commitSend(reservation) + sends += 1 + } + + expect(sends).toBe(spanCount) + // The final span remains the cursor until another source append arrives. + expect(record.sendSpanIndex).toBe(spanCount - 1) + // Linear in spans; the removed Array.find rescanned the sent prefix (524,800 predicate visits). + expect(indexedReads).toBeLessThan(spanCount * 3) + }) + + it('keeps the uncovered-cursor diagnostic when a retained span has a gap', () => { + const ledger = new RelayPtySourceCreditLedger() + const owner = identity('token-gap') + ledger.open(owner, 8) + append(ledger, owner, 'ab', 'span-gap') + const record = getCursorRecord(ledger, owner) + record.spans = [ + Object.freeze({ + ...record.spans[0], + sourceStartSu: 2, + sourceEndSu: 4, + transform: Object.freeze({ ...record.spans[0].transform, rawLengthSu: 2 }) + }) + ] + + expect(() => ledger.reserveNextSend(owner, 2)).toThrow( + 'PTY source delivery cursor is not covered by the retained ledger' + ) + }) + + it('keeps the cursor correct across ACK reclaim, rollback, rotation, and close', () => { + const ledger = new RelayPtySourceCreditLedger() + const oldOwner = identity() + const replacement = identity('token-replacement', { + clientGeneration: 4, + ownerGeneration: 5 + }) + ledger.open(oldOwner, 16) + append(ledger, oldOwner, 'ab', 'span-a') + append(ledger, oldOwner, 'cd', 'span-b') + append(ledger, oldOwner, 'ef', 'span-c') + + const first = ledger.reserveNextSend(oldOwner, 2)! + ledger.commitSend(first) + expect(getCursorRecord(ledger, oldOwner).sendSpanIndex).toBe(0) + ledger.acknowledge(oldOwner, { + id: oldOwner.id, + clientGeneration: oldOwner.clientGeneration, + ownerGeneration: oldOwner.ownerGeneration, + deliveryToken: oldOwner.deliveryToken, + creditedEndSu: 2 + }) + expect(getCursorRecord(ledger, oldOwner).sendSpanIndex).toBe(0) + + const attempted = ledger.reserveNextSend(oldOwner, 1)! + expect(attempted.span.data).toBe('c') + ledger.rollbackSend(attempted) + const retried = ledger.reserveNextSend(oldOwner, 1)! + expect(retried.span.data).toBe('c') + ledger.commitSend(retried) + + ledger.commitSend(ledger.reserveNextSend(oldOwner, 2)!) + const rotation = ledger.rotate(oldOwner, replacement, 2, 16) + expect(rotation.recovery.map((span) => span.data).join('')).toBe('cdef') + expect(getCursorRecord(ledger, replacement).sendSpanIndex).toBe(0) + // Rotation is the only path that removes and re-adds a record in one call. + expectRetentionMatchesRecords(ledger) + ledger.commitSend(ledger.reserveNextSend(replacement, 16)!) + ledger.commitSend(ledger.reserveNextSend(replacement, 16)!) + ledger.seal(replacement) + ledger.settleExitPublication(replacement, { ok: true }) + ledger.acknowledge(replacement, { + id: replacement.id, + clientGeneration: replacement.clientGeneration, + ownerGeneration: replacement.ownerGeneration, + deliveryToken: replacement.deliveryToken, + creditedEndSu: 6 + }) + expect(ledger.snapshotIfKnown(replacement)?.state).toBe('closed') + + expectRetentionMatchesRecords(ledger) + + const canceled = identity('token-canceled') + ledger.open(canceled, 8) + append(ledger, canceled, 'x', 'span-cancel') + expectRetentionMatchesRecords(ledger) + ledger.cancel(canceled, 'test-close') + expect(ledger.snapshotIfKnown(canceled)?.state).toBe('closed') + expectRetentionMatchesRecords(ledger) + expect(ledger.retentionSnapshot()).toEqual({ sourceSu: 0, dataBytes: 0, spans: 0 }) + }) + + it('rebases the send cursor when ACK reclaim removes spans at or past it', () => { + const ledger = new RelayPtySourceCreditLedger() + const owner = identity('token-reclaim') + ledger.open(owner, 16) + append(ledger, owner, 'ab', 'span-a') + append(ledger, owner, 'cd', 'span-b') + append(ledger, owner, 'ef', 'span-c') + ledger.commitSend(ledger.reserveNextSend(owner, 2)!) + ledger.commitSend(ledger.reserveNextSend(owner, 2)!) + expect(getCursorRecord(ledger, owner).sendSpanIndex).toBe(1) + + ledger.acknowledge(owner, { + id: owner.id, + clientGeneration: owner.clientGeneration, + ownerGeneration: owner.ownerGeneration, + deliveryToken: owner.deliveryToken, + creditedEndSu: 4 + }) + + // Reclaim dropped both spans the cursor had passed, so it must rebase onto the new head. + const record = getCursorRecord(ledger, owner) + expect(record.spans.map((span) => span.data)).toEqual(['ef']) + expect(record.sendSpanIndex).toBe(0) + expect(ledger.reserveNextSend(owner, 2)!.span.data).toBe('ef') + }) + it('never exceeds a token source window across generated send/ACK sequences', () => { for (let seed = 1; seed <= 40; seed++) { const ledger = new RelayPtySourceCreditLedger() const owner = identity(`token-${seed}`) const windowSu = 7 + (seed % 17) ledger.open(owner, windowSu) - append(ledger, owner, 'x'.repeat(100), `span-${seed}`) + for (let part = 0; part < 20; part += 1) { + append(ledger, owner, 'x'.repeat(5), `span-${seed}-${part}`) + } for (let turn = 0; turn < 100; turn++) { const reservation = drainOne(ledger, owner, 1 + ((seed * 13 + turn * 7) % 19)) const snapshot = ledger.snapshot(owner) expect(snapshot.sentEndSu - snapshot.creditedEndSu).toBeLessThanOrEqual(windowSu) + const record = getCursorRecord(ledger, owner) + const containingIndex = record.spans.findIndex( + (span) => + span.sourceStartSu <= snapshot.sentEndSu && span.sourceEndSu > snapshot.sentEndSu + ) + // Cursor may lag (advancement is lazy) but must never overshoot the containing span. + if (containingIndex !== -1) { + expect(record.sendSpanIndex).toBeLessThanOrEqual(containingIndex) + } if (snapshot.sentEndSu > snapshot.creditedEndSu && (turn + seed) % 3 === 0) { ledger.acknowledge(owner, { id: owner.id, @@ -73,6 +271,7 @@ describe('RelayPtySourceCreditLedger', () => { creditedEndSu: snapshot.sentEndSu }) } + expectRetentionMatchesRecords(ledger) if (!reservation && snapshot.creditedEndSu === 100) { break } diff --git a/src/relay/pty-source-credit-ledger.ts b/src/relay/pty-source-credit-ledger.ts index be628672c91..db6bd628f41 100644 --- a/src/relay/pty-source-credit-ledger.ts +++ b/src/relay/pty-source-credit-ledger.ts @@ -25,12 +25,10 @@ import { createDeliveryCancellation, createDeliveryRecord, createReplacementDeliveryRecord, + findPtySourceSpanForSend, matchingDeliverySnapshot, MAX_SOURCE_SPAN_DATA_BYTES, ptyOwnerKey, - retainedDataBytesTotal, - retainedSpanTotal, - retainedSourceTotal, sliceForSend, snapshotDeliveryRecord, type DeliveryRecord, @@ -45,6 +43,7 @@ import { type PtySourceAckResult } from './pty-source-credit-settlement' import { chargedPtyRetainedStringBytes } from '../shared/pty-retained-string-memory' +import { PtySourceCreditRetention } from './pty-source-credit-retention' export type { PtySourceSendReservation } from './pty-source-credit-record' export type { PtySourceCreditLedgerOptions } from './pty-source-credit-limits' @@ -54,6 +53,7 @@ export class RelayPtySourceCreditLedger { private readonly upstreamOwnerByPty = new Map() private readonly closedSnapshots = new Map() private readonly limits: PtySourceCreditLimits + private readonly retention = new PtySourceCreditRetention() private nextReservationId = 1 constructor(options: PtySourceCreditLedgerOptions = {}) { @@ -115,6 +115,7 @@ export class RelayPtySourceCreditLedger { record.spans.push(span) record.retainedDataBytes += retainedDataBytes record.receivedEndSu = sourceEndSu + this.retention.addSpan(input.transform.rawLengthSu, retainedDataBytes) return span } @@ -134,10 +135,8 @@ export class RelayPtySourceCreditLedger { if (remainingWindowSu <= 0 || record.sentEndSu >= record.receivedEndSu) { return null } - const containing = record.spans.find( - (span) => span.sourceStartSu <= record.sentEndSu && span.sourceEndSu > record.sentEndSu - ) - if (!containing) { + const containing = findPtySourceSpanForSend(record) + if (!containing || containing.sourceStartSu > record.sentEndSu) { throw new Error('PTY source delivery cursor is not covered by the retained ledger') } const reservedLengthSu = reservedPtySourceSendLength(record, remainingWindowSu, maxSourceSu) @@ -164,7 +163,10 @@ export class RelayPtySourceCreditLedger { commitSend(reservation: PtySourceSendReservation): void { const record = this.requireDelivery(reservation.identity) - if (commitPtySourceSend(record, reservation)) { + const settled = this.retention.trackMutation(record, () => + commitPtySourceSend(record, reservation) + ) + if (settled) { this.maybeClose(record) } } @@ -176,7 +178,9 @@ export class RelayPtySourceCreditLedger { acknowledge(identity: PtySourceDeliveryIdentity, ack: PtySourceCreditAck): PtySourceAckResult { const record = this.requireDelivery(identity) - const result = acknowledgeDeliveryRecord(record, ack) + const result = this.retention.trackMutation(record, () => + acknowledgeDeliveryRecord(record, ack) + ) if (result === 'advanced') { this.maybeClose(record) } @@ -244,6 +248,8 @@ export class RelayPtySourceCreditLedger { ) this.upstreamOwnerByPty.delete(ptyOwnerKey(old.identity)) this.deliveries.set(replacementKey, replacement) + // Retention mirrors the delivery map, so count the replacement as it enters, not after cancel. + this.retention.addRecord(replacement) this.upstreamOwnerByPty.set(ptyOwnerKey(replacement.identity), replacementKey) const cancellation = this.cancel(oldIdentity, 'superseded', newIdentity.deliveryToken) return Object.freeze({ cancellation, recovery: Object.freeze(replacement.spans.slice()) }) @@ -263,23 +269,13 @@ export class RelayPtySourceCreditLedger { return snapshot } - retainedSourceSu = (): number => retainedSourceTotal(this.deliveries.values()) + retainedSourceSu = (): number => this.retention.retainedSourceSu() - retainedDataBytes = (): number => retainedDataBytesTotal(this.deliveries.values()) + retainedDataBytes = (): number => this.retention.retainedDataBytes() - retainedSpans = (): number => retainedSpanTotal(this.deliveries.values()) + retainedSpans = (): number => this.retention.retainedSpans() - retentionSnapshot(): Readonly<{ - sourceSu: number - dataBytes: number - spans: number - }> { - return Object.freeze({ - sourceSu: this.retainedSourceSu(), - dataBytes: this.retainedDataBytes(), - spans: this.retainedSpans() - }) - } + retentionSnapshot = () => this.retention.snapshot() private requireActive(identity: PtySourceDeliveryIdentity): DeliveryRecord { const record = this.requireDelivery(identity) @@ -308,10 +304,15 @@ export class RelayPtySourceCreditLedger { } private closeRecord(record: DeliveryRecord): void { + if (record.state === 'closed') { + return + } + this.retention.removeRecord(record) record.state = 'closed' record.pendingSend = null record.spans = [] record.retainedDataBytes = 0 + record.sendSpanIndex = 0 const key = ptySourceDeliveryKey(record.identity) const ownerKey = ptyOwnerKey(record.identity) if (this.upstreamOwnerByPty.get(ownerKey) === key) { diff --git a/src/relay/pty-source-credit-record.ts b/src/relay/pty-source-credit-record.ts index d2730d25dc7..606edc7dfb0 100644 --- a/src/relay/pty-source-credit-record.ts +++ b/src/relay/pty-source-credit-record.ts @@ -42,6 +42,8 @@ export type DeliveryRecord = { creditedEndSu: number retainedDataBytes: number spans: PtySourceSpan[] + /** First retained span that may contain the next unsent source unit. */ + sendSpanIndex: number sentBoundaries: Set pendingSend: PtySourceSendReservation | null reservedAckEndSu: number | null @@ -68,6 +70,7 @@ export function createDeliveryRecord( creditedEndSu: checkpointSourceEndSu, retainedDataBytes: 0, spans: [], + sendSpanIndex: 0, sentBoundaries: new Set([checkpointSourceEndSu]), pendingSend: null, reservedAckEndSu: null, @@ -160,6 +163,15 @@ export function sliceForSend( }) } +export function findPtySourceSpanForSend(record: DeliveryRecord): PtySourceSpan | undefined { + let span = record.spans[record.sendSpanIndex] + while (span && span.sourceEndSu <= record.sentEndSu) { + record.sendSpanIndex += 1 + span = record.spans[record.sendSpanIndex] + } + return span +} + export function snapshotDeliveryRecord(record: DeliveryRecord): PtySourceDeliverySnapshot { return Object.freeze({ ...record.identity, @@ -188,30 +200,6 @@ export function matchingDeliverySnapshot( return closed && samePtySourceDelivery(closed, identity) ? closed : null } -export function retainedSourceTotal(records: Iterable): number { - let total = 0 - for (const record of records) { - total += record.receivedEndSu - record.creditedEndSu - } - return total -} - -export function retainedDataBytesTotal(records: Iterable): number { - let total = 0 - for (const record of records) { - total += record.retainedDataBytes - } - return total -} - -export function retainedSpanTotal(records: Iterable): number { - let total = 0 - for (const record of records) { - total += record.spans.length - } - return total -} - export function createReplacementDeliveryRecord( old: DeliveryRecord, newIdentity: PtySourceDeliveryIdentity, @@ -246,6 +234,7 @@ export function createReplacementDeliveryRecord( .filter((span) => span.sourceEndSu > acceptedSourceEndSu) .map((span) => sliceAtSourceStart(span, Math.max(span.sourceStartSu, acceptedSourceEndSu))) .map((span) => Object.freeze({ ...span, ...replacement.identity })) + replacement.sendSpanIndex = 0 replacement.retainedDataBytes = replacement.spans.reduce( (bytes, span) => bytes + chargedPtyRetainedStringBytes(span.data), 0 diff --git a/src/relay/pty-source-credit-retention.ts b/src/relay/pty-source-credit-retention.ts new file mode 100644 index 00000000000..742a660e374 --- /dev/null +++ b/src/relay/pty-source-credit-retention.ts @@ -0,0 +1,69 @@ +import type { DeliveryRecord } from './pty-source-credit-record' + +export type PtySourceCreditRetentionSnapshot = Readonly<{ + sourceSu: number + dataBytes: number + spans: number +}> + +type RecordRetention = PtySourceCreditRetentionSnapshot + +export class PtySourceCreditRetention { + private sourceSuTotal = 0 + private dataBytesTotal = 0 + private spansTotal = 0 + + addSpan(sourceSu: number, dataBytes: number): void { + this.sourceSuTotal += sourceSu + this.dataBytesTotal += dataBytes + this.spansTotal += 1 + } + + addRecord(record: DeliveryRecord): void { + this.apply(recordRetention(record), 1) + } + + removeRecord(record: DeliveryRecord): void { + this.apply(recordRetention(record), -1) + } + + trackMutation(record: DeliveryRecord, mutate: () => T): T { + const before = recordRetention(record) + try { + return mutate() + } finally { + const after = recordRetention(record) + this.sourceSuTotal += after.sourceSu - before.sourceSu + this.dataBytesTotal += after.dataBytes - before.dataBytes + this.spansTotal += after.spans - before.spans + } + } + + retainedSourceSu = (): number => this.sourceSuTotal + + retainedDataBytes = (): number => this.dataBytesTotal + + retainedSpans = (): number => this.spansTotal + + snapshot(): PtySourceCreditRetentionSnapshot { + return Object.freeze({ + sourceSu: this.sourceSuTotal, + dataBytes: this.dataBytesTotal, + spans: this.spansTotal + }) + } + + private apply(retention: RecordRetention, direction: 1 | -1): void { + this.sourceSuTotal += direction * retention.sourceSu + this.dataBytesTotal += direction * retention.dataBytes + this.spansTotal += direction * retention.spans + } +} + +function recordRetention(record: DeliveryRecord): RecordRetention { + return { + sourceSu: record.receivedEndSu - record.creditedEndSu, + dataBytes: record.retainedDataBytes, + spans: record.spans.length + } +} diff --git a/src/relay/pty-source-credit-settlement.ts b/src/relay/pty-source-credit-settlement.ts index 464ea9cefe1..e12e73521e5 100644 --- a/src/relay/pty-source-credit-settlement.ts +++ b/src/relay/pty-source-credit-settlement.ts @@ -18,8 +18,13 @@ export function reservedPtySourceSendLength( } function reclaimCreditedSpans(record: DeliveryRecord): void { + let removed = 0 while (record.spans[0]?.sourceEndSu <= record.creditedEndSu) { record.retainedDataBytes -= chargedPtyRetainedStringBytes(record.spans.shift()!.data) + removed += 1 + } + if (removed > 0) { + record.sendSpanIndex = Math.max(0, record.sendSpanIndex - removed) } }