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

@@ -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<string, CursorRecord> }
+  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<string, DeliveryRecord> }
+  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
         }

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<string, string>()
   private readonly closedSnapshots = new Map<string, PtySourceDeliverySnapshot>()
   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) {

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<number>
   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<DeliveryRecord>): number {
-  let total = 0
-  for (const record of records) {
-    total += record.receivedEndSu - record.creditedEndSu
-  }
-  return total
-}
-
-export function retainedDataBytesTotal(records: Iterable<DeliveryRecord>): number {
-  let total = 0
-  for (const record of records) {
-    total += record.retainedDataBytes
-  }
-  return total
-}
-
-export function retainedSpanTotal(records: Iterable<DeliveryRecord>): 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

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<T>(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
+  }
+}

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