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) } }