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/.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 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/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); diff --git a/packages/ws-worker/src/api/claim.ts b/packages/ws-worker/src/api/claim.ts index 8cd782d24..7a64e0d84 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,20 +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; + app.openClaims ??= {}; + + 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')); } if (!app.queueChannel) { @@ -67,15 +89,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 (${ @@ -114,9 +141,12 @@ const claim = ( // 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 fb9b36d2a..4bd197993 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, { resetClaimIdGen, 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: '.' }; @@ -11,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, @@ -117,3 +125,246 @@ 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 = () => ({ runs: [] }), + onExecute = () => {}, + workflows = {}, + } = opts; + + const channel = mockChannel({ + [CLAIM]: (args) => { + return onClaim(args); + }, + }); + + return { + workloop: undefined, // should be safe + openClaims: {}, + 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: any[] = []; + const onClaim = () => ({ runs: [{ id: 'abc' }] }); + const onExecute = (...args: any[]) => { + executeArgs = args; + }; + + const options = { maxWorkers: 1 }; + const app = createMockApp({ + onClaim, + onExecute, + }); + + await claim(app, logger, options); + t.deepEqual(executeArgs[0], { id: 'abc' }); +}); + +test('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 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) => { + // @ts-ignore + setTimeout(resolve({ runs: [] }), 100); + }), + }); + + app.workloop = { + stop: () => { + didStopWorkloop = true; + }, + } as any; + + // @ts-ignore + 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) => { + // @ts-ignore + setTimeout(resolve({ runs: [] }), 100); + }), + }); + + // @ts-ignore + app.execute = ({ id }) => { + app.workflows[id] = true; + }; + + // first claim should be fine + let claimPromise = claim(app, logger, options); + + // 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'); +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)'); 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) => { 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', + }, + }); + }); +});