From 74a05a8de0335988bf1fd59470393f48f8a99952 Mon Sep 17 00:00:00 2001 From: Richard R Date: Wed, 3 Jun 2026 02:49:08 -0600 Subject: [PATCH] feat(worker): extract orphaned operation recovery to module with periodic sweep Move orphaned operation recovery logic into a dedicated orphan-recovery module, introducing a periodic sweep timer that triggers recovery every 15 seconds while the worker is connected. Refactor runtime to delegate orphan detection and handling to the new module, improving modularity and maintainability. Add unit tests for orphan-recovery to ensure correctness. --- compute/worker/src/orphan-recovery.ts | 100 +++++++++++++++ compute/worker/src/runtime.ts | 118 ++++++++---------- .../worker/tests/unit/orphan-recovery.test.ts | 59 +++++++++ 3 files changed, 212 insertions(+), 65 deletions(-) create mode 100644 compute/worker/src/orphan-recovery.ts create mode 100644 compute/worker/tests/unit/orphan-recovery.test.ts 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', + }, + }); + }); +});