diff --git a/src/main/ipc/ssh-pty-legacy-projection-record.ts b/src/main/ipc/ssh-pty-legacy-projection-record.ts index 9f5077c8d4d..ead64d3d6f9 100644 --- a/src/main/ipc/ssh-pty-legacy-projection-record.ts +++ b/src/main/ipc/ssh-pty-legacy-projection-record.ts @@ -113,32 +113,6 @@ export function requireProjectionRecord( return record } -export function rollbackUnpublishedProjection( - records: Map, - cursors: ReadonlyMap, - idsByPty: Map, - id: string -): boolean { - const record = records.get(id) - if (!record || record.state !== 'committed') { - return false - } - const { identity, beforeScanner } = record.semantics - const cursor = cursors.get(identity.ptyId) - if ( - !cursor || - cursor.providerGeneration !== identity.providerGeneration || - cursor.ptyIncarnation !== identity.ptyIncarnation || - cursor.displayEnd !== identity.displayEnd - ) { - return false - } - cursor.displayEnd = identity.displayStart - cursor.scanner = { ...beforeScanner } - reclaimProjectionRecord(records, idsByPty, id, identity.ptyId) - return true -} - export function closeProjectionPty( cursors: Map, ptyId: string, diff --git a/src/main/ipc/ssh-pty-legacy-projection.ts b/src/main/ipc/ssh-pty-legacy-projection.ts index c1042d4815f..435211d5ca6 100644 --- a/src/main/ipc/ssh-pty-legacy-projection.ts +++ b/src/main/ipc/ssh-pty-legacy-projection.ts @@ -9,7 +9,6 @@ import { projectionError, reclaimProjectionRecord, requireProjectionRecord, - rollbackUnpublishedProjection, scannerSnapshot, type LegacySshProjectionDebugSnapshot, type LegacySshProjectionReservation, @@ -212,15 +211,8 @@ export class SshPtyLegacyProjectionLedger { return transferred } - transferUnpublished(id: string, _reason: string): boolean { - const transferred = rollbackUnpublishedProjection( - this.records, - this.cursorByPty, - this.idsByPty, - id - ) - this.transferredCount += Number(transferred) - return transferred + transferUnpublished(id: string, reason: string): boolean { + return this.transfer([id], reason) === 1 } transferGeneration(providerGeneration: number, reason: string): number { diff --git a/src/main/ipc/ssh-pty-model-admission-entry.ts b/src/main/ipc/ssh-pty-model-admission-entry.ts index e73884f1efa..a9894181e71 100644 --- a/src/main/ipc/ssh-pty-model-admission-entry.ts +++ b/src/main/ipc/ssh-pty-model-admission-entry.ts @@ -12,7 +12,7 @@ export type AdmissionEntry = { run: () => { sequence: number; completion: Promise } resolve: (receipt: SshPtyModelAdmissionReceipt) => void reject: (error: Error) => void - state: 'queued' | 'running' | 'pressure' + state: 'queued' | 'running' | 'pressure' | 'settled' } export type PtyUsage = AdmissionCharge & { @@ -74,6 +74,47 @@ export function takePressureEntriesForGeneration( return removed } +export function cancelAdmissionGeneration(args: { + pressure: AdmissionEntry[] + usageByPty: Map + idleWaiters: Map void>> + providerGeneration: number + error: Error + release: (key: SshPtyModelAdmissionKey, charge: AdmissionCharge) => void +}): number { + let pressureBytes = 0 + for (const entry of takePressureEntriesForGeneration(args.pressure, args.providerGeneration)) { + pressureBytes += entry.charge.bytes + entry.state = 'settled' + entry.reject(args.error) + } + for (const [id, usage] of args.usageByPty) { + const canceled = usage.queued.filter( + (entry) => entry.key.providerGeneration === args.providerGeneration + ) + usage.queued = usage.queued.filter( + (entry) => entry.key.providerGeneration !== args.providerGeneration + ) + if (usage.running?.key.providerGeneration === args.providerGeneration) { + canceled.push(usage.running) + usage.running = null + } + for (const entry of canceled) { + if (entry.state === 'settled') { + continue + } + entry.state = 'settled' + args.release(entry.key, entry.charge) + entry.reject(args.error) + } + if (!usage.running && usage.queued.length === 0 && usage.sourceUnits === 0) { + args.usageByPty.delete(id) + } + resolveAdmissionIdleWaiters(args.usageByPty, args.pressure, args.idleWaiters, id) + } + return pressureBytes +} + export function takePausedGeneration( paused: Map, providerGeneration: number diff --git a/src/main/ipc/ssh-pty-model-admission.test.ts b/src/main/ipc/ssh-pty-model-admission.test.ts index 95178659c44..4d267ef83b0 100644 --- a/src/main/ipc/ssh-pty-model-admission.test.ts +++ b/src/main/ipc/ssh-pty-model-admission.test.ts @@ -1,34 +1,71 @@ -import { describe, expect, it } from 'vitest' +import { describe, expect, it, vi } from 'vitest' import { SshPtyModelAdmission } from './ssh-pty-model-admission' function deferred() { let resolve!: () => void - const promise = new Promise((promiseResolve) => { + let reject!: (error: Error) => void + const promise = new Promise((promiseResolve, promiseReject) => { resolve = promiseResolve + reject = promiseReject }) - return { promise, resolve } + return { promise, resolve, reject } +} + +function accept(admission: SshPtyModelAdmission, completion: Promise) { + return admission.accept({ ptyId: 'pty-1', providerGeneration: 7 }, 'data', 4, () => ({ + sequence: 4, + completion + })) } describe('SshPtyModelAdmission', () => { - it('owns the sole running entry through disposal', async () => { - const completion = deferred() + it('cancels a never-settling running entry when its generation closes', async () => { const admission = new SshPtyModelAdmission() - const receipt = admission.accept({ ptyId: 'pty-1', providerGeneration: 7 }, 'data', 4, () => ({ - sequence: 4, - completion: completion.promise - })) - let idle = false - const idleReceipt = admission.whenIdle({ ptyId: 'pty-1', providerGeneration: 7 }).then(() => { - idle = true - }) + const receipt = accept(admission, new Promise(() => {})) + const idle = admission.whenIdle({ ptyId: 'pty-1', providerGeneration: 7 }) - admission.dispose() - await Promise.resolve() - expect(idle).toBe(false) - completion.resolve() + admission.closeGeneration(7, 'provider-closed') - await expect(receipt).rejects.toThrow('ssh_model_admission_generation_closed') - await idleReceipt + await expect(receipt).rejects.toThrow('provider-closed') + await expect(idle).resolves.toBeUndefined() expect(admission.getDebugSnapshot()).toMatchObject({ sourceUnits: 0, bytes: 0 }) }) + + it('cancels a never-settling running entry on disposal', async () => { + const admission = new SshPtyModelAdmission() + const receipt = accept(admission, new Promise(() => {})) + const idle = admission.whenIdle({ ptyId: 'pty-1', providerGeneration: 7 }) + + admission.dispose() + + await expect(receipt).rejects.toThrow('ssh_model_admission_disposed') + await expect(idle).resolves.toBeUndefined() + expect(admission.getDebugSnapshot()).toMatchObject({ sourceUnits: 0, bytes: 0 }) + }) + + it.each(['resolve', 'reject'] as const)( + 'ignores a late completion %s after cancellation', + async (settle) => { + const completion = deferred() + const admission = new SshPtyModelAdmission() + const onResolve = vi.fn() + const onReject = vi.fn() + const receipt = accept(admission, completion.promise) + const observed = receipt.then(onResolve, onReject) + + admission.closeGeneration(7, 'provider-closed') + await observed + if (settle === 'resolve') { + completion.resolve() + } else { + completion.reject(new Error('late emulator failure')) + } + await Promise.resolve() + + expect(onResolve).not.toHaveBeenCalled() + expect(onReject).toHaveBeenCalledTimes(1) + expect(onReject.mock.calls[0]?.[0]).toMatchObject({ message: 'provider-closed' }) + expect(admission.getDebugSnapshot()).toMatchObject({ sourceUnits: 0, bytes: 0 }) + } + ) }) diff --git a/src/main/ipc/ssh-pty-model-admission.ts b/src/main/ipc/ssh-pty-model-admission.ts index 792c5a9cc5b..e074b707aed 100644 --- a/src/main/ipc/ssh-pty-model-admission.ts +++ b/src/main/ipc/ssh-pty-model-admission.ts @@ -8,11 +8,11 @@ import { admissionError, admissionKeyId, canReserveAdmission, + cancelAdmissionGeneration, pressureHasAdmissionKey, retainedBytes, resolveAdmissionIdleWaiters, takePausedGeneration, - takePressureEntriesForGeneration, type AdmissionCharge, type AdmissionEntry, type PtyUsage @@ -83,26 +83,14 @@ export class SshPtyModelAdmission { closeGeneration(providerGeneration: number, reason = 'provider-generation-closed'): void { this.closingGenerations.add(providerGeneration) const error = admissionError(reason) - for (const entry of takePressureEntriesForGeneration(this.pressure, providerGeneration)) { - this.pressureBytes -= entry.charge.bytes - entry.reject(error) - } - for (const [id, usage] of this.usageByPty) { - const queued = usage.queued.filter( - (entry) => entry.key.providerGeneration === providerGeneration - ) - usage.queued = usage.queued.filter( - (entry) => entry.key.providerGeneration !== providerGeneration - ) - for (const entry of queued) { - this.release(entry.key, entry.charge) - entry.reject(error) - } - if (!usage.running && usage.queued.length === 0 && usage.sourceUnits === 0) { - this.usageByPty.delete(id) - } - resolveAdmissionIdleWaiters(this.usageByPty, this.pressure, this.idleWaiters, id) - } + this.pressureBytes -= cancelAdmissionGeneration({ + pressure: this.pressure, + usageByPty: this.usageByPty, + idleWaiters: this.idleWaiters, + providerGeneration, + error, + release: (key, charge) => this.release(key, charge) + }) for (const key of takePausedGeneration(this.pausedKeys, providerGeneration)) { try { this.resumeProvider(key) @@ -228,35 +216,42 @@ export class SshPtyModelAdmission { try { execution = entry.run() } catch (error) { - this.finishEntry(id, usage, entry) - entry.reject(error instanceof Error ? error : new Error(String(error))) + if (this.finishEntry(id, usage, entry)) { + entry.reject(error instanceof Error ? error : new Error(String(error))) + } return } void execution.completion.then( () => { - this.finishEntry(id, usage, entry) - if (this.closingGenerations.has(entry.key.providerGeneration)) { - entry.reject(admissionError('ssh_model_admission_generation_closed')) - } else { - entry.resolve({ - ptyId: entry.key.ptyId, - providerGeneration: entry.key.providerGeneration, - sequence: execution.sequence - }) + if (!this.finishEntry(id, usage, entry)) { + return } + entry.resolve({ + ptyId: entry.key.ptyId, + providerGeneration: entry.key.providerGeneration, + sequence: execution.sequence + }) }, (error) => { - this.finishEntry(id, usage, entry) - entry.reject(error instanceof Error ? error : new Error(String(error))) + if (this.finishEntry(id, usage, entry)) { + entry.reject(error instanceof Error ? error : new Error(String(error))) + } } ) } - private finishEntry(id: string, usage: PtyUsage, entry: AdmissionEntry): void { - if (usage.running === entry) { - usage.running = null + private finishEntry(id: string, usage: PtyUsage, entry: AdmissionEntry): boolean { + if (entry.state !== 'running' || usage.running !== entry) { + return false } + usage.running = null + entry.state = 'settled' this.release(entry.key, entry.charge) + this.cleanupUsage(id, usage) + return true + } + + private cleanupUsage(id: string, usage: PtyUsage): void { this.startNext(id, usage) if (!usage.running && usage.queued.length === 0 && usage.sourceUnits === 0) { this.usageByPty.delete(id) @@ -276,6 +271,9 @@ export class SshPtyModelAdmission { } private maybeResumeAndPromote(): void { + if (this.disposed) { + return + } for (let index = 0; index < this.pressure.length; ) { const entry = this.pressure[index]! const hasEarlierEntryForPty = this.pressure @@ -309,7 +307,9 @@ export class SshPtyModelAdmission { continue } this.pausedKeys.delete(id) - this.resumeProvider(key) + try { + this.resumeProvider(key) + } catch {} } } } diff --git a/src/main/ipc/ssh-pty-output-intake.test.ts b/src/main/ipc/ssh-pty-output-intake.test.ts index 82c7c428943..49f37d4fdd8 100644 --- a/src/main/ipc/ssh-pty-output-intake.test.ts +++ b/src/main/ipc/ssh-pty-output-intake.test.ts @@ -195,8 +195,9 @@ describe('SshPtyOutputIntake', () => { harness.completions[1]!.resolve() await next expect(projections[1]).toMatchObject({ - identity: { displayStart: 0 }, - beforeScanner: { tail: '', pendingSubscribe: false } + identity: { displayStart: 5 }, + beforeScanner: { tail: '\x1b[?20', pendingSubscribe: false }, + decision: 'subscribed' }) }) @@ -319,8 +320,10 @@ describe('SshPtyOutputIntake', () => { await expect(queued).rejects.toThrow('provider-closed') await expect(exit).rejects.toThrow('provider-closed') await expect(harness.intake.acceptData(event())).rejects.toThrow('ssh_output_stale_generation') + await expect(running).rejects.toThrow('provider-closed') + expect(harness.intake.getDebugSnapshot().model).toMatchObject({ sourceUnits: 0, bytes: 0 }) harness.completions[0]!.resolve() - await expect(running).rejects.toThrow('ssh_model_admission_generation_closed') + await Promise.resolve() expect(harness.intake.getDebugSnapshot().exitBarriers).toBe(0) }) @@ -343,7 +346,8 @@ describe('SshPtyOutputIntake', () => { await expect(pressured).rejects.toThrow('provider-closed') expect(harness.dependencies.resumeProvider).toHaveBeenCalledWith(1, 'pty-1') + await expect(running).rejects.toThrow('provider-closed') harness.completions[0]!.resolve() - await expect(running).rejects.toThrow('ssh_model_admission_generation_closed') + await Promise.resolve() }) })