diff --git a/compute/worker/src/orphan-recovery.ts b/compute/worker/src/orphan-recovery.ts new file mode 100644 index 0000000..882ea94 --- /dev/null +++ b/compute/worker/src/orphan-recovery.ts @@ -0,0 +1,100 @@ +import type { + PdfLayoutJobResult, + WhisperAlignJobResult, + WorkerJobTiming, + WorkerJobState, + WorkerOperationState, +} from '@openreader/compute-core/api-contracts'; + +export type StreamedOperationState = WorkerOperationState; + +export interface OrphanRecoveryStateStore { + getOpStateRecord(opId: string): Promise<{ state: StreamedOperationState; revision: number } | null>; + listOpStates(): Promise; +} + +export interface OrphanRecoveryOrchestrator { + markFailedIfUnchanged(input: { + current: StreamedOperationState; + expectedRevision: number; + error: { message: string; code?: string } | string; + updatedAt?: number; + timing?: WorkerJobTiming; + }): Promise; +} + +export interface RecoverOrphanedOperationsInput { + operationStateStore: OrphanRecoveryStateStore; + orchestrator: OrphanRecoveryOrchestrator; + whisperTimeoutMs: number; + pdfTimeoutMs: number; + opStaleMs: number; + nowMs?: number; +} + +function isInflightStatus(status: WorkerJobState): boolean { + return status === 'queued' || status === 'running'; +} + +export function getOrphanRecoveryThresholdMs(input: { + state: StreamedOperationState; + whisperTimeoutMs: number; + pdfTimeoutMs: number; + opStaleMs: number; +}): number | null { + if (!isInflightStatus(input.state.status)) return null; + if (input.state.status === 'running') { + return input.state.kind === 'whisper_align' ? input.whisperTimeoutMs : input.pdfTimeoutMs; + } + if (input.state.kind !== 'pdf_layout') return null; + return input.opStaleMs; +} + +export async function recoverOrphanedOperations( + input: RecoverOrphanedOperationsInput, +): Promise>> { + const nowMs = input.nowMs ?? Date.now(); + const states = await input.operationStateStore.listOpStates(); + const candidateStates = states.filter((state) => ( + getOrphanRecoveryThresholdMs({ + state, + whisperTimeoutMs: input.whisperTimeoutMs, + pdfTimeoutMs: input.pdfTimeoutMs, + opStaleMs: input.opStaleMs, + }) !== null + )); + const recoveredStates: Array> = []; + + for (const candidate of candidateStates) { + const record = await input.operationStateStore.getOpStateRecord(candidate.opId); + if (!record) continue; + const staleAfterMs = getOrphanRecoveryThresholdMs({ + state: record.state, + whisperTimeoutMs: input.whisperTimeoutMs, + pdfTimeoutMs: input.pdfTimeoutMs, + opStaleMs: input.opStaleMs, + }); + if (staleAfterMs === null) continue; + const ageMs = nowMs - record.state.updatedAt; + if (ageMs <= staleAfterMs) continue; + + const recovered = await input.orchestrator.markFailedIfUnchanged({ + current: record.state, + expectedRevision: record.revision, + error: { + code: 'WORKER_ORPHANED_OP', + message: `Worker stopped before completion; stale operation recovered during reconciliation after ${staleAfterMs}ms`, + }, + updatedAt: nowMs, + }); + if (!recovered) continue; + + recoveredStates.push({ + opId: recovered.opId, + kind: recovered.kind, + status: record.state.status, + }); + } + + return recoveredStates; +} diff --git a/compute/worker/src/runtime.ts b/compute/worker/src/runtime.ts index c2bf09d..e208e63 100644 --- a/compute/worker/src/runtime.ts +++ b/compute/worker/src/runtime.ts @@ -56,6 +56,10 @@ import { hashOpKey, } from './control-plane/jetstream'; import { type JsonCodec, createJsonCodec } from './control-plane/json-codec'; +import { + recoverOrphanedOperations, + type StreamedOperationState, +} from './orphan-recovery'; import { buildInferProgressForPageParsed, buildInferProgressForPageStart } from './pdf-progress'; import { buildQueueWaitTiming, decideRetryAction } from './worker-loop-policy'; @@ -79,6 +83,7 @@ const NATS_API_TIMEOUT_MS = 60_000; // put it to sleep. Reconnect happens lazily on the next inbound request. const IDLE_DISCONNECT_MS = 120_000; const IDLE_CHECK_INTERVAL_MS = 5_000; +const ORPHAN_SWEEP_INTERVAL_MS = 15_000; // Bounded pull window so consumer loops yield periodically and can be stopped // cleanly when going idle, instead of blocking on a long-lived pull. const PULL_EXPIRES_MS = 5_000; @@ -107,8 +112,6 @@ interface NatsSession { layoutConsumer: Consumer; } -type StreamedOperationState = WorkerOperationState; - interface OperationEventStreamLike { subscribe(input: { opId: string; @@ -187,24 +190,6 @@ function requireEnv(name: string): string { return value; } -function isInflightStatus(status: WorkerJobState): boolean { - return status === 'queued' || status === 'running'; -} - -function getOrphanRecoveryThresholdMs(input: { - state: StreamedOperationState; - whisperTimeoutMs: number; - pdfTimeoutMs: number; - opStaleMs: number; -}): number | null { - if (!isInflightStatus(input.state.status)) return null; - if (input.state.status === 'running') { - return input.state.kind === 'whisper_align' ? input.whisperTimeoutMs : input.pdfTimeoutMs; - } - if (input.state.kind !== 'pdf_layout') return null; - return input.opStaleMs; -} - function readIntEnv(name: string, fallback: number): number { const raw = process.env[name]?.trim(); if (!raw) return fallback; @@ -568,6 +553,7 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti let connecting: Promise | null = null; let workerLoops: Promise[] = []; let idleTimer: NodeJS.Timeout | null = null; + let orphanSweepTimer: NodeJS.Timeout | null = null; let stopping = false; let loopStopRequested = false; @@ -596,6 +582,19 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti idleTimer.unref?.(); } + function startOrphanSweepTimer(): void { + if (orphanSweepTimer) return; + orphanSweepTimer = setInterval(() => { + if (!session || stopping) return; + void runOrphanedOpRecovery({ force: true }).catch((error) => { + app.log.error({ + error: toErrorMessage(error), + }, 'periodic orphaned operation recovery failed'); + }); + }, ORPHAN_SWEEP_INTERVAL_MS); + orphanSweepTimer.unref?.(); + } + async function disconnect(reason: string): Promise { const current = session; if (!current) return; @@ -607,6 +606,10 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti clearInterval(idleTimer); idleTimer = null; } + if (orphanSweepTimer) { + clearInterval(orphanSweepTimer); + orphanSweepTimer = null; + } try { await current.nc.close(); } catch { @@ -649,6 +652,7 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti markActivity(); startWorkerLoops(next); startIdleTimer(); + startOrphanSweepTimer(); // Safety net: if the connection closes for any reason (network drop after // exhausting reconnects, or our own disconnect), drop the stale session so // the next request reconnects cleanly. @@ -764,59 +768,35 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti let orphanRecoveryDoneForGeneration = -1; let sessionGeneration = options.routeDeps ? 0 : -1; - const ensureOrphanedOpRecovery = async (): Promise => { + const runOrphanedOpRecovery = async (options?: { force?: boolean }): Promise => { if (typeof operationStateStore.listOpStates !== 'function') return; if (typeof operationStateStore.getOpStateRecord !== 'function') return; if (typeof orchestrator.markFailedIfUnchanged !== 'function') return; - if (orphanRecoveryDoneForGeneration === sessionGeneration) return; + if (!options?.force && orphanRecoveryDoneForGeneration === sessionGeneration) return; if (orphanRecoveryPromise) { await orphanRecoveryPromise; return; } orphanRecoveryPromise = (async () => { - const now = Date.now(); - const states = await operationStateStore.listOpStates!(); - const candidateStates = states.filter((state) => ( - getOrphanRecoveryThresholdMs({ - state, - whisperTimeoutMs, - pdfTimeoutMs, - opStaleMs, - }) !== null - )); - const recoveredStates: Array> = []; - - for (const candidate of candidateStates) { - const record = await operationStateStore.getOpStateRecord!(candidate.opId); - if (!record) continue; - const staleAfterMs = getOrphanRecoveryThresholdMs({ - state: record.state, - whisperTimeoutMs, - pdfTimeoutMs, - opStaleMs, - }); - if (staleAfterMs === null) continue; - const ageMs = now - record.state.updatedAt; - if (ageMs <= staleAfterMs) continue; - - const recovered = await orchestrator.markFailedIfUnchanged!({ - current: record.state, - expectedRevision: record.revision, - error: { - code: 'WORKER_ORPHANED_OP', - message: `Worker stopped before completion; stale operation recovered on startup after ${staleAfterMs}ms`, - }, - updatedAt: now, - }); - if (!recovered) continue; - - recoveredStates.push({ - opId: recovered.opId, - kind: recovered.kind, - status: record.state.status, - }); - } + const recoveredStates = await recoverOrphanedOperations({ + operationStateStore: operationStateStore as typeof operationStateStore & { + getOpStateRecord(opId: string): Promise<{ state: StreamedOperationState; revision: number } | null>; + listOpStates(): Promise; + }, + orchestrator: orchestrator as typeof orchestrator & { + markFailedIfUnchanged(input: { + current: StreamedOperationState; + expectedRevision: number; + error: { message: string; code?: string } | string; + updatedAt?: number; + timing?: WorkerJobTiming; + }): Promise; + }, + whisperTimeoutMs, + pdfTimeoutMs, + opStaleMs, + }); if (recoveredStates.length > 0) { app.log.warn({ @@ -826,7 +806,7 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti kind: state.kind, status: state.status, })), - }, 'recovered stale in-flight operations on startup'); + }, 'recovered stale in-flight operations during reconciliation'); } orphanRecoveryDoneForGeneration = sessionGeneration; @@ -837,6 +817,10 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti await orphanRecoveryPromise; }; + const ensureOrphanedOpRecovery = async (): Promise => { + await runOrphanedOpRecovery(); + }; + const getOpState = async (opId: string): Promise => { await ensureOrphanedOpRecovery(); return await operationStateStore.getOpState(opId); @@ -1509,6 +1493,10 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti clearInterval(idleTimer); idleTimer = null; } + if (orphanSweepTimer) { + clearInterval(orphanSweepTimer); + orphanSweepTimer = null; + } await app.close(); await Promise.allSettled(workerLoops); const current = session; diff --git a/compute/worker/tests/unit/orphan-recovery.test.ts b/compute/worker/tests/unit/orphan-recovery.test.ts new file mode 100644 index 0000000..712fdf6 --- /dev/null +++ b/compute/worker/tests/unit/orphan-recovery.test.ts @@ -0,0 +1,59 @@ +import { afterEach, describe, expect, test, vi } from 'vitest'; +import { recoverOrphanedOperations } from '../../src/orphan-recovery'; +import { FakeControlPlane } from '../fixtures/fake-control-plane'; + +describe('orphan recovery', () => { + afterEach(() => { + vi.useRealTimers(); + }); + + test('recovers a running whisper op when a later sweep crosses the timeout in the same session', async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-06-03T08:00:00.000Z')); + + const fake = new FakeControlPlane(); + const startedAt = Date.now(); + fake.seedState({ + opId: 'op-whisper-running', + opKey: 'k-whisper-running', + kind: 'whisper_align', + jobId: 'job-op-whisper-running', + status: 'running', + queuedAt: startedAt, + updatedAt: startedAt, + }); + + const firstSweep = await recoverOrphanedOperations({ + operationStateStore: fake.deps.operationStateStore, + orchestrator: fake.deps.orchestrator, + whisperTimeoutMs: 30_000, + pdfTimeoutMs: 300_000, + opStaleMs: 1_800_000, + }); + expect(firstSweep).toEqual([]); + expect(fake.getState('op-whisper-running')).toMatchObject({ + status: 'running', + }); + + vi.advanceTimersByTime(31_000); + + const secondSweep = await recoverOrphanedOperations({ + operationStateStore: fake.deps.operationStateStore, + orchestrator: fake.deps.orchestrator, + whisperTimeoutMs: 30_000, + pdfTimeoutMs: 300_000, + opStaleMs: 1_800_000, + }); + expect(secondSweep).toEqual([{ + opId: 'op-whisper-running', + kind: 'whisper_align', + status: 'running', + }]); + expect(fake.getState('op-whisper-running')).toMatchObject({ + status: 'failed', + error: { + code: 'WORKER_ORPHANED_OP', + }, + }); + }); +});