Merge pull request #1112 from OpenFn/lost-run-on-claim

worker: fix an edge case where a run can be lost between claim and st…

Author: josephjclark

.changeset/bumpy-sloths-act.md

@@ -0,0 +1,5 @@
+---
+'@openfn/ws-worker': patch
+---
+
+Fix an edge case where the server can shut down in-between claim and start (lost runs edge case)

packages/ws-worker/src/api/claim.ts

@@ -113,7 +113,7 @@ const claim = (
         demand,
         worker_name: NAME || null,
       })
-      .receive('ok', ({ runs }: ClaimReply) => {
+      .receive('ok', async ({ runs }: ClaimReply) => {
         delete app.openClaims[claimId];
         const duration = Date.now() - start;
         logger.debug(
@@ -128,7 +128,7 @@ const claim = (
           return reject(new Error('No runs returned'));
         }
 
-        runs.forEach(async (run) => {
+        for (const run of runs) {
           if (app.options?.runPublicKey) {
             try {
               await verifyToken(run.token, app.options.runPublicKey);
@@ -146,7 +146,8 @@ const claim = (
 
           logger.debug(`${podName} starting run ${run.id}`);
           app.execute(run);
-        });
+        }
+
         // Don't trigger claim complete until all runs are registered
         resolve();
         app.events.emit(INTERNAL_CLAIM_COMPLETE, { runs });

packages/ws-worker/test/api/destroy.test.ts

@@ -1,6 +1,9 @@
 import test from 'ava';
 import crypto from 'node:crypto';
-import createLightningServer from '@openfn/lightning-mock';
+import createLightningServer, {
+  generateKeys,
+  toBase64,
+} from '@openfn/lightning-mock';
 import { createMockLogger } from '@openfn/logger';
 import { LightningPlan } from '@openfn/lexicon/lightning';
 
@@ -18,8 +21,8 @@ const lightningPort = workerPort + 1;
 
 const logger = createMockLogger();
 let lightning: any;
-
 let worker: any;
+let keys = { private: '.', public: '.' };
 
 const initLightning = (options = {}) => {
   lightning = createLightningServer({ port: lightningPort, ...options });
@@ -43,6 +46,10 @@ const initWorker = (options = {}) => {
   });
 };
 
+test.before(async () => {
+  keys = await generateKeys();
+});
+
 test.afterEach(async () => {
   await lightning.destroy();
   logger._reset();
@@ -107,11 +114,8 @@ test.serial('destroy a worker with no active runs', async (t) => {
   t.false(await ping());
   // should not be claiming
   t.false(await waitForClaim());
-
-  // TODO how can I test the socket is closed?
 });
 
-// WARNING this might be flaky in CI
 test.serial('destroy a worker while one run is active', async (t) => {
   initLightning();
   await initWorker();
@@ -203,7 +207,7 @@ test.serial(
       const doDestroy = async () => {
         await destroy(worker, logger);
 
-        t.true(didFinish);
+        t.true(didFinish, 'Run did not finish');
 
         // should not respond to get
         t.false(await ping());
@@ -212,7 +216,56 @@ test.serial(
 
         done();
       };
-      // As  soon as the claim starts, kill the worker
+      // 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();
+    });
+  }
+);
+
+// The async bit is important because we can actually lose a run between claim and start
+test.serial(
+  'destroy a worker while a claim is outstanding and wait for the run to complete with async token validation',
+  async (t) => {
+    t.plan(4);
+
+    initLightning({
+      socketDelay: 50,
+      runPrivateKey: toBase64(keys.private),
+    });
+    await initWorker({ noLoop: true, runPublicKey: keys.public });
+
+    return new Promise((done) => {
+      let didFinish = false;
+
+      const doDestroy = async () => {
+        await destroy(worker, logger);
+
+        t.true(didFinish, 'Run did not finish');
+
+        // 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();
       });