From 2e451bc93c0998b7d441ecef92fb90f166412fbf Mon Sep 17 00:00:00 2001 From: Roland Schlaefli Date: Tue, 30 Jun 2026 01:21:10 +0200 Subject: [PATCH] fix(codeapi-service): poll job results when queue events lag --- service/src/queue-wait.test.ts | 118 +++++++++++++ service/src/queue-wait.ts | 187 +++++++++++++++++++++ service/src/queue.ts | 3 +- service/src/service/programmatic-router.ts | 10 +- service/src/service/router.ts | 6 +- 5 files changed, 315 insertions(+), 9 deletions(-) create mode 100644 service/src/queue-wait.test.ts create mode 100644 service/src/queue-wait.ts diff --git a/service/src/queue-wait.test.ts b/service/src/queue-wait.test.ts new file mode 100644 index 0000000..070dd7f --- /dev/null +++ b/service/src/queue-wait.test.ts @@ -0,0 +1,118 @@ +import { describe, expect, test } from 'bun:test'; +import type { Job, Queue, QueueEvents } from 'bullmq'; +import { EventEmitter } from 'events'; +import { waitForJobFinished } from './queue-wait'; + +type TestResult = { ok: true; source: string }; + +function pending(): Promise { + return new Promise(() => undefined); +} + +function wait(ms: number): Promise { + return new Promise(resolve => setTimeout(resolve, ms)); +} + +function fakeJob(overrides: Partial>): Job { + return { + id: 'job-1', + waitUntilFinished: () => pending(), + ...overrides, + } as Job; +} + +function fakeQueue(getJob: Queue['getJob']): Queue { + return { getJob } as Queue; +} + +function fakeQueueEvents(): QueueEvents & EventEmitter { + return new EventEmitter() as QueueEvents & EventEmitter; +} + +describe('waitForJobFinished', () => { + test('uses the QueueEvents result when it arrives first', async () => { + const result = { ok: true, source: 'events' } satisfies TestResult; + const job = fakeJob({}); + const queue = fakeQueue(async () => ({ + returnvalue: result, + } as Job)); + const queueEvents = fakeQueueEvents(); + const waitForResult = waitForJobFinished(job, queue, queueEvents, 1000); + queueEvents.emit('completed', { jobId: 'job-1', returnvalue: result }); + + await expect(waitForResult).resolves.toEqual(result); + }); + + test('stops polling after the QueueEvents result wins', async () => { + const result = { ok: true, source: 'events' } satisfies TestResult; + let getJobCalls = 0; + const job = fakeJob({}); + const queue = fakeQueue(async () => { + getJobCalls += 1; + if (getJobCalls === 1) { + return { + getState: async () => 'active', + } as Job; + } + return { + returnvalue: result, + } as Job; + }); + const queueEvents = fakeQueueEvents(); + const waitForResult = waitForJobFinished(job, queue, queueEvents, 1000); + setTimeout(() => { + queueEvents.emit('completed', { jobId: 'job-1', returnvalue: result }); + }, 5); + + await expect(waitForResult).resolves.toEqual(result); + await wait(300); + expect(getJobCalls).toBe(2); + }); + + test('polls completed jobs when QueueEvents do not arrive', async () => { + const result = { ok: true, source: 'poll' } satisfies TestResult; + const job = fakeJob({}); + const queue = fakeQueue(async () => ({ + getState: async () => 'completed', + returnvalue: result, + } as Job)); + const queueEvents = fakeQueueEvents(); + + await expect(waitForJobFinished(job, queue, queueEvents, 1000)).resolves.toEqual(result); + }); + + test('rejects failed jobs found by polling', async () => { + const job = fakeJob({}); + const queue = fakeQueue(async () => ({ + getState: async () => 'failed', + failedReason: 'sandbox failed', + } as Job)); + const queueEvents = fakeQueueEvents(); + + await expect(waitForJobFinished(job, queue, queueEvents, 1000)).rejects.toThrow('sandbox failed'); + }); + + test('rejects missing jobs instead of polling until timeout', async () => { + const job = fakeJob({}); + const queue = fakeQueue(async () => undefined); + const queueEvents = fakeQueueEvents(); + + await expect(waitForJobFinished(job, queue, queueEvents, 1000)).rejects.toThrow( + 'Job job-1 no longer exists before reaching a terminal state', + ); + }); + + test('removes QueueEvents listeners after polling wins', async () => { + const result = { ok: true, source: 'poll' } satisfies TestResult; + const job = fakeJob({}); + const queue = fakeQueue(async () => ({ + getState: async () => 'completed', + returnvalue: result, + } as Job)); + const queueEvents = fakeQueueEvents(); + + await expect(waitForJobFinished(job, queue, queueEvents, 1000)).resolves.toEqual(result); + expect(queueEvents.listenerCount('completed')).toBe(0); + expect(queueEvents.listenerCount('failed')).toBe(0); + }); +}); diff --git a/service/src/queue-wait.ts b/service/src/queue-wait.ts new file mode 100644 index 0000000..883dec2 --- /dev/null +++ b/service/src/queue-wait.ts @@ -0,0 +1,187 @@ +import type { Job, Queue, QueueEvents } from 'bullmq'; + +const JOB_RESULT_POLL_INTERVAL_MS = 250; + +class TerminalJobError extends Error {} + +function wait(ms: number, signal?: AbortSignal): Promise { + if (signal?.aborted || ms <= 0) { + return Promise.resolve(); + } + + return new Promise(resolve => { + let timeout: ReturnType | undefined; + const done = () => { + if (timeout !== undefined) { + clearTimeout(timeout); + } + signal?.removeEventListener('abort', done); + resolve(); + }; + timeout = setTimeout(done, ms); + signal?.addEventListener('abort', done, { once: true }); + }); +} + +function throwIfAborted(jobId: string, signal?: AbortSignal): void { + if (signal?.aborted) { + throw new Error(`Job wait aborted before finishing (id=${jobId})`); + } +} + +function waitForJobEvent( + jobId: string, + queue: Queue, + queueEvents: QueueEvents, + timeoutMs: number, + signal?: AbortSignal, +): Promise { + return new Promise((resolve, reject) => { + let timeout: ReturnType | undefined; + let settled = false; + + const cleanup = () => { + if (timeout !== undefined) { + clearTimeout(timeout); + } + queueEvents.off('completed', onCompleted); + queueEvents.off('failed', onFailed); + signal?.removeEventListener('abort', onAbort); + }; + + const settle = (callback: () => void) => { + if (settled) return; + settled = true; + cleanup(); + callback(); + }; + + const settleWith = (promise: Promise) => { + if (settled) return; + settled = true; + cleanup(); + void promise.then(resolve, reject); + }; + + const onCompleted = (event: { jobId: string }) => { + if (event.jobId !== jobId) { + return; + } + settleWith((async () => { + const currentJob = await queue.getJob(jobId); + if (!currentJob) { + throw new Error(`Job ${jobId} no longer exists after completion event`); + } + return currentJob.returnvalue; + })()); + }; + const onFailed = (event: { jobId: string; failedReason?: string }) => { + if (event.jobId === jobId) { + settle(() => reject(new Error(event.failedReason || `Job ${jobId} failed`))); + } + }; + const onAbort = () => { + settle(() => reject(new Error(`Job event wait aborted before finishing (id=${jobId})`))); + }; + + if (signal?.aborted) { + onAbort(); + return; + } + + timeout = setTimeout(() => { + settle(() => reject(new Error(`Job wait execute timed out before receiving a finish event after ${timeoutMs}ms (id=${jobId})`))); + }, timeoutMs); + queueEvents.on('completed', onCompleted); + queueEvents.on('failed', onFailed); + signal?.addEventListener('abort', onAbort, { once: true }); + }); +} + +async function pollJobUntilFinished( + job: Job, + queue: Queue, + timeoutMs: number, + signal?: AbortSignal, +): Promise { + const jobId = job.id; + if (!jobId) { + throw new Error('Cannot poll BullMQ job without an id'); + } + + const deadline = Date.now() + timeoutMs; + let lastState = 'unknown'; + let lastError: unknown; + + while (Date.now() < deadline) { + throwIfAborted(jobId, signal); + + try { + const currentJob = await queue.getJob(jobId); + if (!currentJob) { + lastState = 'missing'; + throw new TerminalJobError(`Job ${jobId} no longer exists before reaching a terminal state`); + } + lastState = await currentJob.getState(); + if (lastState === 'completed') { + return currentJob.returnvalue; + } + if (lastState === 'failed') { + throw new TerminalJobError(currentJob.failedReason || `Job ${jobId} failed`); + } + } catch (error) { + if (error instanceof TerminalJobError) { + throw error; + } + lastError = error; + await wait(Math.min(JOB_RESULT_POLL_INTERVAL_MS, Math.max(0, deadline - Date.now())), signal); + continue; + } + + await wait(Math.min(JOB_RESULT_POLL_INTERVAL_MS, Math.max(0, deadline - Date.now())), signal); + } + + throwIfAborted(jobId, signal); + + const suffix = lastError instanceof Error ? `; last poll error: ${lastError.message}` : ''; + throw new Error( + `Job wait execute timed out before finishing, no terminal state observed after ${timeoutMs}ms (id=${jobId}, lastState=${lastState})${suffix}`, + ); +} + +async function waitForJobFinished( + job: Job, + queue: Queue, + queueEvents: QueueEvents, + timeoutMs: number, +): Promise { + const pollAbortController = new AbortController(); + const eventAbortController = new AbortController(); + const jobId = job.id; + if (!jobId) { + throw new Error('Cannot wait for BullMQ job without an id'); + } + + const eventWait = waitForJobEvent( + jobId, + queue, + queueEvents, + timeoutMs, + eventAbortController.signal, + ); + void eventWait.catch(() => undefined); + const pollWait = pollJobUntilFinished(job, queue, timeoutMs, pollAbortController.signal); + void pollWait.catch(() => undefined); + + try { + return await Promise.race([ + eventWait, + pollWait, + ]); + } finally { + pollAbortController.abort(); + eventAbortController.abort(); + } +} + +export { waitForJobFinished }; diff --git a/service/src/queue.ts b/service/src/queue.ts index c6f3226..a27e5fd 100644 --- a/service/src/queue.ts +++ b/service/src/queue.ts @@ -10,6 +10,7 @@ import { env } from './config'; import logger from './logger'; import { redisKeepAliveOptions } from './redis-options'; import { bullmqQueueJobs, registerBullmqQueueMetricsCollector } from './metrics'; +import { waitForJobFinished } from './queue-wait'; const MAX_RECONNECT_ATTEMPTS = 5; const RECONNECT_DELAY = 2000; @@ -109,4 +110,4 @@ registerBullmqQueueMetricsCollector(async () => { * BullMQ coordination objects. */ setMaxListeners(0, pyQueue, otherQueue, pyQueueEvents, otherQueueEvents); -export { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, connection }; +export { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, connection, waitForJobFinished }; diff --git a/service/src/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index 87bce41..06bb513 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -6,7 +6,7 @@ import type { Queue, QueueEvents } from 'bullmq'; import type * as t from '../types'; import { checkServiceStartUp, checkServiceShutDown } from '../lifecycle'; import { executionLimiter } from '../middleware/limits'; -import { pyQueue, pyQueueEvents, otherQueue, otherQueueEvents, connection } from '../queue'; +import { pyQueue, pyQueueEvents, otherQueue, otherQueueEvents, connection, waitForJobFinished } from '../queue'; import { createProgrammaticPayload, extractPendingFromStdout } from '../preamble'; import { findBashToolNameCollision } from '../preamble-bash'; import type { LCTool } from '../preamble'; @@ -405,13 +405,13 @@ async function runReplayIteration( egressGrantClaims: sandboxSecurity.egressGrantClaims, egressGrantToken: sandboxSecurity.egressGrantToken, }, { - removeOnComplete: { age: 60, count: 1 }, + removeOnComplete: { age: 60, count: 100 }, removeOnFail: { age: 180, count: 1 }, attempts: 1, }); jobsSubmitted.inc({ language }); - return job.waitUntilFinished(events, env.JOB_TIMEOUT); + return waitForJobFinished(job, queue, events, env.JOB_TIMEOUT); } function isSandboxRunSuccess(result: t.ExecuteResult): boolean { @@ -1372,7 +1372,7 @@ async function handleBlocking( egressGrantClaims: sandboxSecurity.egressGrantClaims, egressGrantToken: sandboxSecurity.egressGrantToken, }, { - removeOnComplete: { age: 60, count: 1 }, + removeOnComplete: { age: 60, count: 100 }, removeOnFail: { age: 180, count: 1 }, attempts: 1, jobId: session_id, @@ -1394,7 +1394,7 @@ async function handleBlocking( } }); - job.waitUntilFinished(pyQueueEvents, env.JOB_TIMEOUT) + waitForJobFinished(job, pyQueue, pyQueueEvents, env.JOB_TIMEOUT) .then(async (result) => { if (clientDisconnected) return; await setExecutionResult(execution_id, result); diff --git a/service/src/service/router.ts b/service/src/service/router.ts index 9d9fd37..ebc36c5 100644 --- a/service/src/service/router.ts +++ b/service/src/service/router.ts @@ -10,7 +10,7 @@ import { sessionAuth } from '../middleware/auth'; import { executionLimiter, uploadLimiter, downloadLimiter, fetchLimiter } from '../middleware/limits'; import { internalServiceHeaders } from '../internal-service-auth'; import { resolveSessionKey, resolveOutputBucketSessionKey, SessionKeyResolutionError, parseUploadSessionKeyInput, type SessionKeyInput } from '../session-key'; -import { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, connection } from '../queue'; +import { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, connection, waitForJobFinished } from '../queue'; import { sleep, getAxiosErrorDetails, publicExecutionFailure } from '../utils'; import { env, planLimits, resolveLanguage } from '../config'; import { createPayload } from '../payload'; @@ -226,7 +226,7 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res) }, { removeOnComplete: { age: 60, - count: 1, + count: 100, }, removeOnFail: { age: 180, @@ -251,7 +251,7 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res) 'messaging.system': 'bullmq', 'messaging.destination.name': queueName, 'codeapi.language': language, - }, () => job.waitUntilFinished(queueEvents, env.JOB_TIMEOUT), 'CONSUMER'); + }, () => waitForJobFinished(job, queue, queueEvents, env.JOB_TIMEOUT), 'CONSUMER'); if (!isSyntheticRequest) { logger.info('Execution completed', { session_id, user_id });