diff --git a/src/main/persistence.ts b/src/main/persistence.ts index 48321b3b11d..bffa91e7a28 100644 --- a/src/main/persistence.ts +++ b/src/main/persistence.ts @@ -23,6 +23,10 @@ import { writeTerminalTeardownIntentSnapshot } from './terminal-teardown-intent-snapshot' import type { TerminalTeardownIntentMutation } from './terminal-teardown-intent-snapshot' +import { + indexSshRemotePtyLeases, + sshRemotePtyLeaseIdentityKey +} from './ssh/ssh-remote-pty-lease-index' import type { Automation, AutomationCreateInput, @@ -2668,6 +2672,7 @@ export class Store { private firstPendingSaveAt: number | null = null private githubCacheDirty = false private terminalTeardownIntentJournalReady = false + private sshRemotePtyLeasesByIdentity: Map | null = null private gitUsernameCache = new Map() private loadNeedsSave = false private settingsChangeListeners = new Set< @@ -5785,6 +5790,7 @@ export class Store { ) if (remappedLeases.changed) { this.state.sshRemotePtyLeases = remappedLeases.leases + this.sshRemotePtyLeasesByIdentity = null } if (session && prior) { const priorTabs = prior.tabsByWorktree ?? {} @@ -6388,6 +6394,13 @@ export class Store { return leases.filter((lease) => targetId === undefined || lease.targetId === targetId) } + private getSshRemotePtyLeasesByIdentity(): Map { + this.sshRemotePtyLeasesByIdentity ??= indexSshRemotePtyLeases( + this.state.sshRemotePtyLeases ?? [] + ) + return this.sshRemotePtyLeasesByIdentity + } + getPendingLocalPtyShutdowns(): PersistedLocalPtyShutdown[] { return [...(this.state.pendingLocalPtyShutdowns ?? [])] } @@ -6510,6 +6523,7 @@ export class Store { } else { this.state.sshRemotePtyLeases.push(next) } + this.sshRemotePtyLeasesByIdentity = null this.flush() this.persistTerminalTeardownIntents(this.sshTeardownIntentMutationForLease(next)) } @@ -6564,11 +6578,8 @@ export class Store { ): boolean { const requestedGeneration = parseAppSshPtyId(ptyId)?.relayInstanceId const relayPtyId = this.getRelayPtyIdForSshLeaseStorage(targetId, ptyId) - const lease = this.state.sshRemotePtyLeases?.find( - (entry) => - entry.targetId === targetId && - entry.ptyId === relayPtyId && - entry.relayInstanceId === requestedGeneration + const lease = this.getSshRemotePtyLeasesByIdentity().get( + sshRemotePtyLeaseIdentityKey(targetId, relayPtyId, requestedGeneration) ) if (!lease) { return false @@ -6610,17 +6621,16 @@ export class Store { relayInstanceId: string ): boolean { const relayPtyId = this.getRelayPtyIdForSshLeaseStorage(targetId, ptyId) - const lease = this.state.sshRemotePtyLeases?.find( - (entry) => - entry.targetId === targetId && - entry.ptyId === relayPtyId && - entry.relayInstanceId === undefined - ) + const legacyKey = sshRemotePtyLeaseIdentityKey(targetId, relayPtyId) + const leasesByIdentity = this.getSshRemotePtyLeasesByIdentity() + const lease = leasesByIdentity.get(legacyKey) if (!lease) { return false } + leasesByIdentity.delete(legacyKey) const now = Date.now() lease.relayInstanceId = relayInstanceId + leasesByIdentity.set(sshRemotePtyLeaseIdentityKey(targetId, relayPtyId, relayInstanceId), lease) lease.state = 'attached' lease.updatedAt = now lease.lastAttachedAt = now @@ -6639,11 +6649,8 @@ export class Store { markSshRemotePtyShutdownRequested(targetId: string, ptyId: string): boolean { const requestedGeneration = parseAppSshPtyId(ptyId)?.relayInstanceId const relayPtyId = this.getRelayPtyIdForSshLeaseStorage(targetId, ptyId) - const lease = this.state.sshRemotePtyLeases?.find( - (entry) => - entry.targetId === targetId && - entry.ptyId === relayPtyId && - entry.relayInstanceId === requestedGeneration + const lease = this.getSshRemotePtyLeasesByIdentity().get( + sshRemotePtyLeaseIdentityKey(targetId, relayPtyId, requestedGeneration) ) if (!lease) { return false @@ -6672,6 +6679,7 @@ export class Store { lease.relayInstanceId !== requestedGeneration ) if (this.state.sshRemotePtyLeases.length !== before) { + this.sshRemotePtyLeasesByIdentity = null this.persistTerminalTeardownIntents({ kind: 'ssh-set', targetId, @@ -6689,6 +6697,7 @@ export class Store { (lease) => lease.targetId !== targetId ) if (this.state.sshRemotePtyLeases.length !== before) { + this.sshRemotePtyLeasesByIdentity = null this.persistTerminalTeardownIntents({ kind: 'ssh-remove-target', targetId }) } } diff --git a/src/main/ssh/ssh-relay-generation-reattach.test.ts b/src/main/ssh/ssh-relay-generation-reattach.test.ts index 1ca543af27c..1b43a005d77 100644 --- a/src/main/ssh/ssh-relay-generation-reattach.test.ts +++ b/src/main/ssh/ssh-relay-generation-reattach.test.ts @@ -1,5 +1,5 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' -import { mkdtempSync, rmSync } from 'node:fs' +import { mkdtempSync, readFileSync, rmSync } from 'node:fs' import { tmpdir } from 'node:os' import { join } from 'node:path' import { SshRelaySession } from './ssh-relay-session' @@ -152,14 +152,21 @@ describe('SSH relay generation reattach', () => { it('migrates a legacy lease only after the current relay proves it can attach', async () => { const dir = mkdtempSync(join(tmpdir(), 'orca-legacy-ssh-lease-')) const dataFile = join(dir, 'state.json') - const store = new Store({ dataFile }) - const reloadedStores: Store[] = [] + const stores: Store[] = [] try { - store.upsertSshRemotePtyLease({ + const seedStore = new Store({ dataFile }) + stores.push(seedStore) + seedStore.upsertSshRemotePtyLease({ targetId: 'target-1', ptyId: 'pty-legacy', state: 'detached' }) + seedStore.flush() + // Why: an upgraded profile predates the teardown journal entirely. The + // first migration must survive before the debounced full save runs. + rmSync(join(dir, 'orca-terminal-teardown-intents.json'), { force: true }) + const store = new Store({ dataFile }) + stores.push(store) const synchronousFlush = vi.spyOn(store, 'flush') const { mockConn, mockPortForward, getMainWindow } = createMockDeps() const attachForReconnect = vi.fn().mockResolvedValue({}) @@ -184,19 +191,73 @@ describe('SSH relay generation reattach', () => { ]) expect(synchronousFlush).not.toHaveBeenCalled() const reloaded = new Store({ dataFile }) - reloadedStores.push(reloaded) + stores.push(reloaded) expect(reloaded.getSshRemotePtyLeases('target-1')[0]).toEqual( expect.objectContaining({ relayInstanceId: 'relay-current', state: 'attached' }) ) } finally { - store.flush() - for (const reloaded of reloadedStores) { - reloaded.flush() + for (const store of stores) { + store.flush() } rmSync(dir, { recursive: true, force: true }) } }) + it('adopts many legacy leases with one linear index build and bounded records', () => { + const dir = mkdtempSync(join(tmpdir(), 'orca-ssh-lease-scale-')) + const dataFile = join(dir, 'state.json') + const store = new Store({ dataFile }) + const leaseCount = 50 + try { + for (let index = 0; index < leaseCount; index += 1) { + store.upsertSshRemotePtyLease({ + targetId: 'target-1', + ptyId: `pty-${index}`, + state: 'detached' + }) + } + const internals = store as unknown as { + state: { sshRemotePtyLeases: unknown[] } + } + const leases = internals.state.sshRemotePtyLeases + let indexedLeaseReads = 0 + internals.state.sshRemotePtyLeases = new Proxy(leases, { + get(target, property, receiver) { + if (typeof property === 'string' && /^\d+$/.test(property)) { + indexedLeaseReads += 1 + } + return Reflect.get(target, property, receiver) + } + }) + const synchronousFlush = vi.spyOn(store, 'flush') + + for (let index = 0; index < leaseCount; index += 1) { + expect( + store.migrateLegacySshRemotePtyLeaseGeneration( + 'target-1', + `pty-${index}`, + 'relay-current' + ) + ).toBe(true) + } + + expect(indexedLeaseReads).toBe(leaseCount) + expect(synchronousFlush).not.toHaveBeenCalled() + const records = readFileSync(join(dir, 'orca-terminal-teardown-intents.json'), 'utf-8') + .trim() + .split('\n') + .map((line) => JSON.parse(line) as { kind: string }) + const migrations = records.filter((record) => record.kind === 'ssh-migrate-generation') + expect(migrations).toHaveLength(leaseCount) + expect(Math.max(...migrations.map((record) => JSON.stringify(record).length))).toBeLessThan( + 1_024 + ) + } finally { + store.flush() + rmSync(dir, { recursive: true, force: true }) + } + }) + it('expires a failed legacy identity without adopting the current relay generation', async () => { const dir = mkdtempSync(join(tmpdir(), 'orca-stale-ssh-lease-')) const store = new Store({ dataFile: join(dir, 'state.json') }) diff --git a/src/main/ssh/ssh-remote-pty-lease-index.ts b/src/main/ssh/ssh-remote-pty-lease-index.ts new file mode 100644 index 00000000000..4cedeaef4db --- /dev/null +++ b/src/main/ssh/ssh-remote-pty-lease-index.ts @@ -0,0 +1,20 @@ +import type { SshRemotePtyLease } from '../../shared/ssh-types' + +export function sshRemotePtyLeaseIdentityKey( + targetId: string, + ptyId: string, + relayInstanceId?: string +): string { + return JSON.stringify([targetId, ptyId, relayInstanceId ?? null]) +} + +export function indexSshRemotePtyLeases( + leases: SshRemotePtyLease[] +): Map { + return new Map( + leases.map((lease) => [ + sshRemotePtyLeaseIdentityKey(lease.targetId, lease.ptyId, lease.relayInstanceId), + lease + ]) + ) +} diff --git a/src/main/terminal-teardown-intent-snapshot.ts b/src/main/terminal-teardown-intent-snapshot.ts index a6b5d3224d0..ebce871fb51 100644 --- a/src/main/terminal-teardown-intent-snapshot.ts +++ b/src/main/terminal-teardown-intent-snapshot.ts @@ -18,6 +18,11 @@ import type { TerminalTeardownIntentMutation, TerminalTeardownIntentSnapshot } from './terminal-teardown-intent-records' +import { + indexSshRemotePtyLeases, + sshRemotePtyLeaseIdentityKey +} from './ssh/ssh-remote-pty-lease-index' +import type { SshRemotePtyLease } from '../shared/ssh-types' export type { TerminalTeardownIntentMutation } from './terminal-teardown-intent-records' @@ -112,17 +117,15 @@ function readRecords(dataFile: string): { function applySshShutdowns( state: PersistedState, - pendingSshPtyShutdowns: TerminalTeardownIntentSnapshot['pendingSshPtyShutdowns'] + pendingSshPtyShutdowns: TerminalTeardownIntentSnapshot['pendingSshPtyShutdowns'], + leasesByIdentity: Map ): void { for (const lease of state.sshRemotePtyLeases ?? []) { delete lease.shutdownRequestedAt } for (const pending of pendingSshPtyShutdowns) { - const lease = state.sshRemotePtyLeases?.find( - (candidate) => - candidate.targetId === pending.targetId && - candidate.ptyId === pending.ptyId && - candidate.relayInstanceId === pending.relayInstanceId + const lease = leasesByIdentity.get( + sshRemotePtyLeaseIdentityKey(pending.targetId, pending.ptyId, pending.relayInstanceId) ) if (lease) { lease.shutdownRequestedAt = pending.shutdownRequestedAt @@ -133,13 +136,14 @@ function applySshShutdowns( function applyJournalRecord( state: PersistedState, - record: TerminalTeardownIntentJournalRecord + record: TerminalTeardownIntentJournalRecord, + leasesByIdentity: Map ): void { switch (record.kind) { case 'checkpoint': state.pendingLocalPtyShutdowns = record.pendingLocalPtyShutdowns state.pendingRuntimeTerminalCloses = record.pendingRuntimeTerminalCloses - applySshShutdowns(state, record.pendingSshPtyShutdowns) + applySshShutdowns(state, record.pendingSshPtyShutdowns, leasesByIdentity) return case 'local-upsert': { state.pendingLocalPtyShutdowns ??= [] @@ -183,27 +187,25 @@ function applyJournalRecord( ) return case 'ssh-migrate-generation': { - const lease = state.sshRemotePtyLeases?.find( - (candidate) => - candidate.targetId === record.targetId && - candidate.ptyId === record.ptyId && - candidate.relayInstanceId === undefined - ) + const legacyKey = sshRemotePtyLeaseIdentityKey(record.targetId, record.ptyId) + const lease = leasesByIdentity.get(legacyKey) if (!lease) { return } + leasesByIdentity.delete(legacyKey) lease.relayInstanceId = record.relayInstanceId + leasesByIdentity.set( + sshRemotePtyLeaseIdentityKey(record.targetId, record.ptyId, record.relayInstanceId), + lease + ) lease.state = 'attached' lease.updatedAt = Math.max(lease.updatedAt, record.attachedAt) lease.lastAttachedAt = Math.max(lease.lastAttachedAt ?? 0, record.attachedAt) return } case 'ssh-set': { - const lease = state.sshRemotePtyLeases?.find( - (candidate) => - candidate.targetId === record.targetId && - candidate.ptyId === record.ptyId && - candidate.relayInstanceId === record.relayInstanceId + const lease = leasesByIdentity.get( + sshRemotePtyLeaseIdentityKey(record.targetId, record.ptyId, record.relayInstanceId) ) if (!lease) { return @@ -230,6 +232,7 @@ export function applyTerminalTeardownIntentSnapshot( dataFile: string ): boolean { const { records, journalReady } = readRecords(dataFile) + const leasesByIdentity = indexSshRemotePtyLeases(state.sshRemotePtyLeases ?? []) for (const record of records) { // Why: a failed journal append falls back to a synchronous full-store // write. Older valid records must not override that newer durable state. @@ -243,17 +246,17 @@ export function applyTerminalTeardownIntentSnapshot( if (record.version === 1) { state.pendingLocalPtyShutdowns = record.pendingLocalPtyShutdowns state.pendingRuntimeTerminalCloses = record.pendingRuntimeTerminalCloses - applySshShutdowns(state, record.pendingSshPtyShutdowns) + applySshShutdowns(state, record.pendingSshPtyShutdowns, leasesByIdentity) } else { - applyJournalRecord(state, record) + applyJournalRecord(state, record, leasesByIdentity) } state.terminalTeardownIntentRevision = record.revision } return journalReady } -function writeRecordAtomically( - record: TerminalTeardownIntentJournalRecord, +function writeRecordsAtomically( + records: TerminalTeardownIntentJournalRecord[], dataFile: string ): void { const snapshotFile = getSnapshotFile(dataFile) @@ -261,7 +264,8 @@ function writeRecordAtomically( mkdirSync(dirname(snapshotFile), { recursive: true }) let renamed = false try { - writeFileSync(tmpFile, `${JSON.stringify(record)}\n`, 'utf-8') + const serialized = records.map((record) => JSON.stringify(record)).join('\n') + writeFileSync(tmpFile, `${serialized}\n`, 'utf-8') renameSync(tmpFile, snapshotFile) renamed = true } finally { @@ -276,7 +280,7 @@ function writeRecordAtomically( } export function writeTerminalTeardownIntentSnapshot(state: PersistedState, dataFile: string): void { - writeRecordAtomically({ version: 2, ...checkpointForState(state) }, dataFile) + writeRecordsAtomically([{ version: 2, ...checkpointForState(state) }], dataFile) } export function appendTerminalTeardownIntentMutation( @@ -286,6 +290,23 @@ export function appendTerminalTeardownIntentMutation( journalReady: boolean ): void { if (!journalReady) { + if (mutation.kind === 'ssh-migrate-generation') { + const revision = state.terminalTeardownIntentRevision ?? 0 + // Why: the checkpoint schema contains teardown owners, not lease + // generations. The first legacy migration needs its own atomic record. + writeRecordsAtomically( + [ + { + version: 2, + ...checkpointForState(state), + revision: Math.max(0, revision - 1) + }, + { version: 2, revision, ...mutation } + ], + dataFile + ) + return + } writeTerminalTeardownIntentSnapshot(state, dataFile) return }