diff --git a/compute/worker/src/server.ts b/compute/worker/src/server.ts index 0f1de24..e545b83 100644 --- a/compute/worker/src/server.ts +++ b/compute/worker/src/server.ts @@ -61,7 +61,16 @@ const DOCUMENT_ID_REGEX = /^[a-f0-9]{64}$/i; const SAFE_NAMESPACE_REGEX = /^[a-zA-Z0-9._-]{1,128}$/; const WHISPER_MAX_DELIVER = 1; const NATS_API_TIMEOUT_MS = 60_000; +// Disconnect from NATS after this much continuous idle so the worker stops +// generating outbound traffic (pull polling + keepalive PINGs) and Railway can +// put it to sleep. Reconnect happens lazily on the next inbound request. +const IDLE_DISCONNECT_MS = 120_000; +const IDLE_CHECK_INTERVAL_MS = 5_000; +// Bounded pull window so consumer loops yield periodically and can be stopped +// cleanly when going idle, instead of blocking on a long-lived pull. +const PULL_EXPIRES_MS = 5_000; const REQUEST_STARTED_AT_MS_KEY = Symbol('request-started-at-ms'); +const REQUEST_COUNTED_KEY = Symbol('request-activity-counted'); interface QueuedJob { jobId: string; @@ -91,6 +100,15 @@ interface OpIndexEntry { opId: string; } +interface NatsSession { + nc: NatsConnection; + js: JetStreamClient; + jsm: JetStreamManager; + kv: Awaited>; + whisperConsumer: Consumer; + layoutConsumer: Consumer; +} + type JsonCodec = { encode(value: T): Uint8Array; decode(data: Uint8Array): T; @@ -449,17 +467,100 @@ async function main(): Promise { connectOpts.authenticator = credsAuthenticator(credsData); } - const nc: NatsConnection = await connect(connectOpts); - const js: JetStreamClient = jetstream(nc, { timeout: NATS_API_TIMEOUT_MS }); - const jsm: JetStreamManager = await jetstreamManager(nc, { timeout: NATS_API_TIMEOUT_MS }); + // Lazy NATS connection lifecycle. The worker connects on demand (first request + // that needs the queue/KV) and disconnects after IDLE_DISCONNECT_MS of full idle + // so it stops emitting outbound traffic and Railway can sleep it. Reconnect is + // transparent: any inbound /ops request both wakes the container and re-establishes + // the session via ensureConnected(). + let session: NatsSession | null = null; + let connecting: Promise | null = null; + let workerLoops: Promise[] = []; + let idleTimer: NodeJS.Timeout | null = null; + let stopping = false; + let loopStopRequested = false; - await ensureJetStreamResources(jsm, whisperTimeoutMs, pdfTimeoutMs, pdfAttempts, jobsStreamMaxBytes); + // Activity accounting feeding the idle detector. The worker is considered idle + // only when no HTTP request is in flight, no SSE stream is open, no job is + // processing, and nothing has happened for IDLE_DISCONNECT_MS. + let inFlightHttp = 0; + let activeSse = 0; + let inFlightJobs = 0; + let lastActivityAt = Date.now(); + const jobGate = new ConcurrencyGate(jobConcurrency); - const kv = await new Kvm(js).create(COMPUTE_STATE_BUCKET, { - history: 1, - ttl: COMPUTE_STATE_TTL_MS, - max_bytes: jobStatesMaxBytes, - }); + const markActivity = (): void => { + lastActivityAt = Date.now(); + }; + + function startIdleTimer(): void { + if (idleTimer) return; + idleTimer = setInterval(() => { + if (!session || stopping) return; + if (inFlightHttp > 0 || activeSse > 0 || inFlightJobs > 0) return; + if (Date.now() - lastActivityAt < IDLE_DISCONNECT_MS) return; + void disconnect('idle'); + }, IDLE_CHECK_INTERVAL_MS); + // Don't let the idle checker keep the process alive on its own. + idleTimer.unref?.(); + } + + async function disconnect(reason: string): Promise { + const current = session; + if (!current) return; + // Clear synchronously (before any await) so concurrent requests reconnect a + // fresh session instead of using the connection we're about to close. + session = null; + loopStopRequested = true; + if (idleTimer) { + clearInterval(idleTimer); + idleTimer = null; + } + try { + await current.nc.close(); + } catch { + // ignore close errors + } + await Promise.allSettled(workerLoops); + workerLoops = []; + loopStopRequested = false; + app.log.info({ reason }, 'nats disconnected'); + } + + async function ensureConnected(): Promise { + if (session) return session; + if (connecting) return connecting; + connecting = (async () => { + const nc: NatsConnection = await connect(connectOpts); + const js: JetStreamClient = jetstream(nc, { timeout: NATS_API_TIMEOUT_MS }); + const jsm: JetStreamManager = await jetstreamManager(nc, { timeout: NATS_API_TIMEOUT_MS }); + await ensureJetStreamResources(jsm, whisperTimeoutMs, pdfTimeoutMs, pdfAttempts, jobsStreamMaxBytes); + const kv = await new Kvm(js).create(COMPUTE_STATE_BUCKET, { + history: 1, + ttl: COMPUTE_STATE_TTL_MS, + max_bytes: jobStatesMaxBytes, + }); + const whisperConsumer = await js.consumers.get(JOBS_STREAM_NAME, WHISPER_CONSUMER_NAME); + const layoutConsumer = await js.consumers.get(JOBS_STREAM_NAME, LAYOUT_CONSUMER_NAME); + const next: NatsSession = { nc, js, jsm, kv, whisperConsumer, layoutConsumer }; + session = next; + markActivity(); + startWorkerLoops(next); + startIdleTimer(); + // Safety net: if the connection closes for any reason (network drop after + // exhausting reconnects, or our own disconnect), drop the stale session so + // the next request reconnects cleanly. + void nc.closed().then(() => { + if (session?.nc === nc) session = null; + }); + app.log.info('nats connected'); + return next; + })(); + try { + return await connecting; + } finally { + connecting = null; + } + } const s3 = buildS3Client(); const s3Bucket = requireEnv('S3_BUCKET'); @@ -523,16 +624,19 @@ async function main(): Promise { const jobStateCodec = createJsonCodec>(); const putOpState = async (state: WorkerOperationState): Promise => { + const { kv } = await ensureConnected(); await kv.put(opStateKvKey(state.opId), opStateCodec.encode(state)); }; const getOpState = async (opId: string): Promise | 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): Promise => { + const { kv } = await ensureConnected(); await kv.put(jobStateKvKey(state.jobId), jobStateCodec.encode(state)); }; @@ -540,6 +644,7 @@ async function main(): Promise { op: WorkerOperationState, payload: WhisperAlignJobRequest | PdfLayoutJobRequest, ): Promise => { + const { js } = await ensureConnected(); if (op.kind === 'whisper_align') { await js.publish(WHISPER_JOBS_SUBJECT, whisperJobCodec.encode({ jobId: op.jobId, @@ -567,6 +672,7 @@ async function main(): Promise { ): Promise> => { const opKey = req.opKey.trim(); const indexKey = opIndexKvKey(opKey); + const { kv } = await ensureConnected(); for (let attemptNo = 0; attemptNo < 10; attemptNo += 1) { const indexEntry = await kv.get(indexKey); @@ -698,9 +804,23 @@ async function main(): Promise { throw new Error('Unable to reserve operation after repeated CAS conflicts'); }; + const releaseHttp = (request: FastifyRequest): void => { + const counted = request as FastifyRequest & { [REQUEST_COUNTED_KEY]?: boolean }; + if (!counted[REQUEST_COUNTED_KEY]) return; + counted[REQUEST_COUNTED_KEY] = false; + inFlightHttp = Math.max(0, inFlightHttp - 1); + markActivity(); + }; + app.addHook('onRequest', async (request, reply) => { const path = requestPath(request); (request as FastifyRequest & { [REQUEST_STARTED_AT_MS_KEY]?: number })[REQUEST_STARTED_AT_MS_KEY] = Date.now(); + // Count every request as in-flight activity so the idle detector never + // disconnects mid-request. Released in onResponse, or manually after hijack + // for SSE streams (where onResponse does not fire). + (request as FastifyRequest & { [REQUEST_COUNTED_KEY]?: boolean })[REQUEST_COUNTED_KEY] = true; + inFlightHttp += 1; + markActivity(); if (isHealthPath(path)) return; if (!isAuthed(request, workerToken)) { return reply.code(401).send({ error: 'Unauthorized' }); @@ -709,6 +829,7 @@ async function main(): Promise { }); app.addHook('onResponse', async (request, reply) => { + releaseHttp(request); const path = requestPath(request); if (isHealthPath(path)) return; @@ -725,18 +846,11 @@ async function main(): Promise { app.get('/health/live', async () => ({ ok: true })); - app.get('/health/ready', async (_request, reply) => { - try { - await nc.flush(); - return { ok: true }; - } catch (error) { - reply.code(503); - return { - ok: false, - error: toErrorMessage(error), - }; - } - }); + // Reports readiness without forcing a NATS round-trip. Probing NATS here would + // reconnect (and keep) the connection open, defeating idle sleep, so we only + // report the current connection state. The worker reconnects lazily on the next + // /ops request regardless of what this returns. + app.get('/health/ready', async () => ({ ok: true, natsConnected: session !== null })); app.post('/ops', async (request, reply) => { const parsed = operationCreateSchema.safeParse(request.body); @@ -783,6 +897,11 @@ async function main(): Promise { } reply.hijack(); + // onResponse will not fire for a hijacked reply, so release the HTTP in-flight + // count here and track the long-lived stream via activeSse instead. + releaseHttp(request); + activeSse += 1; + markActivity(); reply.raw.setHeader('Content-Type', 'text/event-stream; charset=utf-8'); reply.raw.setHeader('Cache-Control', 'no-cache, no-transform'); reply.raw.setHeader('Connection', 'keep-alive'); @@ -796,6 +915,8 @@ async function main(): Promise { let closed = false; request.raw.on('close', () => { closed = true; + activeSse = Math.max(0, activeSse - 1); + markActivity(); }); try { @@ -828,12 +949,6 @@ async function main(): Promise { return reply; }); - const whisperConsumer = await js.consumers.get(JOBS_STREAM_NAME, WHISPER_CONSUMER_NAME); - const layoutConsumer = await js.consumers.get(JOBS_STREAM_NAME, LAYOUT_CONSUMER_NAME); - const jobGate = new ConcurrencyGate(jobConcurrency); - - let stopping = false; - const runWhisper = async ( payload: WhisperAlignJobRequest, queueWaitMs: number, @@ -1140,6 +1255,7 @@ async function main(): Promise { } async function createWorkerLoop(input: { + owner: NatsSession; consumer: Consumer; codec: JsonCodec>; run: ( @@ -1149,22 +1265,29 @@ async function main(): Promise { ) => Promise; workerLabel: string; }): Promise { - while (!stopping) { + // Exit when the loop's connection is no longer the active session (idle + // disconnect, unexpected close, or replaced by a reconnect). + const detached = (): boolean => stopping || loopStopRequested || session !== input.owner; + while (!detached()) { let msg: JsMsg | null = null; try { try { - msg = await input.consumer.next(); + // Bounded pull so the loop yields periodically and exits promptly when + // the session is torn down for idle (nc.close() rejects the pending pull). + msg = await input.consumer.next({ expires: PULL_EXPIRES_MS }); } catch (error) { - if (stopping) return; + if (detached()) return; app.log.error({ error: toErrorMessage(error), worker: input.workerLabel }, 'worker pull failed'); await sleep(LOOP_ERROR_BACKOFF_MS); continue; } + // An empty pull is not activity; let the idle window advance. if (!msg) continue; + markActivity(); + inFlightJobs += 1; await jobGate.acquire(); - if (stopping) { - jobGate.release(); + if (detached()) { return; } await processMessage({ @@ -1176,37 +1299,60 @@ async function main(): Promise { } finally { if (msg) { jobGate.release(); + inFlightJobs = Math.max(0, inFlightJobs - 1); + markActivity(); } } } } - const workerLoops: Promise[] = []; - - for (let i = 0; i < jobConcurrency; i += 1) { - workerLoops.push(createWorkerLoop({ - consumer: whisperConsumer, - codec: whisperJobCodec, - run: runWhisper, - workerLabel: `whisper-${i + 1}`, - })); - } - - for (let i = 0; i < jobConcurrency; i += 1) { - workerLoops.push(createWorkerLoop({ - consumer: layoutConsumer, - codec: layoutJobCodec, - run: runLayout, - workerLabel: `layout-${i + 1}`, - })); + function startWorkerLoops(active: NatsSession): void { + // Always starts a fresh set bound to the new session. Any loops from a prior + // session self-terminate once they observe session !== their owner. + workerLoops = []; + loopStopRequested = false; + for (let i = 0; i < jobConcurrency; i += 1) { + workerLoops.push(createWorkerLoop({ + owner: active, + consumer: active.whisperConsumer, + codec: whisperJobCodec, + run: runWhisper, + workerLabel: `whisper-${i + 1}`, + })); + } + for (let i = 0; i < jobConcurrency; i += 1) { + workerLoops.push(createWorkerLoop({ + owner: active, + consumer: active.layoutConsumer, + codec: layoutJobCodec, + run: runLayout, + workerLabel: `layout-${i + 1}`, + })); + } } const close = async (): Promise => { if (stopping) return; stopping = true; + if (idleTimer) { + clearInterval(idleTimer); + idleTimer = null; + } await app.close(); await Promise.allSettled(workerLoops); - await nc.drain(); + const current = session; + session = null; + if (current) { + try { + await current.nc.drain(); + } catch { + try { + await current.nc.close(); + } catch { + // ignore close errors + } + } + } }; process.once('SIGINT', () => { diff --git a/docs-site/docs/deploy/compute-worker.md b/docs-site/docs/deploy/compute-worker.md index 26f3d44..13296d4 100644 --- a/docs-site/docs/deploy/compute-worker.md +++ b/docs-site/docs/deploy/compute-worker.md @@ -92,10 +92,25 @@ Set the same value on app + worker envs. - Non-exposed embedded `weed mini` is not supported with external worker mode. - Protect `COMPUTE_WORKER_TOKEN` and avoid exposing worker routes publicly without auth. +## Railway sleep & idle behavior + +The worker connects to NATS lazily (on the first request needing the queue/KV) and +disconnects after **120s** of full idle — no in-flight request, SSE stream, job, or +queued work. This stops outbound pull polling and keepalive PINGs so Railway can sleep +it; the next inbound request transparently reconnects, re-ensures the stream/consumers +and KV (idempotent), and drains anything pending. No separate mode, no extra env vars, +and the `/ops*` contract is unchanged. + +Caveats: inbound HTTP is the wake signal (in OpenReader the app server only enqueues via +`POST /ops`, so this is always satisfied); a continuous external `/health/*` probe keeps +it awake and prevents sleep; and the first request after a cold start re-runs model +prewarm, so it's slower. + ## Health endpoints -- `GET /health/live` -- `GET /health/ready` +- `GET /health/live` — liveness; always returns `{ ok: true }`. +- `GET /health/ready` — returns `{ ok: true, natsConnected }`. It does not probe NATS (that + would reconnect and prevent idle sleep); `natsConnected` just reflects the current session. ## Synadia Cloud + Railway Setup (Complete Guide)