diff --git a/compute/core/src/control-plane/index.ts b/compute/core/src/control-plane/index.ts index 8e6b18c..5b8255d 100644 --- a/compute/core/src/control-plane/index.ts +++ b/compute/core/src/control-plane/index.ts @@ -1,5 +1,4 @@ export * from './types'; export * from './state-machine'; export * from './orchestrator'; -export * from './in-memory'; export * from './sse'; diff --git a/compute/core/src/control-plane/orchestrator.ts b/compute/core/src/control-plane/orchestrator.ts index 950fe71..7f621a4 100644 --- a/compute/core/src/control-plane/orchestrator.ts +++ b/compute/core/src/control-plane/orchestrator.ts @@ -263,6 +263,36 @@ export class OperationOrchestrator { return next; } + async markFailedIfUnchanged(input: { + current: WorkerOperationState; + expectedRevision: number; + error: WorkerJobErrorShape | string; + updatedAt?: number; + timing?: WorkerJobTiming; + }): Promise { + const now = input.updatedAt ?? this.clock.now(); + const shape = typeof input.error === 'string' ? { message: input.error } : input.error; + + const next: WorkerOperationState = { + ...input.current, + status: 'failed', + startedAt: input.current.startedAt ?? now, + updatedAt: now, + error: shape, + ...(input.timing ? { timing: input.timing } : {}), + }; + + const updated = await this.stateStore.compareAndSetOpState({ + opId: input.current.opId, + expectedRevision: input.expectedRevision, + newState: next, + }); + if (!updated) return null; + + await this.eventStream.append(next.opId, next); + return next; + } + async explainReuseDecision(input: { current: WorkerOperationState; requestKind: WorkerOperationKind; diff --git a/compute/core/src/control-plane/types.ts b/compute/core/src/control-plane/types.ts index 79fd16c..9fc66ff 100644 --- a/compute/core/src/control-plane/types.ts +++ b/compute/core/src/control-plane/types.ts @@ -26,9 +26,20 @@ export interface OperationIndexEntry { opId: string; } +export interface OperationStateRecord { + state: OperationState; + revision: number; +} + export interface OperationStateStore { getOpState(opId: string): Promise | null>; + getOpStateRecord(opId: string): Promise | null>; putOpState(state: OperationState): Promise; + compareAndSetOpState(input: { + opId: string; + expectedRevision: number; + newState: OperationState; + }): Promise; getOpIndex(opKey: string): Promise; compareAndSetOpIndex(input: { opKey: string; diff --git a/compute/core/tests/control-plane/orchestrator.test.ts b/compute/core/tests/control-plane/orchestrator.test.ts index 2a03adf..80bd9fb 100644 --- a/compute/core/tests/control-plane/orchestrator.test.ts +++ b/compute/core/tests/control-plane/orchestrator.test.ts @@ -1,11 +1,11 @@ import { describe, expect, test } from 'vitest'; import type { WorkerOperationRequest } from '../../src/api-contracts'; +import { OperationOrchestrator } from '../../src/control-plane'; import { InMemoryOperationEventStream, InMemoryOperationQueue, InMemoryOperationStateStore, - OperationOrchestrator, -} from '../../src/control-plane'; +} from '../helpers/in-memory-control-plane'; function buildRequest(opKey: string): WorkerOperationRequest { return { @@ -65,7 +65,9 @@ describe('operation orchestrator', () => { let firstAttempt = true; const conflictStore = { getOpState: store.getOpState.bind(store), + getOpStateRecord: store.getOpStateRecord.bind(store), putOpState: store.putOpState.bind(store), + compareAndSetOpState: store.compareAndSetOpState.bind(store), getOpIndex: store.getOpIndex.bind(store), compareAndSetOpIndex: async (input: { opKey: string; newOpId: string; expectedOpId: string | null }) => { if (firstAttempt && input.expectedOpId === null) { @@ -92,4 +94,45 @@ describe('operation orchestrator', () => { expect(created.opId).toMatch(/^op-/); expect(await store.getOpIndex('cas-key')).toEqual({ opId: created.opId }); }); + + test('markFailedIfUnchanged only writes once for the expected revision', async () => { + const queue = new InMemoryOperationQueue(); + const stateStore = new InMemoryOperationStateStore(); + const eventStream = new InMemoryOperationEventStream(); + const orchestrator = new OperationOrchestrator({ + queue, + stateStore, + eventStream, + config: { opStaleMs: 2_000, maxCasRetries: 5 }, + }); + + const created = await orchestrator.enqueueOrReuse(buildRequest('stale-op')); + await orchestrator.markRunning({ opId: created.opId, updatedAt: 2_000 }); + + const record = await stateStore.getOpStateRecord(created.opId); + expect(record).not.toBeNull(); + + const first = await orchestrator.markFailedIfUnchanged({ + current: record!.state, + expectedRevision: record!.revision, + error: { code: 'WORKER_ORPHANED_OP', message: 'stale op' }, + updatedAt: 3_000, + }); + const second = await orchestrator.markFailedIfUnchanged({ + current: record!.state, + expectedRevision: record!.revision, + error: { code: 'WORKER_ORPHANED_OP', message: 'stale op' }, + updatedAt: 3_000, + }); + + expect(first).toMatchObject({ + opId: created.opId, + status: 'failed', + error: { code: 'WORKER_ORPHANED_OP' }, + }); + expect(second).toBeNull(); + + const events = await eventStream.listSince(created.opId, 0); + expect(events.filter((event) => event.snapshot.status === 'failed')).toHaveLength(1); + }); }); diff --git a/compute/core/src/control-plane/in-memory.ts b/compute/core/tests/helpers/in-memory-control-plane.ts similarity index 80% rename from compute/core/src/control-plane/in-memory.ts rename to compute/core/tests/helpers/in-memory-control-plane.ts index fc77c0e..cc9b9ce 100644 --- a/compute/core/src/control-plane/in-memory.ts +++ b/compute/core/tests/helpers/in-memory-control-plane.ts @@ -1,4 +1,5 @@ import { EventEmitter } from 'node:events'; +import type { WorkerOperationKind } from '../../src/api-contracts'; import type { OperationEvent, OperationEventStream, @@ -7,8 +8,7 @@ import type { OperationState, OperationStateStore, QueuedOperation, -} from './types'; -import type { WorkerOperationKind } from '../api-contracts'; +} from '../../src/control-plane/types'; function topicFor(opId: string): string { return `op.${opId}`; @@ -49,14 +49,39 @@ export class InMemoryOperationQueue implements OperationQueue { export class InMemoryOperationStateStore implements OperationStateStore { private readonly stateByOpId = new Map(); + private readonly revisionByOpId = new Map(); private readonly opIndexByKey = new Map(); async getOpState(opId: string): Promise { return this.stateByOpId.get(opId) ?? null; } + async getOpStateRecord(opId: string): Promise<{ state: OperationState; revision: number } | null> { + const state = this.stateByOpId.get(opId); + if (!state) return null; + return { + state, + revision: this.revisionByOpId.get(opId) ?? 0, + }; + } + async putOpState(state: OperationState): Promise { this.stateByOpId.set(state.opId, state); + this.revisionByOpId.set(state.opId, (this.revisionByOpId.get(state.opId) ?? 0) + 1); + } + + async compareAndSetOpState(input: { + opId: string; + expectedRevision: number; + newState: OperationState; + }): Promise { + const currentState = this.stateByOpId.get(input.opId); + if (!currentState) return false; + const currentRevision = this.revisionByOpId.get(input.opId) ?? 0; + if (currentRevision !== input.expectedRevision) return false; + this.stateByOpId.set(input.opId, input.newState); + this.revisionByOpId.set(input.opId, currentRevision + 1); + return true; } async getOpIndex(opKey: string): Promise { diff --git a/compute/worker/src/control-plane/jetstream.ts b/compute/worker/src/control-plane/jetstream.ts index 8844673..9225d96 100644 --- a/compute/worker/src/control-plane/jetstream.ts +++ b/compute/worker/src/control-plane/jetstream.ts @@ -81,10 +81,18 @@ export class JetStreamOperationStateStore implements Operation } async getOpState(opId: string): Promise | null> { + const record = await this.getOpStateRecord(opId); + return record?.state ?? null; + } + + async getOpStateRecord(opId: string): Promise<{ state: OperationState; revision: number } | null> { const kv = await this.getKv(); const entry = await kv.get(opStateKvKey(opId)); if (!isPut(entry)) return null; - return this.opStateCodec.decode(entry.value); + return { + state: this.opStateCodec.decode(entry.value), + revision: entry.revision, + }; } async putOpState(state: OperationState): Promise { @@ -92,6 +100,25 @@ export class JetStreamOperationStateStore implements Operation await kv.put(opStateKvKey(state.opId), this.opStateCodec.encode(state)); } + async compareAndSetOpState(input: { + opId: string; + expectedRevision: number; + newState: OperationState; + }): Promise { + const kv = await this.getKv(); + try { + await kv.update( + opStateKvKey(input.opId), + this.opStateCodec.encode(input.newState), + input.expectedRevision, + ); + return true; + } catch (error) { + if (isCasConflictError(error)) return false; + throw error; + } + } + async listOpStates(): Promise[]> { const kv = await this.getKv(); const keys = await kv.keys('op_state.*'); diff --git a/compute/worker/src/runtime.ts b/compute/worker/src/runtime.ts index fb89764..c2bf09d 100644 --- a/compute/worker/src/runtime.ts +++ b/compute/worker/src/runtime.ts @@ -120,6 +120,7 @@ interface OperationEventStreamLike { interface OperationStateStoreLike { getOpState(opId: string): Promise; + getOpStateRecord?(opId: string): Promise<{ state: StreamedOperationState; revision: number } | null>; listOpStates?(): Promise; } @@ -149,6 +150,13 @@ interface OrchestratorLike { updatedAt?: number; timing?: WorkerJobTiming; }): Promise; + markFailedIfUnchanged?(input: { + current: StreamedOperationState; + expectedRevision: number; + error: { message: string; code?: string } | string; + updatedAt?: number; + timing?: WorkerJobTiming; + }): Promise; } export interface ComputeWorkerRouteDeps { @@ -183,6 +191,20 @@ 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; @@ -744,6 +766,8 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti const ensureOrphanedOpRecovery = async (): Promise => { if (typeof operationStateStore.listOpStates !== 'function') return; + if (typeof operationStateStore.getOpStateRecord !== 'function') return; + if (typeof orchestrator.markFailedIfUnchanged !== 'function') return; if (orphanRecoveryDoneForGeneration === sessionGeneration) return; if (orphanRecoveryPromise) { await orphanRecoveryPromise; @@ -753,35 +777,51 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti orphanRecoveryPromise = (async () => { const now = Date.now(); const states = await operationStateStore.listOpStates!(); - const staleStates = states.filter((state) => { - if (!isInflightStatus(state.status)) return false; - const ageMs = now - state.updatedAt; - if (state.status === 'running') { - const runningTimeoutMs = state.kind === 'whisper_align' ? whisperTimeoutMs : pdfTimeoutMs; - return ageMs > runningTimeoutMs; - } - if (state.kind !== 'pdf_layout') return false; - return ageMs > opStaleMs; - }); + const candidateStates = states.filter((state) => ( + getOrphanRecoveryThresholdMs({ + state, + whisperTimeoutMs, + pdfTimeoutMs, + opStaleMs, + }) !== null + )); + const recoveredStates: Array> = []; - for (const state of staleStates) { - const staleAfterMs = state.status === 'running' - ? (state.kind === 'whisper_align' ? whisperTimeoutMs : pdfTimeoutMs) - : opStaleMs; - await orchestrator.markFailed({ - opId: state.opId, + 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, + }); } - if (staleStates.length > 0) { + if (recoveredStates.length > 0) { app.log.warn({ - recoveredCount: staleStates.length, - ops: staleStates.map((state) => ({ + recoveredCount: recoveredStates.length, + ops: recoveredStates.map((state) => ({ opId: state.opId, kind: state.kind, status: state.status, diff --git a/compute/worker/tests/fixtures/fake-control-plane.ts b/compute/worker/tests/fixtures/fake-control-plane.ts index 305304e..9822523 100644 --- a/compute/worker/tests/fixtures/fake-control-plane.ts +++ b/compute/worker/tests/fixtures/fake-control-plane.ts @@ -13,6 +13,7 @@ type ComputeEvent = WorkerOperationEvent; export class FakeControlPlane { private readonly stateByOpId = new Map(); + private readonly revisionByOpId = new Map(); private readonly opIdByOpKey = new Map(); private readonly eventsByOpId = new Map(); private nextOpId = 1; @@ -44,9 +45,18 @@ export class FakeControlPlane { updatedAt: input.updatedAt, timing: input.timing, }), + markFailedIfUnchanged: async (input) => this.compareAndSetFailed(input), }, operationStateStore: { getOpState: async (opId) => this.stateByOpId.get(opId) ?? null, + getOpStateRecord: async (opId) => { + const state = this.stateByOpId.get(opId); + if (!state) return null; + return { + state, + revision: this.revisionByOpId.get(opId) ?? 0, + }; + }, listOpStates: async () => Array.from(this.stateByOpId.values()), }, operationEventStream: { @@ -63,6 +73,7 @@ export class FakeControlPlane { seedState(state: ComputeState): void { this.stateByOpId.set(state.opId, state); + this.revisionByOpId.set(state.opId, (this.revisionByOpId.get(state.opId) ?? 0) + 1); this.opIdByOpKey.set(state.opKey, state.opId); } @@ -96,11 +107,30 @@ export class FakeControlPlane { }; this.stateByOpId.set(opId, state); + this.revisionByOpId.set(opId, 1); this.opIdByOpKey.set(request.opKey, opId); this.seedEvent(opId, { eventId: 1, snapshot: state }); return state; } + private async compareAndSetFailed(input: { + current: ComputeState; + expectedRevision: number; + error: { message: string; code?: string } | string; + updatedAt?: number; + timing?: ComputeState['timing']; + }): Promise { + const currentRevision = this.revisionByOpId.get(input.current.opId) ?? 0; + if (currentRevision !== input.expectedRevision) return null; + return this.updateState(input.current.opId, { + ...input.current, + status: 'failed', + error: typeof input.error === 'string' ? { message: input.error } : input.error, + updatedAt: input.updatedAt, + timing: input.timing, + }); + } + private async updateState( opId: string, patch: Partial, @@ -115,6 +145,7 @@ export class FakeControlPlane { updatedAt: patch.updatedAt ?? Date.now(), }; this.stateByOpId.set(opId, next); + this.revisionByOpId.set(opId, (this.revisionByOpId.get(opId) ?? 0) + 1); const currentEvents = this.eventsByOpId.get(opId) ?? []; const nextEventId = (currentEvents.at(-1)?.eventId ?? 0) + 1; currentEvents.push({ eventId: nextEventId, snapshot: next });