diff --git a/compute/worker/src/control-plane/jetstream.ts b/compute/worker/src/control-plane/jetstream.ts index 52d183f..642906d 100644 --- a/compute/worker/src/control-plane/jetstream.ts +++ b/compute/worker/src/control-plane/jetstream.ts @@ -14,6 +14,7 @@ import type { WhisperAlignJobRequest, WorkerOperationKind, } from '@openreader/compute-core/api-contracts'; +import { createJsonCodec } from './json-codec'; export interface KvEntryLike { operation?: string; @@ -28,24 +29,6 @@ export interface KvStoreLike { update(key: string, data: Uint8Array, version: number): Promise; } -type JsonCodec = { - encode(value: T): Uint8Array; - decode(data: Uint8Array): T; -}; - -function createJsonCodec(): JsonCodec { - 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); @@ -64,7 +47,8 @@ interface OpIndexEntry { opId: string; } -const OP_EVENTS_SUBJECT_PREFIX = 'ops.events'; +export const OP_EVENTS_SUBJECT_PREFIX = 'ops.events'; +export const OP_EVENTS_SUBJECT_WILDCARD = `${OP_EVENTS_SUBJECT_PREFIX}.*`; export function hashOpKey(opKey: string): string { return createHash('sha256').update(opKey).digest('hex'); diff --git a/compute/worker/src/control-plane/json-codec.ts b/compute/worker/src/control-plane/json-codec.ts new file mode 100644 index 0000000..27d3ea0 --- /dev/null +++ b/compute/worker/src/control-plane/json-codec.ts @@ -0,0 +1,17 @@ +export type JsonCodec = { + encode(value: T): Uint8Array; + decode(data: Uint8Array): T; +}; + +export function createJsonCodec(): JsonCodec { + 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; + }, + }; +} diff --git a/compute/worker/src/server.ts b/compute/worker/src/server.ts index 303ac9b..8b82485 100644 --- a/compute/worker/src/server.ts +++ b/compute/worker/src/server.ts @@ -50,10 +50,12 @@ import type { import { GetObjectCommand, PutObjectCommand, S3Client } from '@aws-sdk/client-s3'; import { JetStreamOperationEventStream, + OP_EVENTS_SUBJECT_WILDCARD, JetStreamOperationQueue, JetStreamOperationStateStore, hashOpKey, } from './control-plane/jetstream'; +import { type JsonCodec, createJsonCodec } from './control-plane/json-codec'; const JOBS_STREAM_NAME = 'compute_jobs'; const WHISPER_JOBS_SUBJECT = 'jobs.whisper'; @@ -65,7 +67,6 @@ const COMPUTE_STATE_BUCKET = 'compute_state'; const COMPUTE_STATE_TTL_MS = 24 * 60 * 60 * 1000; const LOOP_ERROR_BACKOFF_MS = 500; const RUNNING_HEARTBEAT_MS = 5000; -const OP_EVENTS_SUBJECT_PREFIX = 'ops.events'; const DOCUMENT_ID_REGEX = /^[a-f0-9]{64}$/i; const SAFE_NAMESPACE_REGEX = /^[a-zA-Z0-9._-]{1,128}$/; const WHISPER_MAX_DELIVER = 1; @@ -105,11 +106,6 @@ interface NatsSession { type StreamedOperationState = WorkerOperationState; -type JsonCodec = { - encode(value: T): Uint8Array; - decode(data: Uint8Array): T; -}; - function requireEnv(name: string): string { const value = process.env[name]?.trim(); if (!value) throw new Error(`${name} is required`); @@ -281,19 +277,6 @@ function isAlreadyExistsError(error: unknown): boolean { return message.includes('already in use') || message.includes('already exists'); } -function createJsonCodec(): JsonCodec { - 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 sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } @@ -395,7 +378,7 @@ async function ensureJetStreamResources( const eventsStreamConfig = { name: EVENTS_STREAM_NAME, - subjects: [`${OP_EVENTS_SUBJECT_PREFIX}.*`], + subjects: [OP_EVENTS_SUBJECT_WILDCARD], retention: RetentionPolicy.Limits, storage: StorageType.File, max_bytes: eventsMaxBytes, @@ -408,7 +391,7 @@ async function ensureJetStreamResources( } catch (error) { if (!isAlreadyExistsError(error)) throw error; await jsm.streams.update(EVENTS_STREAM_NAME, { - subjects: [`${OP_EVENTS_SUBJECT_PREFIX}.*`], + subjects: [OP_EVENTS_SUBJECT_WILDCARD], max_bytes: eventsMaxBytes, max_age: nanos(COMPUTE_STATE_TTL_MS), num_replicas: natsReplicas, @@ -973,7 +956,7 @@ async function main(): Promise { }; }; - async function processMessage(input: { + type ProcessMessageInput = { msg: JsMsg; codec: JsonCodec>; run: ( @@ -982,81 +965,212 @@ async function main(): Promise { hooks?: { onProgress?: (progress: PdfLayoutProgress) => Promise }, ) => Promise; workerLabel: string; - }): Promise { - let decoded: QueuedJob | null = null; - let heartbeat: NodeJS.Timeout | null = null; - let latestProgress: PdfLayoutProgress | undefined; - try { - decoded = input.codec.decode(input.msg.data); - const startedAt = Date.now(); - const queueWaitMs = safeDurationMs(decoded.queuedAt, startedAt); - const queueWaitTiming = typeof queueWaitMs === 'number' ? { queueWaitMs } : undefined; + }; - await orchestrator.markRunning({ - opId: decoded.opId, - startedAt, - updatedAt: startedAt, - ...(queueWaitTiming ? { timing: queueWaitTiming } : {}), + type ProcessMessageContext = { + decoded: QueuedJob; + workerLabel: string; + startedAt: number; + queueWaitTiming?: { queueWaitMs: number }; + latestProgress?: PdfLayoutProgress; + }; + + function buildQueueWaitTiming(queuedAt: number, now: number): { queueWaitMs: number } | undefined { + const queueWaitMs = safeDurationMs(queuedAt, now); + return typeof queueWaitMs === 'number' ? { queueWaitMs } : undefined; + } + + function extractTiming(result: unknown): WorkerJobTiming | undefined { + if (!result || typeof result !== 'object' || !('timing' in result)) return undefined; + return (result as { timing?: WorkerJobTiming }).timing; + } + + async function markRunning( + context: ProcessMessageContext, + updatedAt: number, + options?: { includeStartedAt?: boolean }, + ): Promise { + if (context.latestProgress) { + await orchestrator.markProgress({ + opId: context.decoded.opId, + progress: context.latestProgress, + updatedAt, + ...(context.queueWaitTiming ? { timing: context.queueWaitTiming } : {}), }); + return; + } + + await orchestrator.markRunning({ + opId: context.decoded.opId, + ...(options?.includeStartedAt === false ? {} : { startedAt: context.startedAt }), + updatedAt, + ...(context.queueWaitTiming ? { timing: context.queueWaitTiming } : {}), + }); + } + + async function markProgress( + context: ProcessMessageContext, + progress: PdfLayoutProgress, + updatedAt: number, + ): Promise { + context.latestProgress = progress; + await markRunning(context, updatedAt); + } + + async function markTerminal(input: { + context: ProcessMessageContext; + status: 'succeeded' | 'failed'; + result?: TResult; + errorMessage?: string; + timing?: WorkerJobTiming; + updatedAt: number; + }): Promise { + if (input.status === 'succeeded') { + await orchestrator.markSucceeded({ + opId: input.context.decoded.opId, + result: input.result as WhisperAlignJobResult | PdfLayoutJobResult, + updatedAt: input.updatedAt, + ...(input.timing ? { timing: input.timing } : {}), + }); + return; + } + + await orchestrator.markFailed({ + opId: input.context.decoded.opId, + error: { message: input.errorMessage ?? 'unknown worker failure' }, + updatedAt: input.updatedAt, + ...(input.timing ? { timing: input.timing } : {}), + }); + } + + async function handleRetry(input: { + context: ProcessMessageContext | null; + msg: JsMsg; + errorMessage: string; + }): Promise { + const deliveryCount = input.msg.info.deliveryCount; + const isWhisperAlign = input.context?.decoded.kind === 'whisper_align'; + const maxAttempts = isWhisperAlign ? WHISPER_MAX_DELIVER : pdfAttempts; + const hasRetriesLeft = !isWhisperAlign && deliveryCount < maxAttempts; + + if (input.context) { + const now = Date.now(); + const retryTiming = buildQueueWaitTiming(input.context.decoded.queuedAt, now); + const persistOpUpdate = hasRetriesLeft + ? (input.context.latestProgress + ? orchestrator.markProgress({ + opId: input.context.decoded.opId, + progress: input.context.latestProgress, + updatedAt: now, + ...(retryTiming ? { timing: retryTiming } : {}), + }) + : orchestrator.markRunning({ + opId: input.context.decoded.opId, + updatedAt: now, + ...(retryTiming ? { timing: retryTiming } : {}), + })) + : markTerminal({ + context: input.context, + status: 'failed', + errorMessage: input.errorMessage, + updatedAt: now, + ...(retryTiming ? { timing: retryTiming } : {}), + }); + + await persistOpUpdate.catch((stateError) => { + app.log.error({ + worker: input.context?.workerLabel, + opId: input.context?.decoded.opId, + jobId: input.context?.decoded.jobId, + error: toErrorMessage(stateError), + }, 'failed to persist operation state'); + }); + } + + if (hasRetriesLeft) { + input.msg.nak(); + app.log.error({ + worker: input.context?.workerLabel, + kind: input.context?.decoded.kind, + opId: input.context?.decoded.opId, + jobId: input.context?.decoded.jobId, + status: 'running', + error: input.errorMessage, + deliveryCount, + maxAttempts, + retryAction: 'nack_retry', + }, 'job.terminal'); + return; + } + + input.msg.term(input.errorMessage); + app.log.error({ + worker: input.context?.workerLabel, + kind: input.context?.decoded.kind, + opId: input.context?.decoded.opId, + jobId: input.context?.decoded.jobId, + status: 'failed', + error: input.errorMessage, + deliveryCount, + maxAttempts, + retrySuppressed: isWhisperAlign ? 'whisper_align' : undefined, + retryAction: 'term', + }, 'job.terminal'); + } + + async function processMessage(input: ProcessMessageInput): Promise { + let context: ProcessMessageContext | null = null; + let heartbeat: NodeJS.Timeout | null = null; + try { + const decoded = input.codec.decode(input.msg.data); + const startedAt = Date.now(); + context = { + decoded, + workerLabel: input.workerLabel, + startedAt, + queueWaitTiming: buildQueueWaitTiming(decoded.queuedAt, startedAt), + }; + + await markRunning(context, startedAt); app.log.info({ worker: input.workerLabel, kind: decoded.kind, opId: decoded.opId, jobId: decoded.jobId, - queueWaitMs: queueWaitMs ?? null, + queueWaitMs: context.queueWaitTiming?.queueWaitMs ?? null, deliveryCount: input.msg.info.deliveryCount, }, 'job.started'); - const persistRunningState = async (updatedAt: number): Promise => { - if (latestProgress) { - await orchestrator.markProgress({ - opId: decoded!.opId, - progress: latestProgress, - updatedAt, - ...(queueWaitTiming ? { timing: queueWaitTiming } : {}), - }); - } else { - await orchestrator.markRunning({ - opId: decoded!.opId, - startedAt, - updatedAt, - ...(queueWaitTiming ? { timing: queueWaitTiming } : {}), - }); - } - }; - heartbeat = setInterval(() => { const now = Date.now(); - void persistRunningState(now).catch((stateError) => { + void markRunning(context!, now).catch((stateError) => { app.log.error({ worker: input.workerLabel, - opId: decoded?.opId, - jobId: decoded?.jobId, + opId: context?.decoded.opId, + jobId: context?.decoded.jobId, error: toErrorMessage(stateError), }, 'failed to persist operation heartbeat state'); }); }, RUNNING_HEARTBEAT_MS); - const result = await input.run(decoded.payload, queueWaitMs ?? 0, { + const result = await input.run(decoded.payload, context.queueWaitTiming?.queueWaitMs ?? 0, { onProgress: async (progress) => { - latestProgress = progress; - await persistRunningState(Date.now()); + await markProgress(context!, progress, Date.now()); }, }); - const resultTiming = result && typeof result === 'object' && 'timing' in result - ? (result as { timing?: WorkerJobTiming }).timing - : undefined; + const resultTiming = extractTiming(result); const now = Date.now(); - await orchestrator.markSucceeded({ - opId: decoded.opId, - result: result as WhisperAlignJobResult | PdfLayoutJobResult, + await markTerminal({ + context, + status: 'succeeded', + result, + timing: resultTiming, updatedAt: now, - ...(resultTiming ? { timing: resultTiming } : {}), }); input.msg.ack(); - const terminalDurationMs = safeDurationMs(startedAt, now); + const terminalDurationMs = safeDurationMs(context.startedAt, now); const slowJobLogThresholdMs = SLOW_JOB_LOG_THRESHOLD_MS_BY_KIND[decoded.kind]; if ((terminalDurationMs ?? 0) >= slowJobLogThresholdMs) { app.log.info({ @@ -1079,75 +1193,11 @@ async function main(): Promise { timing: resultTiming ?? null, }, 'job.terminal'); } catch (error) { - const message = toErrorMessage(error); - const deliveryCount = input.msg.info.deliveryCount; - const isWhisperAlign = decoded?.kind === 'whisper_align'; - const maxAttempts = isWhisperAlign ? WHISPER_MAX_DELIVER : pdfAttempts; - const hasRetriesLeft = !isWhisperAlign && deliveryCount < maxAttempts; - - if (decoded) { - const now = Date.now(); - const queueWaitMs = safeDurationMs(decoded.queuedAt, now); - const queueWaitTiming = typeof queueWaitMs === 'number' ? { queueWaitMs } : undefined; - - const persistOpUpdate = hasRetriesLeft - ? (latestProgress - ? orchestrator.markProgress({ - opId: decoded.opId, - progress: latestProgress, - updatedAt: now, - ...(queueWaitTiming ? { timing: queueWaitTiming } : {}), - }) - : orchestrator.markRunning({ - opId: decoded.opId, - updatedAt: now, - ...(queueWaitTiming ? { timing: queueWaitTiming } : {}), - })) - : orchestrator.markFailed({ - opId: decoded.opId, - error: { message }, - updatedAt: now, - ...(queueWaitTiming ? { timing: queueWaitTiming } : {}), - }); - - await persistOpUpdate.catch((stateError) => { - app.log.error({ - worker: input.workerLabel, - opId: decoded?.opId, - jobId: decoded?.jobId, - error: toErrorMessage(stateError), - }, 'failed to persist operation state'); - }); - } - - if (hasRetriesLeft) { - input.msg.nak(); - app.log.error({ - worker: input.workerLabel, - kind: decoded?.kind, - opId: decoded?.opId, - jobId: decoded?.jobId, - status: 'running', - error: message, - deliveryCount, - maxAttempts, - retryAction: 'nack_retry', - }, 'job.terminal'); - } else { - input.msg.term(message); - app.log.error({ - worker: input.workerLabel, - kind: decoded?.kind, - opId: decoded?.opId, - jobId: decoded?.jobId, - status: 'failed', - error: message, - deliveryCount, - maxAttempts, - retrySuppressed: isWhisperAlign ? 'whisper_align' : undefined, - retryAction: 'term', - }, 'job.terminal'); - } + await handleRetry({ + context, + msg: input.msg, + errorMessage: toErrorMessage(error), + }); } finally { if (heartbeat) clearInterval(heartbeat); } diff --git a/src/app/api/documents/[id]/parsed/events/route.ts b/src/app/api/documents/[id]/parsed/events/route.ts index 077b7e0..315b3fc 100644 --- a/src/app/api/documents/[id]/parsed/events/route.ts +++ b/src/app/api/documents/[id]/parsed/events/route.ts @@ -4,6 +4,7 @@ import { db } from '@/db'; import { documents } from '@/db/schema'; import { requireAuthContext } from '@/lib/server/auth/auth'; import { getWorkerClientConfigFromEnv } from '@/lib/server/compute/worker'; +import { snapshotFromWorkerState } from '@/lib/server/compute/worker-parse-state'; import { fetchWorkerOperationState } from '@/lib/server/compute/worker-op-state'; import { isValidDocumentId } from '@/lib/server/documents/blobstore'; import { normalizeParseStatus, parseDocumentParseState } from '@/lib/server/documents/parse-state'; @@ -47,29 +48,6 @@ function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } -function mapWorkerStatusToParseStatus(status: WorkerOperationState['status']): PdfParseStatus { - switch (status) { - case 'queued': - return 'pending'; - case 'running': - return 'running'; - case 'succeeded': - return 'ready'; - case 'failed': - return 'failed'; - default: - return 'pending'; - } -} - -function snapshotFromWorkerState(state: WorkerOperationState): ParsedSnapshot { - const parseStatus = mapWorkerStatusToParseStatus(state.status); - return { - parseStatus, - parseProgress: parseStatus === 'running' ? (state.progress ?? null) : null, - }; -} - async function toSnapshotState(row: ParseRow): Promise { const state = await healStaleDocumentParseState({ documentId: row.id, diff --git a/src/app/api/documents/[id]/parsed/route.ts b/src/app/api/documents/[id]/parsed/route.ts index 5bcb2dc..492f1db 100644 --- a/src/app/api/documents/[id]/parsed/route.ts +++ b/src/app/api/documents/[id]/parsed/route.ts @@ -4,6 +4,7 @@ import { and, eq, inArray } from 'drizzle-orm'; import { db } from '@/db'; import { documents } from '@/db/schema'; import { requireAuthContext } from '@/lib/server/auth/auth'; +import { mergeNonReadyParseSnapshot } from '@/lib/server/compute/worker-parse-state'; import { fetchWorkerOperationState } from '@/lib/server/compute/worker-op-state'; import { getParsedDocumentBlob, @@ -21,7 +22,7 @@ import { healStaleDocumentParseState } from '@/lib/server/documents/parse-state- import { getOpenReaderTestNamespace, getUnclaimedUserIdForNamespace } from '@/lib/server/testing/test-namespace'; import { isS3Configured } from '@/lib/server/storage/s3'; import type { ParsedPdfDocument } from '@/types/parsed-pdf'; -import type { PdfLayoutJobResult, WorkerOperationState } from '@openreader/compute-core/api-contracts'; +import type { PdfLayoutJobResult } from '@openreader/compute-core/api-contracts'; export const dynamic = 'force-dynamic'; @@ -37,21 +38,6 @@ function hasAnyParsedBlocks(doc: ParsedPdfDocument | null): boolean { return doc.pages.some((page) => Array.isArray(page.blocks) && page.blocks.length > 0); } -function mapWorkerStatusToParseStatus(status: WorkerOperationState['status']) { - switch (status) { - case 'queued': - return 'pending' as const; - case 'running': - return 'running' as const; - case 'succeeded': - return 'ready' as const; - case 'failed': - return 'failed' as const; - default: - return 'pending' as const; - } -} - export async function GET(req: NextRequest, ctx: { params: Promise<{ id: string }> }) { try { if (!isS3Configured()) return s3NotConfiguredResponse(); @@ -103,12 +89,13 @@ export async function GET(req: NextRequest, ctx: { params: Promise<{ id: string if (opId && effectiveStatus !== 'ready') { const workerState = await fetchWorkerOperationState(opId); if (workerState && workerState.opId === opId) { - const workerStatus = mapWorkerStatusToParseStatus(workerState.status); - // Keep DB/blob as source of truth for "ready"; prefer worker only for active/failed states. - if (workerStatus === 'pending' || workerStatus === 'running' || workerStatus === 'failed') { - effectiveStatus = workerStatus; - effectiveProgress = workerStatus === 'running' ? (workerState.progress ?? null) : null; - } + const merged = mergeNonReadyParseSnapshot({ + parseStatus: effectiveStatus, + parseProgress: effectiveProgress, + workerState, + }); + effectiveStatus = merged.parseStatus; + effectiveProgress = merged.parseProgress; } } @@ -226,11 +213,13 @@ export async function POST(req: NextRequest, ctx: { params: Promise<{ id: string if (opId && effectiveStatus !== 'ready') { const workerState = await fetchWorkerOperationState(opId); if (workerState && workerState.opId === opId) { - const workerStatus = mapWorkerStatusToParseStatus(workerState.status); - if (workerStatus === 'pending' || workerStatus === 'running' || workerStatus === 'failed') { - effectiveStatus = workerStatus; - effectiveProgress = workerStatus === 'running' ? (workerState.progress ?? null) : null; - } + const merged = mergeNonReadyParseSnapshot({ + parseStatus: effectiveStatus, + parseProgress: effectiveProgress, + workerState, + }); + effectiveStatus = merged.parseStatus; + effectiveProgress = merged.parseProgress; } } diff --git a/src/lib/server/compute/worker-parse-state.ts b/src/lib/server/compute/worker-parse-state.ts new file mode 100644 index 0000000..c0eaa3a --- /dev/null +++ b/src/lib/server/compute/worker-parse-state.ts @@ -0,0 +1,42 @@ +import type { PdfLayoutJobResult, WorkerOperationState } from '@openreader/compute-core/api-contracts'; +import type { PdfParseProgress, PdfParseStatus } from '@/types/parsed-pdf'; + +export function mapWorkerStatusToParseStatus(status: WorkerOperationState['status']): PdfParseStatus { + switch (status) { + case 'queued': + return 'pending'; + case 'running': + return 'running'; + case 'succeeded': + return 'ready'; + case 'failed': + return 'failed'; + default: + return 'pending'; + } +} + +export function snapshotFromWorkerState( + state: WorkerOperationState, +): { parseStatus: PdfParseStatus; parseProgress: PdfParseProgress | null } { + const parseStatus = mapWorkerStatusToParseStatus(state.status); + return { + parseStatus, + parseProgress: parseStatus === 'running' ? (state.progress ?? null) : null, + }; +} + +export function mergeNonReadyParseSnapshot(input: { + parseStatus: PdfParseStatus; + parseProgress: PdfParseProgress | null; + workerState: WorkerOperationState; +}): { parseStatus: PdfParseStatus; parseProgress: PdfParseProgress | null } { + const workerSnapshot = snapshotFromWorkerState(input.workerState); + if (workerSnapshot.parseStatus === 'ready') { + return { + parseStatus: input.parseStatus, + parseProgress: input.parseProgress, + }; + } + return workerSnapshot; +}