From ef063445f31ea56318c0f09e059bf2a31ea63f8f Mon Sep 17 00:00:00 2001 From: Joe Clark Date: Wed, 10 Sep 2025 13:58:27 +0100 Subject: [PATCH 1/7] worker: add claim tests --- packages/ws-worker/src/api/claim.ts | 3 + packages/ws-worker/test/api/claim.test.ts | 70 ++++++++++++++++++++++- 2 files changed, 72 insertions(+), 1 deletion(-) diff --git a/packages/ws-worker/src/api/claim.ts b/packages/ws-worker/src/api/claim.ts index 8cd782d24..5244bc660 100644 --- a/packages/ws-worker/src/api/claim.ts +++ b/packages/ws-worker/src/api/claim.ts @@ -48,11 +48,14 @@ const claim = ( const podName = NAME ? `[${NAME}] ` : ''; const activeWorkers = Object.keys(app.workflows).length; + if (activeWorkers >= maxWorkers) { // Important: stop the workloop so that we don't try and claim any more app.workloop?.stop(`server at capacity (${activeWorkers}/${maxWorkers})`); return reject(new ClaimError('Server at capacity')); } + // TODO if activeWorkers + activeClaims > capacity, silently abort + // This can happen in response to the work-available event if (!app.queueChannel) { logger.warn('skipping claim attempt: websocket unavailable'); diff --git a/packages/ws-worker/test/api/claim.test.ts b/packages/ws-worker/test/api/claim.test.ts index fb9b36d2a..524fac07b 100644 --- a/packages/ws-worker/test/api/claim.test.ts +++ b/packages/ws-worker/test/api/claim.test.ts @@ -2,8 +2,12 @@ import test from 'ava'; import * as jose from 'jose'; import crypto from 'node:crypto'; -import { verifyToken } from '../../src/api/claim'; +import claim, { verifyToken } from '../../src/api/claim'; import { generateKeys } from '@openfn/lightning-mock/src/util'; +import { createMockLogger } from '@openfn/logger'; +import { ServerApp } from '../../src/server'; +import { mockChannel } from '../../src/mock/sockets'; +import { CLAIM } from '../../src'; let keys = { public: '.', private: '.' }; @@ -117,3 +121,67 @@ test('verifyToken should accept a token with NBF exactly 2 seconds in future (us '✅ Token that is 2 seconds off is accepted within 5-second clock tolerance' ); }); + +const createMockApp = (opts: any) => { + const { onClaim = () => {}, onExecute = () => {}, workflows = {} } = opts; + + const channel = mockChannel({ + [CLAIM]: (args) => { + return onClaim(args); + }, + }); + + return { + workloop: undefined, // should be safe + workflows, + queueChannel: channel, + execute: (...args) => { + onExecute(...args); + }, + } as ServerApp; +}; +const logger = createMockLogger(); + +test.todo('claim: should do nothing if no runs returned'); + +test('claim: should call execute for a single run', async (t) => { + let executeArgs; + const onClaim = () => ({ runs: [{ id: 'abc' }] }); + const onExecute = (...args) => { + executeArgs = args; + }; + + const options = { maxWorkers: 1 }; + const app = createMockApp({ + onClaim, + onExecute, + }); + + await claim(app, logger, options); + t.deepEqual(executeArgs[0], { id: 'abc' }); +}); + +test.only('should not claim if worker is at capacity', async (t) => { + const options = { maxWorkers: 1 }; + + const app = createMockApp({ + workflows: { + a: true, + }, + }); + + await t.throwsAsync(() => claim(app, logger, options), { + message: 'Server at capacity', + }); +}); + +test('should not claim if open claims exceeds capacity', () => {}); + +test('should not claim if open claims + active runs exceeds capacity', () => {}); + +test.todo('should handle multiple runs'); +test.todo('claim payload should have a demand'); +test.todo('claim payload should include a worker name'); +test.todo('should stop the workloop if at capacity'); +// TODO I'd rather return true/false really and let the backoff itself decide whether to throw or not +test.todo('should throw if there are no runs available (to trigger backoff)'); From 9c6383f9422efd3f3417ca0f9ac93eee8d36aa43 Mon Sep 17 00:00:00 2001 From: Joe Clark Date: Wed, 10 Sep 2025 16:22:55 +0100 Subject: [PATCH 2/7] implement open claim handling and add tests --- packages/ws-worker/src/api/claim.ts | 33 +++- packages/ws-worker/src/server.ts | 2 + packages/ws-worker/test/api/claim.test.ts | 189 +++++++++++++++++++++- 3 files changed, 216 insertions(+), 8 deletions(-) diff --git a/packages/ws-worker/src/api/claim.ts b/packages/ws-worker/src/api/claim.ts index 5244bc660..47ae00f16 100644 --- a/packages/ws-worker/src/api/claim.ts +++ b/packages/ws-worker/src/api/claim.ts @@ -24,6 +24,7 @@ export const verifyToken = async (token: string, publicKey: string) => { type ClaimOptions = { maxWorkers?: number; + demand?: number; }; // used to report the pod name in logging, for tracking @@ -38,24 +39,41 @@ class ClaimError extends Error { } } +let claimIdGen = 0; + +export const resetClaimIdGen = () => { + claimIdGen = 0; +}; + const claim = ( app: ServerApp, logger: Logger = mockLogger, options: ClaimOptions = {} ) => { return new Promise((resolve, reject) => { - const { maxWorkers = 5 } = options; + const { maxWorkers = 5, demand = 1 } = options; const podName = NAME ? `[${NAME}] ` : ''; const activeWorkers = Object.keys(app.workflows).length; + const pendingClaims = Object.values(app.openClaims).reduce( + (a, b) => a + b, + 0 + ); + if (activeWorkers >= maxWorkers) { // Important: stop the workloop so that we don't try and claim any more app.workloop?.stop(`server at capacity (${activeWorkers}/${maxWorkers})`); return reject(new ClaimError('Server at capacity')); + } else if (activeWorkers + pendingClaims >= maxWorkers) { + // There are active claims which haven't yet been fulfilled + // This can happen in response to the work-available event + app.workloop?.stop( + `server at capacity (${activeWorkers}/${maxWorkers}, ${pendingClaims} pending)` + ); + return reject(new ClaimError('Server at capacity')); } // TODO if activeWorkers + activeClaims > capacity, silently abort - // This can happen in response to the work-available event if (!app.queueChannel) { logger.warn('skipping claim attempt: websocket unavailable'); @@ -70,15 +88,20 @@ const claim = ( return reject(e); } + const claimId = ++claimIdGen; + + app.openClaims[claimId] = demand; + logger.debug(`requesting run (capacity ${activeWorkers}/${maxWorkers})`); const start = Date.now(); app.queueChannel .push(CLAIM, { - demand: 1, + demand, worker_name: NAME || null, }) .receive('ok', ({ runs }: ClaimReply) => { + delete app.openClaims[claimId]; const duration = Date.now() - start; logger.debug( `${podName}claimed ${runs.length} runs in ${duration}ms (${ @@ -111,15 +134,19 @@ const claim = ( logger.debug(`${podName} starting run ${run.id}`); app.execute(run); + console.log('> done'); resolve(); }); }) // TODO need implementations for both of these really // What do we do if we fail to join the worker channel? .receive('error', (e) => { + delete app.openClaims[claimId]; logger.error('Error on claim', e); + reject(new Error('claim error')); }) .receive('timeout', () => { + delete app.openClaims[claimId]; logger.error('TIMEOUT on claim. Runs may be lost.'); reject(new Error('timeout')); }); diff --git a/packages/ws-worker/src/server.ts b/packages/ws-worker/src/server.ts index d7fa9db14..795df5283 100644 --- a/packages/ws-worker/src/server.ts +++ b/packages/ws-worker/src/server.ts @@ -60,6 +60,7 @@ export interface ServerApp extends Koa { socket?: any; queueChannel?: Channel; workflows: Record; + openClaims: Record; destroyed: boolean; events: EventEmitter; server: Server; @@ -221,6 +222,7 @@ function createServer(engine: RuntimeEngine, options: ServerOptions = {}) { }) ); + app.openClaims = {}; app.workflows = {}; app.destroyed = false; diff --git a/packages/ws-worker/test/api/claim.test.ts b/packages/ws-worker/test/api/claim.test.ts index 524fac07b..824d7efe3 100644 --- a/packages/ws-worker/test/api/claim.test.ts +++ b/packages/ws-worker/test/api/claim.test.ts @@ -2,7 +2,7 @@ import test from 'ava'; import * as jose from 'jose'; import crypto from 'node:crypto'; -import claim, { verifyToken } from '../../src/api/claim'; +import claim, { resetClaimIdGen, verifyToken } from '../../src/api/claim'; import { generateKeys } from '@openfn/lightning-mock/src/util'; import { createMockLogger } from '@openfn/logger'; import { ServerApp } from '../../src/server'; @@ -15,6 +15,10 @@ test.before(async () => { keys = await generateKeys(); }); +test.beforeEach(() => { + resetClaimIdGen(); +}); + // Helper function to generate a token with custom iat and nbf offsets const generateToken = async ( runId: string, @@ -123,7 +127,11 @@ test('verifyToken should accept a token with NBF exactly 2 seconds in future (us }); const createMockApp = (opts: any) => { - const { onClaim = () => {}, onExecute = () => {}, workflows = {} } = opts; + const { + onClaim = () => ({ runs: [] }), + onExecute = () => {}, + workflows = {}, + } = opts; const channel = mockChannel({ [CLAIM]: (args) => { @@ -133,6 +141,7 @@ const createMockApp = (opts: any) => { return { workloop: undefined, // should be safe + openClaims: {}, workflows, queueChannel: channel, execute: (...args) => { @@ -161,7 +170,7 @@ test('claim: should call execute for a single run', async (t) => { t.deepEqual(executeArgs[0], { id: 'abc' }); }); -test.only('should not claim if worker is at capacity', async (t) => { +test('should not claim if worker is at capacity', async (t) => { const options = { maxWorkers: 1 }; const app = createMockApp({ @@ -175,9 +184,179 @@ test.only('should not claim if worker is at capacity', async (t) => { }); }); -test('should not claim if open claims exceeds capacity', () => {}); +test('should mark a claim when in flight', async (t) => { + const options = {}; + + const app = createMockApp({ + workflows: {}, + }); + + let claimPromise = claim(app, logger, options); + + t.is(app.openClaims['1'], 1); + t.is(Object.keys(app.openClaims).length, 1); + + await t.throwsAsync(claimPromise, { + message: 'No runs returned', + }); +}); + +test('should remove an open claim when completed', async (t) => { + const options = {}; + + const app = createMockApp({ + workflows: {}, + }); + + await t.throwsAsync(() => claim(app, logger, options), { + message: 'No runs returned', + }); + + t.falsy(app.openClaims['1']); + t.is(Object.keys(app.openClaims).length, 0); +}); + +test('should remove an open claim on error', async (t) => { + const options = {}; + + const app = createMockApp({ + workflows: {}, + onClaim: () => { + throw {}; + }, + }); + + await t.throwsAsync(() => claim(app, logger, options), { + message: 'claim error', + }); + + t.falsy(app.openClaims['1']); + t.is(Object.keys(app.openClaims).length, 0); +}); + +// TODO not really sure how to check this +test.skip('should remove an open claim on timeout', async (t) => { + const options = {}; + + const app = createMockApp({ + workflows: {}, + onClaim: async () => { + await new Promise((resolve) => setTimeout(resolve, 1000 * 60 * 60)); + }, + }); + + await t.throwsAsync(() => claim(app, logger, options), { + message: 'timeout', + }); + + t.falsy(app.openClaims['1']); + t.is(Object.keys(app.openClaims).length, 0); +}); + +test('should mark a claim when in flight with demand: 2', async (t) => { + const options = { + demand: 2, + }; + + const app = createMockApp({ + workflows: { + a: true, + }, + }); + + let claimPromise = claim(app, logger, options); + + t.is(app.openClaims['1'], 2); + t.is(Object.keys(app.openClaims).length, 1); + + await t.throwsAsync(claimPromise, { + message: 'No runs returned', + }); +}); + +test('should not claim if open claims exceeds capacity', async (t) => { + let didStopWorkloop = false; + + const options = { + maxWorkers: 1, + demand: 1, + }; + + const app = createMockApp({ + workflows: {}, + // Slow claim event + onClaim: () => + new Promise((resolve) => { + setTimeout(resolve({ runs: [] }), 100); + }), + }); + + app.workloop = { + stop: () => { + didStopWorkloop = true; + }, + }; + + app.execute = ({ id }) => { + app.workflows[id] = true; + }; + + // first claim should be fine + let claimPromise = claim(app, logger, options); + + // second claim should error and stop the loop actually + await t.throwsAsync(() => claim(app, logger, options), { + message: 'Server at capacity', + }); + t.true(didStopWorkloop); + + // The prior claim should not have counted for anything + t.is(Object.keys(app.workflows).length, 0); + + // wait for the original return to finish + await t.throwsAsync(claimPromise, { + message: 'No runs returned', + }); +}); + +test('should not claim if open claims + active runs exceeds capacity', async (t) => { + const options = { + maxWorkers: 2, + demand: 1, + }; + + const app = createMockApp({ + workflows: { + // pretend a run is in progress + a: true, + }, + // Slow claim event + onClaim: () => + new Promise((resolve) => { + setTimeout(resolve({ runs: [] }), 100); + }), + }); + + app.execute = ({ id }) => { + app.workflows[id] = true; + }; + + // first claim should be fine + let claimPromise = claim(app, logger, options); -test('should not claim if open claims + active runs exceeds capacity', () => {}); + // second claim should error + await t.throwsAsync(() => claim(app, logger, options), { + message: 'Server at capacity', + }); + + // The prior claim should not have counted for anything + t.is(Object.keys(app.workflows).length, 1); + + // wait for the original return to finish + await t.throwsAsync(claimPromise, { + message: 'No runs returned', + }); +}); test.todo('should handle multiple runs'); test.todo('claim payload should have a demand'); From 1afadbc5777fd7e7379c9cc51eab6c71966b6e45 Mon Sep 17 00:00:00 2001 From: Joe Clark Date: Thu, 11 Sep 2025 09:06:23 +0100 Subject: [PATCH 3/7] add test --- packages/ws-worker/src/api/claim.ts | 1 - packages/ws-worker/test/wake-up.test.ts | 71 ++++++++++++++++++++++++- 2 files changed, 69 insertions(+), 3 deletions(-) diff --git a/packages/ws-worker/src/api/claim.ts b/packages/ws-worker/src/api/claim.ts index 47ae00f16..8b440b792 100644 --- a/packages/ws-worker/src/api/claim.ts +++ b/packages/ws-worker/src/api/claim.ts @@ -134,7 +134,6 @@ const claim = ( logger.debug(`${podName} starting run ${run.id}`); app.execute(run); - console.log('> done'); resolve(); }); }) diff --git a/packages/ws-worker/test/wake-up.test.ts b/packages/ws-worker/test/wake-up.test.ts index 44932a148..963cd1180 100644 --- a/packages/ws-worker/test/wake-up.test.ts +++ b/packages/ws-worker/test/wake-up.test.ts @@ -23,12 +23,13 @@ const urls = { const ONE_HOUR = 1000 * 60 * 60; let rollingRunId = 1; -const getRun = (): LightningPlan => + +const getRun = (body = 'fn(s => s)'): LightningPlan => ({ id: `a${++rollingRunId}`, jobs: [ { - body: 'fn(s => s)', + body, }, ], } as LightningPlan); @@ -37,6 +38,7 @@ test.before(async () => { engine = await createMockRTE(); lng = createLightningServer({ port: 7654, + socketDelay: 10, }); createWorkerServer(engine, { @@ -106,3 +108,68 @@ test('should initiate a claim when /run?wakeup=true', async (t) => { }); }); }); + +// This tests creates two runs, call wake_up on them both, +// and ensures they are executed sequentially +// This doesn't feel like a great test, but if the pending claims validation +// is removed from claim.ts, it fails. So it is kinda working +test.only('should not claim beyond capacity', async (t) => { + t.plan(2); + let claimCount = 0; + return new Promise(async (done) => { + const run_1 = getRun( + 's => new Promise(resolve => setTimeout(resolve, 200))' + ); + const run_2 = getRun( + 's => new Promise(resolve => setTimeout(resolve, 200))' + ); + let run_1_finished = false; + + lng.onSocketEvent(e.RUN_COMPLETE, run_1.id, (evt) => { + if (evt.runId === run_1.id) { + t.log(' RUN 1 complete'); + run_1_finished = true; + } + }); + + lng.onSocketEvent(e.RUN_START, run_2.id, (evt) => { + if (evt.runId === run_2.id) { + t.log(' RUN 2 start'); + // Ensure that run 1 is finished when run 2 starts + t.true(run_1_finished); + } + }); + + // TODO there's a terrible bug in the mock and these events don't bind right + lng.onSocketEvent(e.RUN_COMPLETE, run_2.id, (evt) => { + if (evt.runId === run_2.id) { + t.log(' RUN 2 complete'); + // Only 2 claims should be made + t.is(claimCount, 2); + done(); + } + }); + + lng.onSocketEvent(e.CLAIM, run_1.id, () => { + claimCount++; + }); + + // Send two runs up at the same time + fetch(`${urls.lngServer}/run?wakeup=true`, { + method: 'POST', + body: JSON.stringify(run_1), + headers: { + Accept: 'application/json', + 'Content-Type': 'application/json', + }, + }); + fetch(`${urls.lngServer}/run?wakeup=true`, { + method: 'POST', + body: JSON.stringify(run_2), + headers: { + Accept: 'application/json', + 'Content-Type': 'application/json', + }, + }); + }); +}); From 56e3326367a7610dc622e7f76e4e541cab047ad2 Mon Sep 17 00:00:00 2001 From: Joe Clark Date: Thu, 11 Sep 2025 09:09:39 +0100 Subject: [PATCH 4/7] fix types --- packages/lightning-mock/src/socket-server.ts | 6 ++++-- packages/ws-worker/test/api/claim.test.ts | 10 +++++++--- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/packages/lightning-mock/src/socket-server.ts b/packages/lightning-mock/src/socket-server.ts index cd30c4549..4a0746ba4 100644 --- a/packages/lightning-mock/src/socket-server.ts +++ b/packages/lightning-mock/src/socket-server.ts @@ -190,8 +190,10 @@ function createServer({ topic, payload: stringify(payload), // TODO do we stringify this? All of it? }); - // @ts-ignore - ws.send(evt); + setTimeout(() => { + // @ts-ignore + ws.send(evt); + }, socketDelay); }; ws.on('message', async function (data: string) { diff --git a/packages/ws-worker/test/api/claim.test.ts b/packages/ws-worker/test/api/claim.test.ts index 824d7efe3..4bd197993 100644 --- a/packages/ws-worker/test/api/claim.test.ts +++ b/packages/ws-worker/test/api/claim.test.ts @@ -154,9 +154,9 @@ const logger = createMockLogger(); test.todo('claim: should do nothing if no runs returned'); test('claim: should call execute for a single run', async (t) => { - let executeArgs; + let executeArgs: any[] = []; const onClaim = () => ({ runs: [{ id: 'abc' }] }); - const onExecute = (...args) => { + const onExecute = (...args: any[]) => { executeArgs = args; }; @@ -287,6 +287,7 @@ test('should not claim if open claims exceeds capacity', async (t) => { // Slow claim event onClaim: () => new Promise((resolve) => { + // @ts-ignore setTimeout(resolve({ runs: [] }), 100); }), }); @@ -295,8 +296,9 @@ test('should not claim if open claims exceeds capacity', async (t) => { stop: () => { didStopWorkloop = true; }, - }; + } as any; + // @ts-ignore app.execute = ({ id }) => { app.workflows[id] = true; }; @@ -333,10 +335,12 @@ test('should not claim if open claims + active runs exceeds capacity', async (t) // Slow claim event onClaim: () => new Promise((resolve) => { + // @ts-ignore setTimeout(resolve({ runs: [] }), 100); }), }); + // @ts-ignore app.execute = ({ id }) => { app.workflows[id] = true; }; From 9f8ebaa81c6312361484f0c899d2fa28372e89f5 Mon Sep 17 00:00:00 2001 From: Joe Clark Date: Thu, 11 Sep 2025 12:12:55 +0100 Subject: [PATCH 5/7] fix tests --- packages/ws-worker/src/api/claim.ts | 3 ++- packages/ws-worker/test/api/workloop.test.ts | 6 +++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/packages/ws-worker/src/api/claim.ts b/packages/ws-worker/src/api/claim.ts index 8b440b792..7a64e0d84 100644 --- a/packages/ws-worker/src/api/claim.ts +++ b/packages/ws-worker/src/api/claim.ts @@ -51,6 +51,8 @@ const claim = ( options: ClaimOptions = {} ) => { return new Promise((resolve, reject) => { + app.openClaims ??= {}; + const { maxWorkers = 5, demand = 1 } = options; const podName = NAME ? `[${NAME}] ` : ''; @@ -73,7 +75,6 @@ const claim = ( ); return reject(new ClaimError('Server at capacity')); } - // TODO if activeWorkers + activeClaims > capacity, silently abort if (!app.queueChannel) { logger.warn('skipping claim attempt: websocket unavailable'); diff --git a/packages/ws-worker/test/api/workloop.test.ts b/packages/ws-worker/test/api/workloop.test.ts index f67e5afb4..6a920e9af 100644 --- a/packages/ws-worker/test/api/workloop.test.ts +++ b/packages/ws-worker/test/api/workloop.test.ts @@ -16,7 +16,7 @@ test.afterEach(() => { test('workloop can be cancelled', async (t) => { let count = 0; - + const app = { workflows: {}, queueChannel: mockChannel({ @@ -30,12 +30,12 @@ test('workloop can be cancelled', async (t) => { }; workloop = startWorkloop(app as any, logger, 1, 1); - t.false(workloop.isStopped()) + t.false(workloop.isStopped()); await sleep(100); // A quirk of how cancel works is that the loop will be called a few times t.true(count <= 5); - t.true(workloop.isStopped()) + t.true(workloop.isStopped()); }); test('workloop sends the runs:claim event', (t) => { From 3a33557b3e626235311f300158e9a723f7b18e81 Mon Sep 17 00:00:00 2001 From: Joe Clark Date: Thu, 11 Sep 2025 13:40:39 +0100 Subject: [PATCH 6/7] changeset --- .changeset/rude-squids-mix.md | 5 +++++ packages/lightning-mock/test.js | 27 +++++++++++++++++++++++++++ 2 files changed, 32 insertions(+) create mode 100644 .changeset/rude-squids-mix.md create mode 100644 packages/lightning-mock/test.js diff --git a/.changeset/rude-squids-mix.md b/.changeset/rude-squids-mix.md new file mode 100644 index 000000000..c03578c49 --- /dev/null +++ b/.changeset/rude-squids-mix.md @@ -0,0 +1,5 @@ +--- +'@openfn/ws-worker': patch +--- + +Count outstanding claim requests as capacity. This fixes an issue where `work-available` messages can cause a worker to over-claim, particularly during periods of high load on the Lightning database. diff --git a/packages/lightning-mock/test.js b/packages/lightning-mock/test.js new file mode 100644 index 000000000..99a54d361 --- /dev/null +++ b/packages/lightning-mock/test.js @@ -0,0 +1,27 @@ +let id = 0; + +const wf = { + id: 'my-attempt', + triggers: [], + edges: [], + jobs: [ + { + id: 'job1', + adaptor: '@openfn/language-common@1.7.7', + body: 'fn(() => new Promise((resolve) => setTimeout(resolve, 5000)) )', + }, + ], +}; + +// setInterval(() => { +fetch('http://localhost:8888/run', { + method: 'POST', + body: JSON.stringify({ + ...wf, + id: `${++id}`, + }), + headers: { + 'content-type': 'application/json', + }, +}); +// }, 500); From c47b65fe4da1e1b44246a50a0b70e51e28cae94c Mon Sep 17 00:00:00 2001 From: Joe Clark Date: Thu, 11 Sep 2025 13:45:57 +0100 Subject: [PATCH 7/7] changeset --- .changeset/tough-ties-tease.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .changeset/tough-ties-tease.md diff --git a/.changeset/tough-ties-tease.md b/.changeset/tough-ties-tease.md new file mode 100644 index 000000000..19591e4b5 --- /dev/null +++ b/.changeset/tough-ties-tease.md @@ -0,0 +1,5 @@ +--- +'@openfn/lightning-mock': patch +--- + +Ensure json payloads respect sockey delay