From df05a7d7a34b16aa10b11d68025044cc1147b425 Mon Sep 17 00:00:00 2001 From: Richard R Date: Tue, 26 May 2026 11:41:33 -0600 Subject: [PATCH] refactor(control-plane): hard-cut worker events and shared SSE helpers --- compute/core/src/control-plane/index.ts | 2 + .../core/src/control-plane/orchestrator.ts | 71 ++---- compute/core/src/control-plane/sse.ts | 49 ++++ .../core/src/control-plane/state-machine.ts | 65 ++++++ .../jetstream.ts} | 116 +++++++++- compute/worker/src/server.ts | 213 ++++++++---------- .../api/documents/[id]/parsed/events/route.ts | 27 +-- src/lib/server/compute/worker.ts | 29 +-- tests/unit/compute-control-plane.spec.ts | 48 ++++ ...ute-worker-control-plane-jetstream.spec.ts | 25 +- 10 files changed, 417 insertions(+), 228 deletions(-) create mode 100644 compute/core/src/control-plane/sse.ts create mode 100644 compute/core/src/control-plane/state-machine.ts rename compute/worker/src/{control-plane-jetstream.ts => control-plane/jetstream.ts} (63%) diff --git a/compute/core/src/control-plane/index.ts b/compute/core/src/control-plane/index.ts index 9b67931..8e6b18c 100644 --- a/compute/core/src/control-plane/index.ts +++ b/compute/core/src/control-plane/index.ts @@ -1,3 +1,5 @@ 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 3b45372..950fe71 100644 --- a/compute/core/src/control-plane/orchestrator.ts +++ b/compute/core/src/control-plane/orchestrator.ts @@ -6,6 +6,13 @@ import type { WorkerOperationRequest, WorkerOperationState, } from '../api-contracts'; +import { + buildQueuedState, + createErrorShape, + explainReplacementReason, + isTerminalStatus, + shouldReuseExistingOperation, +} from './state-machine'; import type { OperationClock, OperationEventStream, @@ -22,51 +29,6 @@ function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } -function isTerminalStatus(status: WorkerOperationState['status']): boolean { - return status === 'succeeded' || status === 'failed'; -} - -function isInflightStatus(status: WorkerOperationState['status']): boolean { - return status === 'queued' || status === 'running'; -} - -function createErrorShape(error: unknown): WorkerJobErrorShape { - if (error && typeof error === 'object' && 'message' in error && typeof (error as { message: unknown }).message === 'string') { - return { message: (error as { message: string }).message }; - } - return { message: String(error) }; -} - -function queuedStateFromRequest(input: { - request: WorkerOperationRequest; - opId: string; - jobId: string; - queuedAt: number; -}): WorkerOperationState { - return { - opId: input.opId, - opKey: input.request.opKey, - kind: input.request.kind, - jobId: input.jobId, - status: 'queued', - queuedAt: input.queuedAt, - updatedAt: input.queuedAt, - }; -} - -function mapReplacementReason(input: { - current: WorkerOperationState; - requestKind: WorkerOperationKind; - now: number; - opStaleMs: number; -}): string { - if (input.current.kind !== input.requestKind) return 'kind_mismatch'; - const ageMs = input.now - input.current.updatedAt; - if (isInflightStatus(input.current.status) && ageMs > input.opStaleMs) return 'stale_running'; - if (input.current.status === 'failed') return 'failed_prior'; - return `status_${input.current.status}`; -} - export interface OperationOrchestratorDeps { queue: OperationQueue; stateStore: OperationStateStore; @@ -127,15 +89,16 @@ export class OperationOrchestrator { } const now = this.clock.now(); - const ageMs = now - current.updatedAt; - if (current.kind === request.kind) { - if (current.status === 'succeeded') return current; - if (isInflightStatus(current.status) && ageMs <= this.opStaleMs) { - return current; - } + if (shouldReuseExistingOperation({ + current, + requestKind: request.kind, + now, + opStaleMs: this.opStaleMs, + })) { + return current; } - const replacement = queuedStateFromRequest({ + const replacement = buildQueuedState({ request, opId: this.ids.opId(), jobId: this.ids.jobId(), @@ -170,7 +133,7 @@ export class OperationOrchestrator { } const now = this.clock.now(); - const created = queuedStateFromRequest({ + const created = buildQueuedState({ request, opId: this.ids.opId(), jobId: this.ids.jobId(), @@ -304,7 +267,7 @@ export class OperationOrchestrator { current: WorkerOperationState; requestKind: WorkerOperationKind; }): Promise { - return mapReplacementReason({ + return explainReplacementReason({ current: input.current, requestKind: input.requestKind, now: this.clock.now(), diff --git a/compute/core/src/control-plane/sse.ts b/compute/core/src/control-plane/sse.ts new file mode 100644 index 0000000..712f6af --- /dev/null +++ b/compute/core/src/control-plane/sse.ts @@ -0,0 +1,49 @@ +export interface SseFrameInput { + event?: string; + id?: string | number; + data?: T; + comment?: string; +} + +export function encodeSseFrame(input: SseFrameInput): string { + const lines: string[] = []; + if (typeof input.comment === 'string') { + lines.push(`: ${input.comment}`); + } + if (typeof input.id !== 'undefined') { + lines.push(`id: ${String(input.id)}`); + } + if (typeof input.event === 'string' && input.event.trim()) { + lines.push(`event: ${input.event}`); + } + if (typeof input.data !== 'undefined') { + const serialized = typeof input.data === 'string' ? input.data : JSON.stringify(input.data); + for (const line of serialized.replace(/\r\n/g, '\n').split('\n')) { + lines.push(`data: ${line}`); + } + } + return `${lines.join('\n')}\n\n`; +} + +export function parseSsePayload(frame: string): string | null { + const lines = frame.replace(/\r\n/g, '\n').split('\n'); + const dataLines: string[] = []; + for (const line of lines) { + if (!line.startsWith('data:')) continue; + dataLines.push(line.slice('data:'.length).trimStart()); + } + return dataLines.length > 0 ? dataLines.join('\n') : null; +} + +export function parseSseEventId(frame: string): number | null { + const lines = frame.replace(/\r\n/g, '\n').split('\n'); + for (const line of lines) { + if (!line.startsWith('id:')) continue; + const value = Number(line.slice('id:'.length).trim()); + if (Number.isFinite(value) && value > 0) { + return Math.floor(value); + } + } + return null; +} + diff --git a/compute/core/src/control-plane/state-machine.ts b/compute/core/src/control-plane/state-machine.ts new file mode 100644 index 0000000..d280215 --- /dev/null +++ b/compute/core/src/control-plane/state-machine.ts @@ -0,0 +1,65 @@ +import type { + WorkerJobErrorShape, + WorkerJobState, + WorkerOperationKind, + WorkerOperationRequest, + WorkerOperationState, +} from '../api-contracts'; + +export function isTerminalStatus(status: WorkerJobState): boolean { + return status === 'succeeded' || status === 'failed'; +} + +export function isInflightStatus(status: WorkerJobState): boolean { + return status === 'queued' || status === 'running'; +} + +export function createErrorShape(error: unknown): WorkerJobErrorShape { + if (error && typeof error === 'object' && 'message' in error && typeof (error as { message: unknown }).message === 'string') { + return { message: (error as { message: string }).message }; + } + return { message: String(error) }; +} + +export function buildQueuedState(input: { + request: WorkerOperationRequest; + opId: string; + jobId: string; + queuedAt: number; +}): WorkerOperationState { + return { + opId: input.opId, + opKey: input.request.opKey, + kind: input.request.kind, + jobId: input.jobId, + status: 'queued', + queuedAt: input.queuedAt, + updatedAt: input.queuedAt, + }; +} + +export function explainReplacementReason(input: { + current: WorkerOperationState; + requestKind: WorkerOperationKind; + now: number; + opStaleMs: number; +}): string { + if (input.current.kind !== input.requestKind) return 'kind_mismatch'; + const ageMs = input.now - input.current.updatedAt; + if (isInflightStatus(input.current.status) && ageMs > input.opStaleMs) return 'stale_running'; + if (input.current.status === 'failed') return 'failed_prior'; + return `status_${input.current.status}`; +} + +export function shouldReuseExistingOperation(input: { + current: WorkerOperationState; + requestKind: WorkerOperationKind; + now: number; + opStaleMs: number; +}): boolean { + if (input.current.kind !== input.requestKind) return false; + if (input.current.status === 'succeeded') return true; + if (!isInflightStatus(input.current.status)) return false; + const ageMs = input.now - input.current.updatedAt; + return ageMs <= input.opStaleMs; +} diff --git a/compute/worker/src/control-plane-jetstream.ts b/compute/worker/src/control-plane/jetstream.ts similarity index 63% rename from compute/worker/src/control-plane-jetstream.ts rename to compute/worker/src/control-plane/jetstream.ts index deab1de..52d183f 100644 --- a/compute/worker/src/control-plane-jetstream.ts +++ b/compute/worker/src/control-plane/jetstream.ts @@ -1,5 +1,6 @@ import { createHash } from 'node:crypto'; -import type { JetStreamClient } from '@nats-io/jetstream'; +import { AckPolicy, DeliverPolicy, ReplayPolicy, type JetStreamClient, type JetStreamManager } from '@nats-io/jetstream'; +import { nanos } from '@nats-io/transport-node'; import type { OperationEvent, OperationEventStream, @@ -148,14 +149,24 @@ export class JetStreamOperationStateStore implements Operation } export interface JetStreamOperationEventStreamDeps { - getJs: () => Promise>; + getJs: () => Promise>; + getJsm: () => Promise>; + eventsStreamName: string; + inactiveThresholdMs?: number; } export class JetStreamOperationEventStream implements OperationEventStream { - private readonly getJs: () => Promise>; + private readonly getJs: () => Promise>; + private readonly getJsm: () => Promise>; + private readonly eventsStreamName: string; + private readonly inactiveThresholdNanos: number; + private readonly opStateCodec = createJsonCodec>(); constructor(deps: JetStreamOperationEventStreamDeps) { this.getJs = deps.getJs; + this.getJsm = deps.getJsm; + this.eventsStreamName = deps.eventsStreamName; + this.inactiveThresholdNanos = nanos((deps.inactiveThresholdMs ?? 60_000)); } async append(opId: string, snapshot: OperationState): Promise> { @@ -168,12 +179,103 @@ export class JetStreamOperationEventStream implements Operatio }; } - async listSince(): Promise[]> { - return []; + private async createConsumer(input: { + opId: string; + sinceEventId?: number; + replayOnly: boolean; + }): Promise<{ name: string; js: Pick }> { + const js = await this.getJs(); + const jsm = await this.getJsm(); + const subject = opEventsSubject(input.opId); + const since = Math.max(0, Math.floor(input.sinceEventId ?? 0)); + const name = `op_events_${input.opId.slice(0, 12)}_${crypto.randomUUID().replaceAll('-', '').slice(0, 12)}`; + const config = { + name, + ack_policy: AckPolicy.None, + deliver_policy: since > 0 ? DeliverPolicy.StartSequence : (input.replayOnly ? DeliverPolicy.All : DeliverPolicy.New), + replay_policy: ReplayPolicy.Instant, + filter_subject: subject, + max_deliver: 1, + inactive_threshold: this.inactiveThresholdNanos, + ...(since > 0 ? { opt_start_seq: since + 1 } : {}), + }; + await jsm.consumers.add(this.eventsStreamName, config); + return { name, js }; } - async subscribe(): Promise<() => void> { - return () => undefined; + private async deleteConsumer(name: string): Promise { + const jsm = await this.getJsm(); + await jsm.consumers.delete(this.eventsStreamName, name).catch(() => undefined); + } + + async listSince(opId: string, sinceEventId: number, limit = 200): Promise[]> { + const boundedLimit = Number.isFinite(limit) && limit > 0 ? Math.floor(limit) : 200; + const { name, js } = await this.createConsumer({ + opId, + sinceEventId, + replayOnly: true, + }); + try { + const consumer = await js.consumers.get(this.eventsStreamName, name); + const output: OperationEvent[] = []; + while (output.length < boundedLimit) { + const msg = await consumer.next({ expires: 250 }); + if (!msg) break; + output.push({ + eventId: msg.seq, + snapshot: this.opStateCodec.decode(msg.data), + }); + } + return output; + } finally { + await this.deleteConsumer(name); + } + } + + async subscribe(input: { + opId: string; + sinceEventId?: number; + onEvent: (event: OperationEvent) => void | Promise; + onError?: (error: unknown) => void; + }): Promise<() => void> { + const { name, js } = await this.createConsumer({ + opId: input.opId, + sinceEventId: input.sinceEventId, + replayOnly: false, + }); + const consumer = await js.consumers.get(this.eventsStreamName, name); + const messages = await consumer.consume(); + let closed = false; + + void (async () => { + try { + for await (const msg of messages) { + if (closed) break; + try { + await input.onEvent({ + eventId: msg.seq, + snapshot: this.opStateCodec.decode(msg.data), + }); + } catch (error) { + input.onError?.(error); + } + } + } catch (error) { + if (!closed) input.onError?.(error); + } finally { + if (!closed) { + closed = true; + await this.deleteConsumer(name); + } + } + })(); + + return () => { + if (closed) return; + closed = true; + void messages.close().catch(() => undefined); + void this.deleteConsumer(name); + }; } } diff --git a/compute/worker/src/server.ts b/compute/worker/src/server.ts index fafa45b..3df73f5 100644 --- a/compute/worker/src/server.ts +++ b/compute/worker/src/server.ts @@ -33,7 +33,7 @@ import { withIdleTimeoutAndHardCap, withTimeout, } from '@openreader/compute-core'; -import { OperationOrchestrator } from '@openreader/compute-core/control-plane'; +import { encodeSseFrame, OperationOrchestrator } from '@openreader/compute-core/control-plane'; import type { PdfLayoutJobRequest, PdfLayoutJobResult, @@ -54,8 +54,7 @@ import { JetStreamOperationQueue, JetStreamOperationStateStore, hashOpKey, - opEventsSubject, -} from './control-plane-jetstream'; +} from './control-plane/jetstream'; const JOBS_STREAM_NAME = 'compute_jobs'; const WHISPER_JOBS_SUBJECT = 'jobs.whisper'; @@ -667,7 +666,6 @@ async function main(): Promise { const whisperJobCodec = createJsonCodec>(); const layoutJobCodec = createJsonCodec>(); const jobStateCodec = createJsonCodec>(); - const opEventCodec = createJsonCodec(); const putJobState = async (state: StoredJobState): Promise => { const { kv } = await ensureConnected(); @@ -680,6 +678,8 @@ async function main(): Promise { const operationEventStream = new JetStreamOperationEventStream({ getJs: async () => (await ensureConnected()).js, + getJsm: async () => (await ensureConnected()).jsm, + eventsStreamName: EVENTS_STREAM_NAME, }); const operationQueue = new JetStreamOperationQueue({ @@ -709,19 +709,6 @@ async function main(): Promise { }, }); - const putOpState = async (state: StreamedOperationState): Promise => { - await operationStateStore.putOpState(state); - try { - await operationEventStream.append(state.opId, state); - } catch (error) { - app.log.warn({ - opId: state.opId, - status: state.status, - error: toErrorMessage(error), - }, 'failed to publish op event'); - } - }; - const getOpState = async (opId: string): Promise => { return await operationStateStore.getOpState(opId); }; @@ -852,22 +839,38 @@ async function main(): Promise { reply.raw.setHeader('Connection', 'keep-alive'); reply.raw.setHeader('X-Accel-Buffering', 'no'); + let closed = false; + let unsubscribe: (() => void) | null = null; + const writeSnapshot = (snapshot: StreamedOperationState, eventId: number): void => { if (closed || reply.raw.writableEnded) return; const frameEvent: WorkerOperationEvent = { eventId, snapshot, }; - reply.raw.write(`id: ${eventId}\nevent: snapshot\ndata: ${JSON.stringify(frameEvent)}\n\n`); + reply.raw.write(encodeSseFrame({ + id: eventId, + event: 'snapshot', + data: frameEvent, + })); }; - let closed = false; - let consumerName: string | null = null; - let messages: Awaited> | null = null; - request.raw.on('close', () => { + const closeStream = (): void => { + if (closed) return; closed = true; + if (unsubscribe) { + unsubscribe(); + unsubscribe = null; + } activeSse = Math.max(0, activeSse - 1); markActivity(); + if (!reply.raw.writableEnded) { + reply.raw.end(); + } + }; + + request.raw.on('close', () => { + closeStream(); }); try { @@ -878,52 +881,41 @@ async function main(): Promise { return reply; } - const { jsm, js } = await ensureConnected(); - const consumerConfig = { - name: `op_events_${params.data.opId.slice(0, 12)}_${crypto.randomUUID().replaceAll('-', '').slice(0, 12)}`, - ack_policy: AckPolicy.None, - deliver_policy: sinceEventId > 0 ? DeliverPolicy.StartSequence : DeliverPolicy.New, - replay_policy: ReplayPolicy.Instant, - filter_subject: opEventsSubject(params.data.opId), - max_deliver: 1, - inactive_threshold: nanos(60_000_000_000), - ...(sinceEventId > 0 ? { opt_start_seq: sinceEventId + 1 } : {}), - }; - const info = await jsm.consumers.add(EVENTS_STREAM_NAME, consumerConfig); - consumerName = info.name; - const consumer = await js.consumers.get(EVENTS_STREAM_NAME, info.name); - messages = await consumer.consume(); + unsubscribe = await operationEventStream.subscribe({ + opId: params.data.opId, + sinceEventId, + onEvent: (event) => { + if (closed) return; + if (event.snapshot.opId !== params.data.opId) return; + const nextSignature = JSON.stringify(event.snapshot); + if (nextSignature !== signature) { + current = event.snapshot; + signature = nextSignature; + writeSnapshot(current, event.eventId); + } + if (isTerminalStatus(event.snapshot.status)) { + closeStream(); + } + }, + onError: (error) => { + app.log.warn({ + opId: params.data.opId, + error: toErrorMessage(error), + }, 'op events stream loop error'); + closeStream(); + }, + }); - for await (const msg of messages) { - if (closed) break; - const state = opEventCodec.decode(msg.data); - if (state.opId !== params.data.opId) continue; - - const nextSignature = JSON.stringify(state); - if (nextSignature !== signature) { - current = state; - signature = nextSignature; - writeSnapshot(current, msg.seq); - } - if (isTerminalStatus(state.status)) break; - } + await new Promise((resolve) => { + request.raw.once('close', () => resolve()); + }); } catch (error) { app.log.warn({ opId: params.data.opId, error: toErrorMessage(error), }, 'op events stream loop error'); } finally { - if (messages) { - await messages.close().catch(() => undefined); - } - if (consumerName) { - const { jsm } = await ensureConnected(); - await jsm.consumers.delete(EVENTS_STREAM_NAME, consumerName).catch(() => undefined); - } - } - - if (!reply.raw.writableEnded) { - reply.raw.end(); + closeStream(); } return reply; @@ -1035,21 +1027,14 @@ async function main(): Promise { decoded = input.codec.decode(input.msg.data); const startedAt = Date.now(); const queueWaitMs = safeDurationMs(decoded.queuedAt, startedAt); + const queueWaitTiming = typeof queueWaitMs === 'number' ? { queueWaitMs } : undefined; - const runningState: WorkerOperationState = { + await orchestrator.markRunning({ opId: decoded.opId, - opKey: decoded.opKey, - kind: decoded.kind, - jobId: decoded.jobId, - status: 'running', - queuedAt: decoded.queuedAt, startedAt, updatedAt: startedAt, - ...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}), - ...(latestProgress ? { progress: latestProgress } : {}), - }; - - await putOpState(runningState); + ...(queueWaitTiming ? { timing: queueWaitTiming } : {}), + }); await putJobState({ jobId: decoded.jobId, opId: decoded.opId, @@ -1059,7 +1044,7 @@ async function main(): Promise { timestamp: decoded.queuedAt, startedAt, updatedAt: startedAt, - ...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}), + ...(queueWaitTiming ? { timing: queueWaitTiming } : {}), ...(latestProgress ? { progress: latestProgress } : {}), }); app.log.info({ @@ -1072,20 +1057,21 @@ async function main(): Promise { }, 'job.started'); const persistRunningState = async (updatedAt: number): Promise => { - const runningOpState: WorkerOperationState = { - opId: decoded!.opId, - opKey: decoded!.opKey, - kind: decoded!.kind, - jobId: decoded!.jobId, - status: 'running', - queuedAt: decoded!.queuedAt, - startedAt, - updatedAt, - ...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}), - ...(latestProgress ? { progress: latestProgress } : {}), - }; - - await putOpState(runningOpState); + 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 } : {}), + }); + } await putJobState({ jobId: decoded!.jobId, opId: decoded!.opId, @@ -1095,7 +1081,7 @@ async function main(): Promise { timestamp: decoded!.queuedAt, startedAt, updatedAt, - ...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}), + ...(queueWaitTiming ? { timing: queueWaitTiming } : {}), ...(latestProgress ? { progress: latestProgress } : {}), }); }; @@ -1123,21 +1109,12 @@ async function main(): Promise { : undefined; const now = Date.now(); - const succeededState: WorkerOperationState = { + await orchestrator.markSucceeded({ opId: decoded.opId, - opKey: decoded.opKey, - kind: decoded.kind, - jobId: decoded.jobId, - status: 'succeeded', - queuedAt: decoded.queuedAt, - startedAt, - updatedAt: now, result: result as WhisperAlignJobResult | PdfLayoutJobResult, + updatedAt: now, ...(resultTiming ? { timing: resultTiming } : {}), - ...(latestProgress ? { progress: latestProgress } : {}), - }; - - await putOpState(succeededState); + }); await putJobState({ jobId: decoded.jobId, opId: decoded.opId, @@ -1185,22 +1162,30 @@ async function main(): Promise { if (decoded) { const now = Date.now(); const queueWaitMs = safeDurationMs(decoded.queuedAt, now); + const queueWaitTiming = typeof queueWaitMs === 'number' ? { queueWaitMs } : undefined; const status: WorkerJobState = hasRetriesLeft ? 'running' : 'failed'; - const opState: WorkerOperationState = { - opId: decoded.opId, - opKey: decoded.opKey, - kind: decoded.kind, - jobId: decoded.jobId, - status, - queuedAt: decoded.queuedAt, - updatedAt: now, - ...(status === 'failed' ? { error: { message } } : {}), - ...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}), - ...(latestProgress ? { progress: latestProgress } : {}), - }; + 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 putOpState(opState).catch((stateError) => { + await persistOpUpdate.catch((stateError) => { app.log.error({ worker: input.workerLabel, opId: decoded?.opId, @@ -1218,7 +1203,7 @@ async function main(): Promise { timestamp: decoded.queuedAt, updatedAt: now, ...(status === 'failed' ? { error: { message } } : {}), - ...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}), + ...(queueWaitTiming ? { timing: queueWaitTiming } : {}), ...(latestProgress ? { progress: latestProgress } : {}), }).catch((stateError) => { app.log.error({ diff --git a/src/app/api/documents/[id]/parsed/events/route.ts b/src/app/api/documents/[id]/parsed/events/route.ts index 536b509..abd8e0f 100644 --- a/src/app/api/documents/[id]/parsed/events/route.ts +++ b/src/app/api/documents/[id]/parsed/events/route.ts @@ -13,12 +13,12 @@ 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 { PdfParseProgress, PdfParseStatus } from '@/types/parsed-pdf'; +import { parseSseEventId, parseSsePayload } from '@openreader/compute-core'; import type { PdfLayoutJobResult, WorkerOperationEvent, WorkerOperationState } from '@openreader/compute-core/api-contracts'; export const dynamic = 'force-dynamic'; export const runtime = 'nodejs'; -const SSE_POLL_INTERVAL_MS = 1200; const SSE_KEEPALIVE_MS = 15_000; const SSE_RESYNC_INTERVAL_MS = 30_000; @@ -49,29 +49,6 @@ function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } -function parseSsePayload(frame: string): string | null { - const lines = frame.replace(/\r\n/g, '\n').split('\n'); - const dataLines: string[] = []; - for (const line of lines) { - if (!line.startsWith('data:')) continue; - dataLines.push(line.slice('data:'.length).trimStart()); - } - if (dataLines.length === 0) return null; - return dataLines.join('\n'); -} - -function parseSseEventId(frame: string): number | null { - const lines = frame.replace(/\r\n/g, '\n').split('\n'); - for (const line of lines) { - if (!line.startsWith('id:')) continue; - const value = Number(line.slice('id:'.length).trim()); - if (Number.isFinite(value) && value > 0) { - return Math.floor(value); - } - } - return null; -} - function mapWorkerStatusToParseStatus(status: WorkerOperationState['status']): PdfParseStatus { switch (status) { case 'queued': @@ -361,7 +338,7 @@ export async function GET(req: NextRequest, ctx: { params: Promise<{ id: string while (!closed) { if (!currentOpId) { - await sleep(SSE_POLL_INTERVAL_MS); + await sleep(SSE_RESYNC_INTERVAL_MS); const next = await syncFromDb({ documentId: id, storageUserId, diff --git a/src/lib/server/compute/worker.ts b/src/lib/server/compute/worker.ts index 10d4dc2..3be3897 100644 --- a/src/lib/server/compute/worker.ts +++ b/src/lib/server/compute/worker.ts @@ -1,5 +1,6 @@ import { createHash, randomUUID } from 'node:crypto'; import type { ComputeBackend, PdfLayoutInput, WhisperAlignInput, WhisperAlignResult } from '@/lib/server/compute/types'; +import { parseSseEventId, parseSsePayload } from '@openreader/compute-core'; import { getWorkerClientWaitTimeoutMs } from '@openreader/compute-core'; import type { PdfLayoutJobRequest, @@ -185,30 +186,6 @@ export function buildPdfOpKey(input: PdfLayoutInput): string { ].join('|'); } -function extractSsePayload(frame: string): string | null { - const dataLines: string[] = []; - const normalized = frame.replace(/\r\n/g, '\n'); - for (const line of normalized.split('\n')) { - if (line.startsWith('data:')) { - dataLines.push(line.slice('data:'.length).trimStart()); - } - } - if (dataLines.length === 0) return null; - return dataLines.join('\n'); -} - -function extractSseId(frame: string): number | null { - const normalized = frame.replace(/\r\n/g, '\n'); - for (const line of normalized.split('\n')) { - if (!line.startsWith('id:')) continue; - const value = Number(line.slice('id:'.length).trim()); - if (Number.isFinite(value) && value > 0) { - return Math.floor(value); - } - } - return null; -} - type RetryMeta = { attempt: number, maxAttempts: number, @@ -606,11 +583,11 @@ export class WorkerComputeBackend implements ComputeBackend { const frame = buffer.slice(0, frameEnd); buffer = buffer.slice(frameEnd + 2); - const eventId = extractSseId(frame); + const eventId = parseSseEventId(frame); if (eventId && eventId > 0) { lastEventId = eventId; } - const payload = extractSsePayload(frame); + const payload = parseSsePayload(frame); if (!payload) continue; let event: WorkerOperationEvent | WorkerOperationState; diff --git a/tests/unit/compute-control-plane.spec.ts b/tests/unit/compute-control-plane.spec.ts index 482faa2..7a9e338 100644 --- a/tests/unit/compute-control-plane.spec.ts +++ b/tests/unit/compute-control-plane.spec.ts @@ -1,10 +1,15 @@ import { expect, test } from '@playwright/test'; import type { WorkerOperationRequest, WorkerOperationState } from '../../compute/core/src/api-contracts'; import { + encodeSseFrame, + explainReplacementReason, InMemoryOperationEventStream, InMemoryOperationQueue, InMemoryOperationStateStore, OperationOrchestrator, + parseSseEventId, + parseSsePayload, + shouldReuseExistingOperation, } from '../../compute/core/src/control-plane'; function buildPdfLayoutRequest(opKey: string): WorkerOperationRequest { @@ -150,4 +155,47 @@ test.describe('compute control-plane', () => { expect(op1Events.map((event) => event.eventId)).toEqual([1, 2]); expect(op2Events.map((event) => event.eventId)).toEqual([1]); }); + + test('state-machine helpers return consistent reuse/replacement decisions', () => { + const current: WorkerOperationState = { + opId: 'op-1', + opKey: 'same-op-key', + kind: 'pdf_layout', + jobId: 'job-1', + status: 'running', + queuedAt: 1000, + updatedAt: 2000, + }; + + expect(shouldReuseExistingOperation({ + current, + requestKind: 'pdf_layout', + now: 2500, + opStaleMs: 1_000, + })).toBeTruthy(); + + expect(shouldReuseExistingOperation({ + current, + requestKind: 'pdf_layout', + now: 4005, + opStaleMs: 1_000, + })).toBeFalsy(); + + expect(explainReplacementReason({ + current, + requestKind: 'pdf_layout', + now: 4005, + opStaleMs: 1_000, + })).toBe('stale_running'); + }); + + test('sse helpers encode and decode frame payload and id', () => { + const frame = encodeSseFrame({ + id: 42, + event: 'snapshot', + data: { ok: true }, + }); + expect(parseSseEventId(frame)).toBe(42); + expect(parseSsePayload(frame)).toBe('{"ok":true}'); + }); }); diff --git a/tests/unit/compute-worker-control-plane-jetstream.spec.ts b/tests/unit/compute-worker-control-plane-jetstream.spec.ts index 699534c..641b78e 100644 --- a/tests/unit/compute-worker-control-plane-jetstream.spec.ts +++ b/tests/unit/compute-worker-control-plane-jetstream.spec.ts @@ -10,7 +10,7 @@ import { opStateKvKey, type KvEntryLike, type KvStoreLike, -} from '../../compute/worker/src/control-plane-jetstream'; +} from '../../compute/worker/src/control-plane/jetstream'; class FakeKvStore implements KvStoreLike { private readonly data = new Map(); @@ -64,6 +64,11 @@ class FakeKvStore implements KvStoreLike { class FakeJetStream { private seq = 0; readonly published: Array<{ subject: string; payload: unknown; seq: number }> = []; + readonly consumers = { + get: async () => { + throw new Error('not implemented in fake'); + }, + }; async publish(subject: string, data: Uint8Array): Promise<{ seq: number }> { this.seq += 1; @@ -127,6 +132,13 @@ test.describe('worker jetstream control-plane adapters', () => { }); const events = new JetStreamOperationEventStream({ getJs: async () => js, + getJsm: async () => ({ + consumers: { + add: async () => ({ name: 'noop' }), + delete: async () => true, + }, + }), + eventsStreamName: 'compute_events', }); await queue.enqueue({ @@ -157,7 +169,16 @@ test.describe('worker jetstream control-plane adapters', () => { const js = new FakeJetStream(); const store = new JetStreamOperationStateStore({ getKv: async () => kv }); - const events = new JetStreamOperationEventStream({ getJs: async () => js }); + const events = new JetStreamOperationEventStream({ + getJs: async () => js, + getJsm: async () => ({ + consumers: { + add: async () => ({ name: 'noop' }), + delete: async () => true, + }, + }), + eventsStreamName: 'compute_events', + }); const queue = new JetStreamOperationQueue({ getJs: async () => js, whisperSubject: 'jobs.whisper',