From 5688813a4c1548ec02424c1e0acb31cbae17b158 Mon Sep 17 00:00:00 2001 From: Joe Clark Date: Fri, 3 Oct 2025 17:00:18 +0100 Subject: [PATCH 1/5] worker: don't destroy until all outstanding claims have come home --- .changeset/cuddly-apples-buy.md | 5 + packages/lightning-mock/src/server.ts | 11 +- packages/lightning-mock/src/types.ts | 2 +- packages/ws-worker/src/api/claim.ts | 14 ++- packages/ws-worker/src/api/destroy.ts | 31 ++++-- packages/ws-worker/src/events.ts | 3 + packages/ws-worker/src/server.ts | 25 ++++- packages/ws-worker/test/api/destroy.test.ts | 111 +++++++++++++++++--- 8 files changed, 162 insertions(+), 40 deletions(-) create mode 100644 .changeset/cuddly-apples-buy.md diff --git a/.changeset/cuddly-apples-buy.md b/.changeset/cuddly-apples-buy.md new file mode 100644 index 000000000..1204dc9eb --- /dev/null +++ b/.changeset/cuddly-apples-buy.md @@ -0,0 +1,5 @@ +--- +'@openfn/ws-worker': patch +--- + +Allow the worker to shutdown gracefully while claims are still in-flight. Runs will be completed before the server closes diff --git a/packages/lightning-mock/src/server.ts b/packages/lightning-mock/src/server.ts index 1e08f34ad..5ef0e96f8 100644 --- a/packages/lightning-mock/src/server.ts +++ b/packages/lightning-mock/src/server.ts @@ -117,10 +117,13 @@ const createLightningServer = (options: LightningOptions = {}) => { app.use(createRestAPI(app as any, state, logger, api)); - app.destroy = () => { - server.close(); - api.close(); - }; + app.destroy = () => + new Promise(async (resolve) => { + api.close(); + server.close(() => { + resolve(); + }); + }); return app; }; diff --git a/packages/lightning-mock/src/types.ts b/packages/lightning-mock/src/types.ts index 1a86c5350..92d7e5626 100644 --- a/packages/lightning-mock/src/types.ts +++ b/packages/lightning-mock/src/types.ts @@ -14,7 +14,7 @@ export type DevServer = Koa & { addCredential(id: string, cred: Credential): void; addDataclip(id: string, data: DataClip): void; enqueueRun(run: LightningPlan): void; - destroy: () => void; + destroy: () => Promise; getRun(id: string): LightningPlan; getCredential(id: string): Credential; getDataclip(id: string): DataClip; diff --git a/packages/ws-worker/src/api/claim.ts b/packages/ws-worker/src/api/claim.ts index 7a64e0d84..f24c38dc1 100644 --- a/packages/ws-worker/src/api/claim.ts +++ b/packages/ws-worker/src/api/claim.ts @@ -4,7 +4,11 @@ import * as jose from 'jose'; import { Logger, createMockLogger } from '@openfn/logger'; import { ClaimPayload, ClaimReply } from '@openfn/lexicon/lightning'; -import { CLAIM } from '../events'; +import { + CLAIM, + INTERNAL_CLAIM_COMPLETE, + INTERNAL_CLAIM_START, +} from '../events'; import type { ServerApp } from '../server'; @@ -95,6 +99,7 @@ const claim = ( logger.debug(`requesting run (capacity ${activeWorkers}/${maxWorkers})`); + app.events.emit(INTERNAL_CLAIM_START); const start = Date.now(); app.queueChannel .push(CLAIM, { @@ -109,10 +114,9 @@ const claim = ( runs.length ? runs.map((r) => r.id).join(',') : '-' })` ); - // TODO what if we get here after we've been cancelled? - // the events have already been claimed... if (!runs?.length) { + app.events.emit(INTERNAL_CLAIM_COMPLETE, { runs }); // throw to backoff and try again return reject(new Error('No runs returned')); } @@ -135,8 +139,10 @@ const claim = ( logger.debug(`${podName} starting run ${run.id}`); app.execute(run); - resolve(); }); + // Don't trigger claim complete until all runs are registered + resolve(); + app.events.emit(INTERNAL_CLAIM_COMPLETE, { runs }); }) // TODO need implementations for both of these really // What do we do if we fail to join the worker channel? diff --git a/packages/ws-worker/src/api/destroy.ts b/packages/ws-worker/src/api/destroy.ts index 130d92a21..52e51b415 100644 --- a/packages/ws-worker/src/api/destroy.ts +++ b/packages/ws-worker/src/api/destroy.ts @@ -1,5 +1,5 @@ import { ServerApp } from '../server'; -import { INTERNAL_RUN_COMPLETE } from '../events'; +import { INTERNAL_CLAIM_COMPLETE, INTERNAL_RUN_COMPLETE } from '../events'; import type { Logger } from '@openfn/logger'; @@ -15,7 +15,6 @@ const destroy = async (app: ServerApp, logger: Logger) => { // Immediately stop asking for more work app.workloop?.stop('server closed'); - app.queueChannel?.leave(); // Shut down the HTTP server app.server.close(async () => { @@ -24,8 +23,9 @@ const destroy = async (app: ServerApp, logger: Logger) => { }), new Promise(async (resolve) => { // Let any active runs complete - await waitForRuns(app, logger); + await waitForRunsAndClaims(app, logger); + app.queueChannel?.leave(); // Kill the engine and socket await app.engine.destroy(); app.socket?.disconnect(); @@ -37,27 +37,38 @@ const destroy = async (app: ServerApp, logger: Logger) => { logger.success('Server closed'); }; -const waitForRuns = (app: ServerApp, logger: Logger) => +const waitForRunsAndClaims = (app: ServerApp, logger: Logger) => new Promise((resolve) => { const log = () => { logger.debug( - `Waiting for ${Object.keys(app.workflows).length} runs to complete...` + `Waiting for ${Object.keys(app.workflows).length} runs and ${ + Object.keys(app.openClaims).length + } claims to complete...` ); }; - const onRunComplete = () => { - if (Object.keys(app.workflows).length === 0) { + const checkAllClear = () => { + if ( + Object.keys(app.workflows).length + + Object.keys(app.openClaims).length === + 0 + ) { logger.debug('All runs completed!'); - app.events.off(INTERNAL_RUN_COMPLETE, onRunComplete); + app.events.off(INTERNAL_RUN_COMPLETE, checkAllClear); + app.events.off(INTERNAL_CLAIM_COMPLETE, checkAllClear); resolve(); } else { log(); } }; - if (Object.keys(app.workflows).length) { + if ( + Object.keys(app.workflows).length || + Object.keys(app.openClaims).length + ) { log(); - app.events.on(INTERNAL_RUN_COMPLETE, onRunComplete); + app.events.on(INTERNAL_RUN_COMPLETE, checkAllClear); + app.events.on(INTERNAL_CLAIM_COMPLETE, checkAllClear); } else { logger.debug('No active runs detected, closing immediately'); resolve(); diff --git a/packages/ws-worker/src/events.ts b/packages/ws-worker/src/events.ts index e9c29b20a..de757d0d7 100644 --- a/packages/ws-worker/src/events.ts +++ b/packages/ws-worker/src/events.ts @@ -14,6 +14,9 @@ export const RUN_LOG = 'run:log'; export const STEP_START = 'step:start'; export const STEP_COMPLETE = 'step:complete'; export const INTERNAL_RUN_COMPLETE = 'server:run-complete'; +export const INTERNAL_CLAIM_START = 'server:claim-start'; +export const INTERNAL_CLAIM_COMPLETE = 'server:claim-complete'; +export const INTERNAL_SOCKET_READY = 'server:socket-ready'; export type QueueEvents = { [CLAIM]: l.ClaimPayload; diff --git a/packages/ws-worker/src/server.ts b/packages/ws-worker/src/server.ts index ab23d7e23..fed858e9f 100644 --- a/packages/ws-worker/src/server.ts +++ b/packages/ws-worker/src/server.ts @@ -11,7 +11,11 @@ import Router from '@koa/router'; import { humanId } from 'human-id'; import { createMockLogger, Logger } from '@openfn/logger'; import { ClaimRun } from '@openfn/lexicon/lightning'; -import { INTERNAL_RUN_COMPLETE, WORK_AVAILABLE } from './events'; +import { + INTERNAL_RUN_COMPLETE, + INTERNAL_SOCKET_READY, + WORK_AVAILABLE, +} from './events'; import destroy from './api/destroy'; import startWorkloop, { Workloop } from './api/workloop'; import claim from './api/claim'; @@ -71,6 +75,9 @@ export interface ServerApp extends Koa { execute: ({ id, token }: ClaimRun) => Promise; destroy: () => void; resumeWorkloop: () => void; + + // debug API + claim: () => Promise; } type SocketAndChannel = { @@ -112,6 +119,7 @@ function connect(app: ServerApp, logger: Logger, options: ServerOptions = {}) { logger.break(); } + app.events.emit(INTERNAL_SOCKET_READY); app.resumeWorkloop(); }; @@ -147,9 +155,12 @@ function connect(app: ServerApp, logger: Logger, options: ServerOptions = {}) { // handles messages for the worker:queue const onMessage = (event: string) => { if (event === WORK_AVAILABLE) { - claim(app, logger, { maxWorkers: options.maxWorkflows }).catch(() => { - // do nothing - it's fine if claim throws here - }); + // TODO ought to have a unit test against this + if (!app.destroyed) { + claim(app, logger, { maxWorkers: options.maxWorkflows }).catch(() => { + // do nothing - it's fine if claim throws here + }); + } } }; @@ -351,6 +362,12 @@ function createServer(engine: RuntimeEngine, options: ServerOptions = {}) { }); }); + app.claim = () => { + return claim(app, logger, { + maxWorkers: options.maxWorkflows, + }); + }; + app.destroy = () => destroy(app, logger); app.use(router.routes()); diff --git a/packages/ws-worker/test/api/destroy.test.ts b/packages/ws-worker/test/api/destroy.test.ts index 29ef88092..fcbc344ed 100644 --- a/packages/ws-worker/test/api/destroy.test.ts +++ b/packages/ws-worker/test/api/destroy.test.ts @@ -7,29 +7,45 @@ import { LightningPlan } from '@openfn/lexicon/lightning'; import createWorker from '../../src/server'; import createMockRTE from '../../src/mock/runtime-engine'; import destroy from '../../src/api/destroy'; +import { + INTERNAL_CLAIM_COMPLETE, + INTERNAL_CLAIM_START, + INTERNAL_SOCKET_READY, +} from '../../src'; const workerPort = 9876; const lightningPort = workerPort + 1; const logger = createMockLogger(); -const lightning = createLightningServer({ port: lightningPort }); +let lightning: any; let worker: any; -test.beforeEach(async () => { - const engine: any = await createMockRTE(); +const initLightning = (options = {}) => { + lightning = createLightningServer({ port: lightningPort, ...options }); +}; - worker = createWorker(engine, { - logger, - lightning: `ws://localhost:${lightningPort}/worker`, - port: workerPort, - backoff: { min: 10, max: 20 }, - collectionsVersion: '1.0.0', +const initWorker = (options = {}) => { + return new Promise(async (resolve) => { + const engine: any = await createMockRTE(); + + worker = createWorker(engine, { + logger, + lightning: `ws://localhost:${lightningPort}/worker`, + port: workerPort, + backoff: { min: 10, max: 20 }, + collectionsVersion: '1.0.0', + ...options, + }); + worker.events.on(INTERNAL_SOCKET_READY, () => { + resolve(); + }); }); -}); +}; -test.afterEach(() => { - lightning.reset(); +test.afterEach(async () => { + await lightning.destroy(); + logger._reset(); }); const createRun = () => @@ -45,17 +61,22 @@ const createRun = () => const waitForClaim = (timeout: number = 1000) => new Promise((resolve) => { + let didTimeout = false; + let didResolve = false; const handler = () => { + didResolve = true; + lightning.off('claim', handler); if (!didTimeout) { resolve(true); } }; - let didTimeout = false; setTimeout(() => { - didTimeout = true; - lightning.off('claim', handler); - resolve(false); + if (!didResolve) { + didTimeout = true; + lightning.off('claim', handler); + resolve(false); + } }, timeout); lightning.once('claim', handler); @@ -72,6 +93,9 @@ const ping = async () => { }; test.serial('destroy a worker with no active runs', async (t) => { + initLightning(); + await initWorker(); + // should respond to get t.true(await ping()); // should be claiming @@ -89,6 +113,9 @@ test.serial('destroy a worker with no active runs', async (t) => { // WARNING this might be flaky in CI test.serial('destroy a worker while one run is active', async (t) => { + initLightning(); + await initWorker(); + return new Promise((done) => { let didFinish = false; @@ -122,6 +149,9 @@ test.serial('destroy a worker while one run is active', async (t) => { }); test.serial('destroy a worker while multiple runs are active', async (t) => { + initLightning(); + await initWorker(); + return new Promise((done) => { let completeCount = 0; let startCount = 0; @@ -159,7 +189,54 @@ test.serial('destroy a worker while multiple runs are active', async (t) => { }); }); -test("don't claim after destroy", (t) => { +test.serial( + 'destroy a worker while a claim is outstanding and wait for the run to complete', + async (t) => { + t.plan(4); + initLightning({ + socketDelay: 50, + }); + await initWorker({ noLoop: true }); + return new Promise((done) => { + let didFinish = false; + + const doDestroy = async () => { + await destroy(worker, logger); + + t.true(didFinish); + + // should not respond to get + t.false(await ping()); + // should not be claiming + t.false(await waitForClaim()); + + done(); + }; + // As soon as the claim starts, kill the worker + worker.events.once(INTERNAL_CLAIM_START, () => { + doDestroy(); + }); + + // We still expect the run to complete + lightning.once('run:complete', () => { + didFinish = true; + }); + + // By the time the claim is complete, the worker should be marked destroyed + worker.events.once(INTERNAL_CLAIM_COMPLETE, () => { + t.true(worker.destroyed); + }); + + lightning.enqueueRun(createRun()); + worker.claim().catch(); + }); + } +); + +test("don't claim after destroy", async (t) => { + initLightning(); + await initWorker(); + return new Promise((done) => { let completeCount = 0; From ffa76897200cf12b19c54c8ece96e681f321fb12 Mon Sep 17 00:00:00 2001 From: Joe Clark Date: Fri, 3 Oct 2025 17:05:27 +0100 Subject: [PATCH 2/5] version@1.15.3 --- .changeset/cuddly-apples-buy.md | 5 ----- packages/ws-worker/CHANGELOG.md | 6 ++++++ packages/ws-worker/package.json | 2 +- packages/ws-worker/src/server.ts | 1 - 4 files changed, 7 insertions(+), 7 deletions(-) delete mode 100644 .changeset/cuddly-apples-buy.md diff --git a/.changeset/cuddly-apples-buy.md b/.changeset/cuddly-apples-buy.md deleted file mode 100644 index 1204dc9eb..000000000 --- a/.changeset/cuddly-apples-buy.md +++ /dev/null @@ -1,5 +0,0 @@ ---- -'@openfn/ws-worker': patch ---- - -Allow the worker to shutdown gracefully while claims are still in-flight. Runs will be completed before the server closes diff --git a/packages/ws-worker/CHANGELOG.md b/packages/ws-worker/CHANGELOG.md index ae84fe36c..835c95f54 100644 --- a/packages/ws-worker/CHANGELOG.md +++ b/packages/ws-worker/CHANGELOG.md @@ -1,5 +1,11 @@ # ws-worker +## 1.15.3 + +### Patch Changes + +- 5688813: Allow the worker to shutdown gracefully while claims are still in-flight. Runs will be completed before the server closes + ## 1.15.2 ### Patch Changes diff --git a/packages/ws-worker/package.json b/packages/ws-worker/package.json index 1a3f9bedd..9837e269c 100644 --- a/packages/ws-worker/package.json +++ b/packages/ws-worker/package.json @@ -1,6 +1,6 @@ { "name": "@openfn/ws-worker", - "version": "1.15.2", + "version": "1.15.3", "description": "A Websocket Worker to connect Lightning to a Runtime Engine", "main": "dist/index.js", "type": "module", diff --git a/packages/ws-worker/src/server.ts b/packages/ws-worker/src/server.ts index fed858e9f..57e41ec41 100644 --- a/packages/ws-worker/src/server.ts +++ b/packages/ws-worker/src/server.ts @@ -155,7 +155,6 @@ function connect(app: ServerApp, logger: Logger, options: ServerOptions = {}) { // handles messages for the worker:queue const onMessage = (event: string) => { if (event === WORK_AVAILABLE) { - // TODO ought to have a unit test against this if (!app.destroyed) { claim(app, logger, { maxWorkers: options.maxWorkflows }).catch(() => { // do nothing - it's fine if claim throws here From 30f9bf20311f247d255726e2ea5ffd05bee1721a Mon Sep 17 00:00:00 2001 From: Joe Clark Date: Fri, 3 Oct 2025 17:20:00 +0100 Subject: [PATCH 3/5] fix test --- packages/ws-worker/test/api/claim.test.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/ws-worker/test/api/claim.test.ts b/packages/ws-worker/test/api/claim.test.ts index 4bd197993..46345b4c4 100644 --- a/packages/ws-worker/test/api/claim.test.ts +++ b/packages/ws-worker/test/api/claim.test.ts @@ -8,6 +8,7 @@ import { createMockLogger } from '@openfn/logger'; import { ServerApp } from '../../src/server'; import { mockChannel } from '../../src/mock/sockets'; import { CLAIM } from '../../src'; +import EventEmitter from 'node:events'; let keys = { public: '.', private: '.' }; @@ -147,7 +148,8 @@ const createMockApp = (opts: any) => { execute: (...args) => { onExecute(...args); }, - } as ServerApp; + events: new EventEmitter(), + } as unknown as ServerApp; }; const logger = createMockLogger(); From 14e4ec19e1efff1b5ca45d102ba1a56e0131afd8 Mon Sep 17 00:00:00 2001 From: Joe Clark Date: Fri, 3 Oct 2025 17:26:42 +0100 Subject: [PATCH 4/5] types --- packages/ws-worker/test/api/claim.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/ws-worker/test/api/claim.test.ts b/packages/ws-worker/test/api/claim.test.ts index 46345b4c4..bb6783669 100644 --- a/packages/ws-worker/test/api/claim.test.ts +++ b/packages/ws-worker/test/api/claim.test.ts @@ -145,7 +145,7 @@ const createMockApp = (opts: any) => { openClaims: {}, workflows, queueChannel: channel, - execute: (...args) => { + execute: (...args: any) => { onExecute(...args); }, events: new EventEmitter(), From 682c5ab71fb43e5803a27c6e40d71adb6e96fff1 Mon Sep 17 00:00:00 2001 From: Joe Clark Date: Fri, 3 Oct 2025 17:42:05 +0100 Subject: [PATCH 5/5] fix more tests --- packages/ws-worker/test/api/workloop.test.ts | 37 ++++++++++++-------- 1 file changed, 22 insertions(+), 15 deletions(-) diff --git a/packages/ws-worker/test/api/workloop.test.ts b/packages/ws-worker/test/api/workloop.test.ts index 6a920e9af..091cfdf18 100644 --- a/packages/ws-worker/test/api/workloop.test.ts +++ b/packages/ws-worker/test/api/workloop.test.ts @@ -5,6 +5,7 @@ import { sleep } from '../util'; import { mockChannel } from '../../src/mock/sockets'; import startWorkloop from '../../src/api/workloop'; import { CLAIM } from '../../src/events'; +import EventEmitter from 'node:events'; let workloop: any; @@ -14,11 +15,23 @@ test.afterEach(() => { workloop?.stop(); // cancel any workloops }); +const createMockApp = (props: any) => ({ + workflows: {}, + queueChannel: mockChannel({ + [CLAIM]: () => { + return { runs: [] }; + }, + }), + execute: () => {}, + + events: new EventEmitter(), + ...props, +}); + test('workloop can be cancelled', async (t) => { let count = 0; - const app = { - workflows: {}, + const app = createMockApp({ queueChannel: mockChannel({ [CLAIM]: () => { count++; @@ -26,8 +39,7 @@ test('workloop can be cancelled', async (t) => { return { runs: [] }; }, }), - execute: () => {}, - }; + }); workloop = startWorkloop(app as any, logger, 1, 1); t.false(workloop.isStopped()); @@ -40,8 +52,7 @@ test('workloop can be cancelled', async (t) => { test('workloop sends the runs:claim event', (t) => { return new Promise((done) => { - const app = { - workflows: {}, + const app = createMockApp({ queueChannel: mockChannel({ [CLAIM]: () => { t.pass(); @@ -49,8 +60,7 @@ test('workloop sends the runs:claim event', (t) => { return { runs: [] }; }, }), - execute: () => {}, - }; + }); workloop = startWorkloop(app as any, logger, 1, 1); }); }); @@ -58,8 +68,7 @@ test('workloop sends the runs:claim event', (t) => { test('workloop sends the runs:claim event several times ', (t) => { return new Promise((done) => { let count = 0; - const app = { - workflows: {}, + const app = createMockApp({ queueChannel: mockChannel({ [CLAIM]: () => { count++; @@ -70,16 +79,14 @@ test('workloop sends the runs:claim event several times ', (t) => { return { runs: [] }; }, }), - execute: () => {}, - }; + }); workloop = startWorkloop(app as any, logger, 1, 1); }); }); test('workloop calls execute if runs:claim returns runs', (t) => { return new Promise((done) => { - const app = { - workflows: {}, + const app = createMockApp({ queueChannel: mockChannel({ [CLAIM]: () => ({ runs: [{ id: 'a', token: 'x.y.z' }], @@ -90,7 +97,7 @@ test('workloop calls execute if runs:claim returns runs', (t) => { t.pass(); done(); }, - }; + }); workloop = startWorkloop(app as any, logger, 1, 1); });