refactor(worker-control-plane): route ops through core orchestrator
This commit is contained in:
Richard R
parent
00fc2d5e36
commit
084cbdfac1
3 changed files with 484 additions and 289 deletions
229
compute/worker/src/control-plane-jetstream.ts
Normal file
229
compute/worker/src/control-plane-jetstream.ts
Normal file
|
|
@ -0,0 +1,229 @@
|
|||
import { createHash } from 'node:crypto';
|
||||
import type { JetStreamClient } from '@nats-io/jetstream';
|
||||
import type {
|
||||
OperationEvent,
|
||||
OperationEventStream,
|
||||
OperationQueue,
|
||||
OperationState,
|
||||
OperationStateStore,
|
||||
QueuedOperation,
|
||||
} from '@openreader/compute-core/control-plane';
|
||||
import type {
|
||||
PdfLayoutJobRequest,
|
||||
WhisperAlignJobRequest,
|
||||
WorkerOperationKind,
|
||||
} from '@openreader/compute-core/api-contracts';
|
||||
|
||||
export interface KvEntryLike {
|
||||
operation?: string;
|
||||
value: Uint8Array;
|
||||
revision: number;
|
||||
}
|
||||
|
||||
export interface KvStoreLike {
|
||||
get(key: string): Promise<KvEntryLike | null>;
|
||||
put(key: string, data: Uint8Array): Promise<unknown>;
|
||||
create(key: string, data: Uint8Array): Promise<unknown>;
|
||||
update(key: string, data: Uint8Array, version: number): Promise<unknown>;
|
||||
}
|
||||
|
||||
type JsonCodec<T> = {
|
||||
encode(value: T): Uint8Array;
|
||||
decode(data: Uint8Array): T;
|
||||
};
|
||||
|
||||
function createJsonCodec<T>(): JsonCodec<T> {
|
||||
const encoder = new TextEncoder();
|
||||
const decoder = new TextDecoder();
|
||||
return {
|
||||
encode(value: T): Uint8Array {
|
||||
return encoder.encode(JSON.stringify(value));
|
||||
},
|
||||
decode(data: Uint8Array): T {
|
||||
return JSON.parse(decoder.decode(data)) as T;
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function toErrorMessage(error: unknown): string {
|
||||
if (error instanceof Error && error.message) return error.message;
|
||||
return String(error);
|
||||
}
|
||||
|
||||
function isCasConflictError(error: unknown): boolean {
|
||||
const message = toErrorMessage(error).toLowerCase();
|
||||
return message.includes('wrong last sequence') || message.includes('key exists') || message.includes('wrong last');
|
||||
}
|
||||
|
||||
function isPut(entry: KvEntryLike | null): entry is KvEntryLike {
|
||||
return Boolean(entry && entry.operation === 'PUT');
|
||||
}
|
||||
|
||||
interface OpIndexEntry {
|
||||
opId: string;
|
||||
}
|
||||
|
||||
const OP_EVENTS_SUBJECT_PREFIX = 'ops.events';
|
||||
|
||||
export function hashOpKey(opKey: string): string {
|
||||
return createHash('sha256').update(opKey).digest('hex');
|
||||
}
|
||||
|
||||
export function opIndexKvKey(opKey: string): string {
|
||||
return `op_index.${hashOpKey(opKey)}`;
|
||||
}
|
||||
|
||||
export function opStateKvKey(opId: string): string {
|
||||
return `op_state.${opId}`;
|
||||
}
|
||||
|
||||
export function opEventsSubject(opId: string): string {
|
||||
return `${OP_EVENTS_SUBJECT_PREFIX}.${opId}`;
|
||||
}
|
||||
|
||||
export interface JetStreamOperationStateStoreDeps {
|
||||
getKv: () => Promise<KvStoreLike>;
|
||||
}
|
||||
|
||||
export class JetStreamOperationStateStore<Result = unknown> implements OperationStateStore<Result> {
|
||||
private readonly getKv: () => Promise<KvStoreLike>;
|
||||
private readonly opStateCodec = createJsonCodec<OperationState<Result>>();
|
||||
private readonly opIndexCodec = createJsonCodec<OpIndexEntry>();
|
||||
|
||||
constructor(deps: JetStreamOperationStateStoreDeps) {
|
||||
this.getKv = deps.getKv;
|
||||
}
|
||||
|
||||
async getOpState(opId: string): Promise<OperationState<Result> | null> {
|
||||
const kv = await this.getKv();
|
||||
const entry = await kv.get(opStateKvKey(opId));
|
||||
if (!isPut(entry)) return null;
|
||||
return this.opStateCodec.decode(entry.value);
|
||||
}
|
||||
|
||||
async putOpState(state: OperationState<Result>): Promise<void> {
|
||||
const kv = await this.getKv();
|
||||
await kv.put(opStateKvKey(state.opId), this.opStateCodec.encode(state));
|
||||
}
|
||||
|
||||
async getOpIndex(opKey: string): Promise<{ opId: string } | null> {
|
||||
const kv = await this.getKv();
|
||||
const entry = await kv.get(opIndexKvKey(opKey));
|
||||
if (!isPut(entry)) return null;
|
||||
return this.opIndexCodec.decode(entry.value);
|
||||
}
|
||||
|
||||
async compareAndSetOpIndex(input: {
|
||||
opKey: string;
|
||||
newOpId: string;
|
||||
expectedOpId: string | null;
|
||||
}): Promise<boolean> {
|
||||
const kv = await this.getKv();
|
||||
const key = opIndexKvKey(input.opKey);
|
||||
const value = this.opIndexCodec.encode({ opId: input.newOpId });
|
||||
|
||||
if (input.expectedOpId === null) {
|
||||
try {
|
||||
await kv.create(key, value);
|
||||
return true;
|
||||
} catch (error) {
|
||||
if (isCasConflictError(error)) return false;
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
const current = await kv.get(key);
|
||||
if (!isPut(current)) return false;
|
||||
const decoded = this.opIndexCodec.decode(current.value);
|
||||
if (decoded.opId !== input.expectedOpId) return false;
|
||||
|
||||
try {
|
||||
await kv.update(key, value, current.revision);
|
||||
return true;
|
||||
} catch (error) {
|
||||
if (isCasConflictError(error)) return false;
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export interface JetStreamOperationEventStreamDeps {
|
||||
getJs: () => Promise<Pick<JetStreamClient, 'publish'>>;
|
||||
}
|
||||
|
||||
export class JetStreamOperationEventStream<Result = unknown> implements OperationEventStream<Result> {
|
||||
private readonly getJs: () => Promise<Pick<JetStreamClient, 'publish'>>;
|
||||
|
||||
constructor(deps: JetStreamOperationEventStreamDeps) {
|
||||
this.getJs = deps.getJs;
|
||||
}
|
||||
|
||||
async append(opId: string, snapshot: OperationState<Result>): Promise<OperationEvent<Result>> {
|
||||
const js = await this.getJs();
|
||||
const encoder = new TextEncoder();
|
||||
const ack = await js.publish(opEventsSubject(opId), encoder.encode(JSON.stringify(snapshot)));
|
||||
return {
|
||||
eventId: ack.seq,
|
||||
snapshot,
|
||||
};
|
||||
}
|
||||
|
||||
async listSince(): Promise<OperationEvent<Result>[]> {
|
||||
return [];
|
||||
}
|
||||
|
||||
async subscribe(): Promise<() => void> {
|
||||
return () => undefined;
|
||||
}
|
||||
}
|
||||
|
||||
export interface JetStreamOperationQueueDeps<TPayload> {
|
||||
getJs: () => Promise<Pick<JetStreamClient, 'publish'>>;
|
||||
whisperSubject: string;
|
||||
layoutSubject: string;
|
||||
onEnqueued?: (job: QueuedOperation<TPayload>) => Promise<void> | void;
|
||||
}
|
||||
|
||||
export class JetStreamOperationQueue implements OperationQueue<WhisperAlignJobRequest | PdfLayoutJobRequest> {
|
||||
private readonly getJs: () => Promise<Pick<JetStreamClient, 'publish'>>;
|
||||
private readonly whisperSubject: string;
|
||||
private readonly layoutSubject: string;
|
||||
private readonly onEnqueued?: (job: QueuedOperation<WhisperAlignJobRequest | PdfLayoutJobRequest>) => Promise<void> | void;
|
||||
private readonly whisperCodec = createJsonCodec<QueuedOperation<WhisperAlignJobRequest>>();
|
||||
private readonly layoutCodec = createJsonCodec<QueuedOperation<PdfLayoutJobRequest>>();
|
||||
|
||||
constructor(deps: JetStreamOperationQueueDeps<WhisperAlignJobRequest | PdfLayoutJobRequest>) {
|
||||
this.getJs = deps.getJs;
|
||||
this.whisperSubject = deps.whisperSubject;
|
||||
this.layoutSubject = deps.layoutSubject;
|
||||
this.onEnqueued = deps.onEnqueued;
|
||||
}
|
||||
|
||||
async enqueue(job: QueuedOperation<WhisperAlignJobRequest | PdfLayoutJobRequest>): Promise<void> {
|
||||
const js = await this.getJs();
|
||||
if (job.kind === 'whisper_align') {
|
||||
await js.publish(
|
||||
this.whisperSubject,
|
||||
this.whisperCodec.encode(job as QueuedOperation<WhisperAlignJobRequest>),
|
||||
);
|
||||
} else if (job.kind === 'pdf_layout') {
|
||||
await js.publish(
|
||||
this.layoutSubject,
|
||||
this.layoutCodec.encode(job as QueuedOperation<PdfLayoutJobRequest>),
|
||||
);
|
||||
} else {
|
||||
const exhaustive: never = job.kind;
|
||||
throw new Error(`Unsupported operation kind: ${String(exhaustive)}`);
|
||||
}
|
||||
|
||||
await this.onEnqueued?.(job);
|
||||
}
|
||||
|
||||
async claimNext(_kind: WorkerOperationKind): Promise<QueuedOperation<WhisperAlignJobRequest | PdfLayoutJobRequest> | null> {
|
||||
throw new Error('JetStreamOperationQueue.claimNext is not used by the worker runtime');
|
||||
}
|
||||
|
||||
size(): number {
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
|
|
@ -1,4 +1,3 @@
|
|||
import { createHash } from 'node:crypto';
|
||||
import Fastify, { type FastifyRequest } from 'fastify';
|
||||
import { z } from 'zod';
|
||||
import {
|
||||
|
|
@ -34,6 +33,7 @@ import {
|
|||
withIdleTimeoutAndHardCap,
|
||||
withTimeout,
|
||||
} from '@openreader/compute-core';
|
||||
import { OperationOrchestrator } from '@openreader/compute-core/control-plane';
|
||||
import type {
|
||||
PdfLayoutJobRequest,
|
||||
PdfLayoutJobResult,
|
||||
|
|
@ -49,6 +49,13 @@ import type {
|
|||
PdfLayoutProgress,
|
||||
} from '@openreader/compute-core/api-contracts';
|
||||
import { GetObjectCommand, PutObjectCommand, S3Client } from '@aws-sdk/client-s3';
|
||||
import {
|
||||
JetStreamOperationEventStream,
|
||||
JetStreamOperationQueue,
|
||||
JetStreamOperationStateStore,
|
||||
hashOpKey,
|
||||
opEventsSubject,
|
||||
} from './control-plane-jetstream';
|
||||
|
||||
const JOBS_STREAM_NAME = 'compute_jobs';
|
||||
const WHISPER_JOBS_SUBJECT = 'jobs.whisper';
|
||||
|
|
@ -104,10 +111,6 @@ interface StoredJobState<Result> {
|
|||
progress?: PdfLayoutProgress;
|
||||
}
|
||||
|
||||
interface OpIndexEntry {
|
||||
opId: string;
|
||||
}
|
||||
|
||||
interface NatsSession {
|
||||
nc: NatsConnection;
|
||||
js: JetStreamClient;
|
||||
|
|
@ -295,11 +298,6 @@ function isAlreadyExistsError(error: unknown): boolean {
|
|||
return message.includes('already in use') || message.includes('already exists');
|
||||
}
|
||||
|
||||
function isCasConflictError(error: unknown): boolean {
|
||||
const message = toErrorMessage(error).toLowerCase();
|
||||
return message.includes('wrong last sequence') || message.includes('key exists') || message.includes('wrong last');
|
||||
}
|
||||
|
||||
function createJsonCodec<T>(): JsonCodec<T> {
|
||||
const encoder = new TextEncoder();
|
||||
const decoder = new TextDecoder();
|
||||
|
|
@ -351,26 +349,10 @@ function isTerminalStatus(status: WorkerJobState): boolean {
|
|||
return status === 'succeeded' || status === 'failed';
|
||||
}
|
||||
|
||||
function hashOpKey(opKey: string): string {
|
||||
return createHash('sha256').update(opKey).digest('hex');
|
||||
}
|
||||
|
||||
function opIndexKvKey(opKey: string): string {
|
||||
return `op_index.${hashOpKey(opKey)}`;
|
||||
}
|
||||
|
||||
function opStateKvKey(opId: string): string {
|
||||
return `op_state.${opId}`;
|
||||
}
|
||||
|
||||
function jobStateKvKey(jobId: string): string {
|
||||
return `job_state.${jobId}`;
|
||||
}
|
||||
|
||||
function opEventsSubject(opId: string): string {
|
||||
return `${OP_EVENTS_SUBJECT_PREFIX}.${opId}`;
|
||||
}
|
||||
|
||||
function extractResultRef(kind: WorkerOperationKind, result: unknown): string | undefined {
|
||||
if (kind !== 'pdf_layout' || !result || typeof result !== 'object') return undefined;
|
||||
const maybe = result as { parsedObjectKey?: unknown };
|
||||
|
|
@ -682,18 +664,55 @@ async function main(): Promise<void> {
|
|||
pdfLayoutHardCapMs: pdfHardCapMs,
|
||||
}, 'compute runtime config');
|
||||
|
||||
const opIndexCodec = createJsonCodec<OpIndexEntry>();
|
||||
const opStateCodec = createJsonCodec<StreamedOperationState>();
|
||||
const whisperJobCodec = createJsonCodec<QueuedJob<WhisperAlignJobRequest>>();
|
||||
const layoutJobCodec = createJsonCodec<QueuedJob<PdfLayoutJobRequest>>();
|
||||
const jobStateCodec = createJsonCodec<StoredJobState<WhisperAlignJobResult | PdfLayoutJobResult>>();
|
||||
const opEventCodec = createJsonCodec<StreamedOperationState>();
|
||||
|
||||
const putJobState = async (state: StoredJobState<WhisperAlignJobResult | PdfLayoutJobResult>): Promise<void> => {
|
||||
const { kv } = await ensureConnected();
|
||||
await kv.put(jobStateKvKey(state.jobId), jobStateCodec.encode(state));
|
||||
};
|
||||
|
||||
const operationStateStore = new JetStreamOperationStateStore<WhisperAlignJobResult | PdfLayoutJobResult>({
|
||||
getKv: async () => (await ensureConnected()).kv,
|
||||
});
|
||||
|
||||
const operationEventStream = new JetStreamOperationEventStream<WhisperAlignJobResult | PdfLayoutJobResult>({
|
||||
getJs: async () => (await ensureConnected()).js,
|
||||
});
|
||||
|
||||
const operationQueue = new JetStreamOperationQueue({
|
||||
getJs: async () => (await ensureConnected()).js,
|
||||
whisperSubject: WHISPER_JOBS_SUBJECT,
|
||||
layoutSubject: LAYOUT_JOBS_SUBJECT,
|
||||
onEnqueued: async (job) => {
|
||||
await putJobState({
|
||||
jobId: job.jobId,
|
||||
opId: job.opId,
|
||||
opKey: job.opKey,
|
||||
kind: job.kind,
|
||||
status: 'queued',
|
||||
timestamp: job.queuedAt,
|
||||
updatedAt: job.queuedAt,
|
||||
});
|
||||
},
|
||||
});
|
||||
|
||||
const orchestrator = new OperationOrchestrator({
|
||||
queue: operationQueue,
|
||||
stateStore: operationStateStore,
|
||||
eventStream: operationEventStream,
|
||||
config: {
|
||||
opStaleMs,
|
||||
maxCasRetries: 10,
|
||||
},
|
||||
});
|
||||
|
||||
const putOpState = async (state: StreamedOperationState): Promise<void> => {
|
||||
const { kv, js } = await ensureConnected();
|
||||
await kv.put(opStateKvKey(state.opId), opStateCodec.encode(state));
|
||||
await operationStateStore.putOpState(state);
|
||||
try {
|
||||
await js.publish(opEventsSubject(state.opId), opEventCodec.encode(state));
|
||||
await operationEventStream.append(state.opId, state);
|
||||
} catch (error) {
|
||||
app.log.warn({
|
||||
opId: state.opId,
|
||||
|
|
@ -704,263 +723,7 @@ async function main(): Promise<void> {
|
|||
};
|
||||
|
||||
const getOpState = async (opId: string): Promise<StreamedOperationState | null> => {
|
||||
const { kv } = await ensureConnected();
|
||||
const entry = await kv.get(opStateKvKey(opId));
|
||||
if (!entry || entry.operation !== 'PUT') return null;
|
||||
return opStateCodec.decode(entry.value);
|
||||
};
|
||||
|
||||
const putJobState = async (state: StoredJobState<WhisperAlignJobResult | PdfLayoutJobResult>): Promise<void> => {
|
||||
const { kv } = await ensureConnected();
|
||||
await kv.put(jobStateKvKey(state.jobId), jobStateCodec.encode(state));
|
||||
};
|
||||
|
||||
const publishQueuedJob = async (
|
||||
op: WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult>,
|
||||
payload: WhisperAlignJobRequest | PdfLayoutJobRequest,
|
||||
): Promise<void> => {
|
||||
const { js } = await ensureConnected();
|
||||
if (op.kind === 'whisper_align') {
|
||||
await js.publish(WHISPER_JOBS_SUBJECT, whisperJobCodec.encode({
|
||||
jobId: op.jobId,
|
||||
opId: op.opId,
|
||||
opKey: op.opKey,
|
||||
kind: 'whisper_align',
|
||||
queuedAt: op.queuedAt,
|
||||
payload: payload as WhisperAlignJobRequest,
|
||||
}));
|
||||
return;
|
||||
}
|
||||
|
||||
await js.publish(LAYOUT_JOBS_SUBJECT, layoutJobCodec.encode({
|
||||
jobId: op.jobId,
|
||||
opId: op.opId,
|
||||
opKey: op.opKey,
|
||||
kind: 'pdf_layout',
|
||||
queuedAt: op.queuedAt,
|
||||
payload: payload as PdfLayoutJobRequest,
|
||||
}));
|
||||
};
|
||||
|
||||
const enqueueOrReuseOperation = async (
|
||||
req: WorkerOperationRequest,
|
||||
): Promise<WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult>> => {
|
||||
const opKey = req.opKey.trim();
|
||||
const indexKey = opIndexKvKey(opKey);
|
||||
const opKeyHash = hashOpKey(opKey).slice(0, 16);
|
||||
const { kv } = await ensureConnected();
|
||||
|
||||
for (let attemptNo = 0; attemptNo < 10; attemptNo += 1) {
|
||||
const indexEntry = await kv.get(indexKey);
|
||||
if (indexEntry && indexEntry.operation === 'PUT') {
|
||||
const pointer = opIndexCodec.decode(indexEntry.value);
|
||||
const current = await getOpState(pointer.opId);
|
||||
|
||||
if (!current) {
|
||||
await sleep(25);
|
||||
continue;
|
||||
}
|
||||
|
||||
if (current && current.kind === req.kind) {
|
||||
const ageMs = Date.now() - current.updatedAt;
|
||||
if (current.status === 'succeeded') {
|
||||
app.log.info({
|
||||
kind: req.kind,
|
||||
opId: current.opId,
|
||||
jobId: current.jobId,
|
||||
opKeyHash,
|
||||
action: 'reused_terminal',
|
||||
status: current.status,
|
||||
ageMs,
|
||||
}, 'op.accepted');
|
||||
return current;
|
||||
}
|
||||
if ((current.status === 'queued' || current.status === 'running') && ageMs <= opStaleMs) {
|
||||
app.log.info({
|
||||
kind: req.kind,
|
||||
opId: current.opId,
|
||||
jobId: current.jobId,
|
||||
opKeyHash,
|
||||
action: 'reused_inflight',
|
||||
status: current.status,
|
||||
ageMs,
|
||||
}, 'op.accepted');
|
||||
return current;
|
||||
}
|
||||
if ((current.status === 'queued' || current.status === 'running') && ageMs > opStaleMs) {
|
||||
app.log.warn({
|
||||
kind: req.kind,
|
||||
opId: current.opId,
|
||||
jobId: current.jobId,
|
||||
opKeyHash,
|
||||
status: current.status,
|
||||
ageMs,
|
||||
opStaleMs,
|
||||
}, 'op.stuck_detected');
|
||||
}
|
||||
}
|
||||
|
||||
const replacementReason = (() => {
|
||||
if (current.kind !== req.kind) return 'kind_mismatch';
|
||||
const ageMs = Date.now() - current.updatedAt;
|
||||
if ((current.status === 'queued' || current.status === 'running') && ageMs > opStaleMs) return 'stale_running';
|
||||
if (current.status === 'failed') return 'failed_prior';
|
||||
return `status_${current.status}`;
|
||||
})();
|
||||
|
||||
const now = Date.now();
|
||||
const replacement: WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult> = {
|
||||
opId: crypto.randomUUID(),
|
||||
opKey,
|
||||
kind: req.kind,
|
||||
jobId: crypto.randomUUID(),
|
||||
status: 'queued',
|
||||
queuedAt: now,
|
||||
updatedAt: now,
|
||||
};
|
||||
|
||||
try {
|
||||
await kv.update(indexKey, opIndexCodec.encode({ opId: replacement.opId }), indexEntry.revision);
|
||||
} catch (error) {
|
||||
if (isCasConflictError(error)) continue;
|
||||
throw error;
|
||||
}
|
||||
|
||||
await putOpState(replacement);
|
||||
await putJobState({
|
||||
jobId: replacement.jobId,
|
||||
opId: replacement.opId,
|
||||
opKey: replacement.opKey,
|
||||
kind: replacement.kind,
|
||||
status: 'queued',
|
||||
timestamp: now,
|
||||
updatedAt: now,
|
||||
});
|
||||
|
||||
try {
|
||||
await publishQueuedJob(replacement, req.payload);
|
||||
app.log.info({
|
||||
kind: req.kind,
|
||||
opId: replacement.opId,
|
||||
jobId: replacement.jobId,
|
||||
opKeyHash,
|
||||
reason: replacementReason,
|
||||
status: replacement.status,
|
||||
}, 'op.replaced');
|
||||
app.log.info({
|
||||
kind: req.kind,
|
||||
opId: replacement.opId,
|
||||
jobId: replacement.jobId,
|
||||
opKeyHash,
|
||||
action: 'replaced',
|
||||
status: replacement.status,
|
||||
reason: replacementReason,
|
||||
}, 'op.accepted');
|
||||
return replacement;
|
||||
} catch (error) {
|
||||
const failed: WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult> = {
|
||||
...replacement,
|
||||
status: 'failed',
|
||||
updatedAt: Date.now(),
|
||||
error: { message: toErrorMessage(error) },
|
||||
};
|
||||
await putOpState(failed);
|
||||
await putJobState({
|
||||
jobId: replacement.jobId,
|
||||
opId: replacement.opId,
|
||||
opKey: replacement.opKey,
|
||||
kind: replacement.kind,
|
||||
status: 'failed',
|
||||
timestamp: replacement.queuedAt,
|
||||
updatedAt: failed.updatedAt,
|
||||
error: failed.error,
|
||||
});
|
||||
app.log.error({
|
||||
kind: req.kind,
|
||||
opId: failed.opId,
|
||||
jobId: failed.jobId,
|
||||
opKeyHash,
|
||||
action: 'replaced',
|
||||
status: failed.status,
|
||||
reason: replacementReason,
|
||||
error: failed.error?.message ?? null,
|
||||
}, 'op.accepted');
|
||||
return failed;
|
||||
}
|
||||
}
|
||||
|
||||
const now = Date.now();
|
||||
const created: WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult> = {
|
||||
opId: crypto.randomUUID(),
|
||||
opKey,
|
||||
kind: req.kind,
|
||||
jobId: crypto.randomUUID(),
|
||||
status: 'queued',
|
||||
queuedAt: now,
|
||||
updatedAt: now,
|
||||
};
|
||||
|
||||
try {
|
||||
await kv.create(indexKey, opIndexCodec.encode({ opId: created.opId }));
|
||||
} catch (error) {
|
||||
if (isCasConflictError(error)) continue;
|
||||
throw error;
|
||||
}
|
||||
|
||||
await putOpState(created);
|
||||
await putJobState({
|
||||
jobId: created.jobId,
|
||||
opId: created.opId,
|
||||
opKey: created.opKey,
|
||||
kind: created.kind,
|
||||
status: 'queued',
|
||||
timestamp: now,
|
||||
updatedAt: now,
|
||||
});
|
||||
|
||||
try {
|
||||
await publishQueuedJob(created, req.payload);
|
||||
app.log.info({
|
||||
kind: req.kind,
|
||||
opId: created.opId,
|
||||
jobId: created.jobId,
|
||||
opKeyHash,
|
||||
action: 'created',
|
||||
status: created.status,
|
||||
}, 'op.accepted');
|
||||
return created;
|
||||
} catch (error) {
|
||||
const failed: WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult> = {
|
||||
...created,
|
||||
status: 'failed',
|
||||
updatedAt: Date.now(),
|
||||
error: { message: toErrorMessage(error) },
|
||||
};
|
||||
await putOpState(failed);
|
||||
await putJobState({
|
||||
jobId: created.jobId,
|
||||
opId: created.opId,
|
||||
opKey: created.opKey,
|
||||
kind: created.kind,
|
||||
status: 'failed',
|
||||
timestamp: created.queuedAt,
|
||||
updatedAt: failed.updatedAt,
|
||||
error: failed.error,
|
||||
});
|
||||
app.log.error({
|
||||
kind: req.kind,
|
||||
opId: failed.opId,
|
||||
jobId: failed.jobId,
|
||||
opKeyHash,
|
||||
action: 'created',
|
||||
status: failed.status,
|
||||
error: failed.error?.message ?? null,
|
||||
}, 'op.accepted');
|
||||
return failed;
|
||||
}
|
||||
}
|
||||
|
||||
throw new Error('Unable to reserve operation after repeated CAS conflicts');
|
||||
return await operationStateStore.getOpState(opId);
|
||||
};
|
||||
|
||||
const releaseHttp = (request: FastifyRequest): void => {
|
||||
|
|
@ -1024,7 +787,15 @@ async function main(): Promise<void> {
|
|||
};
|
||||
}
|
||||
|
||||
const op = await enqueueOrReuseOperation(parsed.data as WorkerOperationRequest);
|
||||
const requestOp = parsed.data as WorkerOperationRequest;
|
||||
const op = await orchestrator.enqueueOrReuse(requestOp);
|
||||
app.log.info({
|
||||
kind: requestOp.kind,
|
||||
opId: op.opId,
|
||||
jobId: op.jobId,
|
||||
status: op.status,
|
||||
opKeyHash: hashOpKey(requestOp.opKey.trim()).slice(0, 16),
|
||||
}, 'op.accepted');
|
||||
reply.code(202);
|
||||
return op;
|
||||
});
|
||||
|
|
|
|||
195
tests/unit/compute-worker-control-plane-jetstream.spec.ts
Normal file
195
tests/unit/compute-worker-control-plane-jetstream.spec.ts
Normal file
|
|
@ -0,0 +1,195 @@
|
|||
import { expect, test } from '@playwright/test';
|
||||
import { OperationOrchestrator } from '../../compute/core/src/control-plane';
|
||||
import type { WorkerOperationRequest } from '../../compute/core/src/api-contracts';
|
||||
import {
|
||||
JetStreamOperationEventStream,
|
||||
JetStreamOperationQueue,
|
||||
JetStreamOperationStateStore,
|
||||
opEventsSubject,
|
||||
opIndexKvKey,
|
||||
opStateKvKey,
|
||||
type KvEntryLike,
|
||||
type KvStoreLike,
|
||||
} from '../../compute/worker/src/control-plane-jetstream';
|
||||
|
||||
class FakeKvStore implements KvStoreLike {
|
||||
private readonly data = new Map<string, KvEntryLike>();
|
||||
private revision = 0;
|
||||
|
||||
async get(key: string): Promise<KvEntryLike | null> {
|
||||
const value = this.data.get(key);
|
||||
if (!value) return null;
|
||||
return {
|
||||
operation: value.operation,
|
||||
value: value.value.slice(),
|
||||
revision: value.revision,
|
||||
};
|
||||
}
|
||||
|
||||
async put(key: string, data: Uint8Array): Promise<void> {
|
||||
this.revision += 1;
|
||||
this.data.set(key, {
|
||||
operation: 'PUT',
|
||||
value: data.slice(),
|
||||
revision: this.revision,
|
||||
});
|
||||
}
|
||||
|
||||
async create(key: string, data: Uint8Array): Promise<void> {
|
||||
if (this.data.has(key)) {
|
||||
throw new Error('key exists');
|
||||
}
|
||||
this.revision += 1;
|
||||
this.data.set(key, {
|
||||
operation: 'PUT',
|
||||
value: data.slice(),
|
||||
revision: this.revision,
|
||||
});
|
||||
}
|
||||
|
||||
async update(key: string, data: Uint8Array, version: number): Promise<void> {
|
||||
const current = this.data.get(key);
|
||||
if (!current || current.revision !== version) {
|
||||
throw new Error('wrong last sequence');
|
||||
}
|
||||
this.revision += 1;
|
||||
this.data.set(key, {
|
||||
operation: 'PUT',
|
||||
value: data.slice(),
|
||||
revision: this.revision,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
class FakeJetStream {
|
||||
private seq = 0;
|
||||
readonly published: Array<{ subject: string; payload: unknown; seq: number }> = [];
|
||||
|
||||
async publish(subject: string, data: Uint8Array): Promise<{ seq: number }> {
|
||||
this.seq += 1;
|
||||
const payload = JSON.parse(new TextDecoder().decode(data)) as unknown;
|
||||
this.published.push({ subject, payload, seq: this.seq });
|
||||
return { seq: this.seq };
|
||||
}
|
||||
}
|
||||
|
||||
function buildPdfRequest(opKey: string): WorkerOperationRequest {
|
||||
return {
|
||||
kind: 'pdf_layout',
|
||||
opKey,
|
||||
payload: {
|
||||
documentId: 'd1',
|
||||
namespace: null,
|
||||
documentObjectKey: 's3://bucket/doc.pdf',
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
test.describe('worker jetstream control-plane adapters', () => {
|
||||
test('state store compareAndSet handles create/update semantics', async () => {
|
||||
const kv = new FakeKvStore();
|
||||
const store = new JetStreamOperationStateStore({ getKv: async () => kv });
|
||||
|
||||
const created = await store.compareAndSetOpIndex({
|
||||
opKey: 'k1',
|
||||
newOpId: 'op-1',
|
||||
expectedOpId: null,
|
||||
});
|
||||
const failedCreate = await store.compareAndSetOpIndex({
|
||||
opKey: 'k1',
|
||||
newOpId: 'op-2',
|
||||
expectedOpId: null,
|
||||
});
|
||||
const wrongExpected = await store.compareAndSetOpIndex({
|
||||
opKey: 'k1',
|
||||
newOpId: 'op-2',
|
||||
expectedOpId: 'op-x',
|
||||
});
|
||||
const updated = await store.compareAndSetOpIndex({
|
||||
opKey: 'k1',
|
||||
newOpId: 'op-2',
|
||||
expectedOpId: 'op-1',
|
||||
});
|
||||
|
||||
expect(created).toBeTruthy();
|
||||
expect(failedCreate).toBeFalsy();
|
||||
expect(wrongExpected).toBeFalsy();
|
||||
expect(updated).toBeTruthy();
|
||||
expect(await store.getOpIndex('k1')).toEqual({ opId: 'op-2' });
|
||||
});
|
||||
|
||||
test('queue and event adapters publish expected JetStream subjects', async () => {
|
||||
const js = new FakeJetStream();
|
||||
const queue = new JetStreamOperationQueue({
|
||||
getJs: async () => js,
|
||||
whisperSubject: 'jobs.whisper',
|
||||
layoutSubject: 'jobs.layout',
|
||||
});
|
||||
const events = new JetStreamOperationEventStream({
|
||||
getJs: async () => js,
|
||||
});
|
||||
|
||||
await queue.enqueue({
|
||||
jobId: 'j1',
|
||||
opId: 'o1',
|
||||
opKey: 'k1',
|
||||
kind: 'pdf_layout',
|
||||
queuedAt: 1000,
|
||||
payload: { documentId: 'd1', namespace: null, documentObjectKey: 'obj' },
|
||||
});
|
||||
|
||||
const appended = await events.append('o1', {
|
||||
opId: 'o1',
|
||||
opKey: 'k1',
|
||||
kind: 'pdf_layout',
|
||||
jobId: 'j1',
|
||||
status: 'queued',
|
||||
queuedAt: 1000,
|
||||
updatedAt: 1000,
|
||||
});
|
||||
|
||||
expect(js.published.map((entry) => entry.subject)).toEqual(['jobs.layout', opEventsSubject('o1')]);
|
||||
expect(appended.eventId).toBe(2);
|
||||
});
|
||||
|
||||
test('orchestrator integration writes index/state and reuses active op', async () => {
|
||||
const kv = new FakeKvStore();
|
||||
const js = new FakeJetStream();
|
||||
|
||||
const store = new JetStreamOperationStateStore({ getKv: async () => kv });
|
||||
const events = new JetStreamOperationEventStream({ getJs: async () => js });
|
||||
const queue = new JetStreamOperationQueue({
|
||||
getJs: async () => js,
|
||||
whisperSubject: 'jobs.whisper',
|
||||
layoutSubject: 'jobs.layout',
|
||||
});
|
||||
|
||||
let now = 1_000;
|
||||
let nextId = 1;
|
||||
const orchestrator = new OperationOrchestrator({
|
||||
queue,
|
||||
stateStore: store,
|
||||
eventStream: events,
|
||||
config: { opStaleMs: 10_000, maxCasRetries: 3 },
|
||||
clock: { now: () => now },
|
||||
idFactory: {
|
||||
opId: () => `op-${nextId}`,
|
||||
jobId: () => `job-${nextId++}`,
|
||||
},
|
||||
});
|
||||
|
||||
const req = buildPdfRequest('k-integration');
|
||||
const first = await orchestrator.enqueueOrReuse(req);
|
||||
now = 2_000;
|
||||
const reused = await orchestrator.enqueueOrReuse(req);
|
||||
|
||||
expect(first.opId).toBe('op-1');
|
||||
expect(reused.opId).toBe('op-1');
|
||||
|
||||
const indexEntry = await kv.get(opIndexKvKey('k-integration'));
|
||||
const stateEntry = await kv.get(opStateKvKey('op-1'));
|
||||
expect(indexEntry?.operation).toBe('PUT');
|
||||
expect(stateEntry?.operation).toBe('PUT');
|
||||
expect(js.published.map((entry) => entry.subject)).toEqual(['ops.events.op-1', 'jobs.layout']);
|
||||
});
|
||||
});
|
||||
Loading…
Reference in a new issue