From 21a8a71b038b20441469e559c5b5d96becec1b2e Mon Sep 17 00:00:00 2001 From: Richard R Date: Wed, 3 Jun 2026 02:15:29 -0600 Subject: [PATCH] refactor(control-plane): introduce revision-based CAS for operation state updates Add revision tracking and compare-and-set (CAS) semantics to operation state stores, enabling atomic state transitions and preventing lost updates. Extend the OperationStateStore interface with getOpStateRecord and compareAndSetOpState methods. Update orchestrator and worker runtime to utilize CAS for marking operations as failed only if the state is unchanged. Enhance in-memory, JetStream, and test control plane implementations to support revision logic. This change improves concurrency safety and correctness of operation state management across distributed components. --- compute/core/src/control-plane/index.ts | 1 - .../core/src/control-plane/orchestrator.ts | 30 +++++++ compute/core/src/control-plane/types.ts | 11 +++ .../tests/control-plane/orchestrator.test.ts | 47 ++++++++++- .../helpers/in-memory-control-plane.ts} | 29 ++++++- compute/worker/src/control-plane/jetstream.ts | 29 ++++++- compute/worker/src/runtime.ts | 78 ++++++++++++++----- .../tests/fixtures/fake-control-plane.ts | 31 ++++++++ 8 files changed, 231 insertions(+), 25 deletions(-) rename compute/core/{src/control-plane/in-memory.ts => tests/helpers/in-memory-control-plane.ts} (80%) 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 });