Fix flaky tests maybe

.changeset/quick-pens-close.md

@@ -0,0 +1,5 @@
+---
+'@openfn/ws-worker': minor
+---
+
+Add an `immediate` option to worker destroy to tear down without waiting for runs

integration-tests/worker/src/init.ts

@@ -53,7 +53,7 @@ export const initWorker = async (
     lightning: `ws://localhost:${lightningPort}/worker`,
     secret: crypto.randomUUID(),
     collectionsVersion: '1.0.0',
-    messageTimeoutSeconds: 0.01,
+    messageTimeoutSeconds: 0.1,
     batchLogs: true,
     ...workerArgs,
   });

integration-tests/worker/test/integration.test.ts

@@ -12,7 +12,7 @@ let engine;
 let engineLogger;
 let lightningPort;
 
-test.before(async () => {
+test.beforeEach(async () => {
   const keys = await generateKeys();
   lightningPort = nextPort();
   lightning = initLightning(lightningPort, keys.private);
@@ -33,13 +33,10 @@ test.before(async () => {
   ));
 });
 
-test.afterEach(() => {
+test.afterEach(async () => {
   engineLogger._reset();
-});
-
-test.after(async () => {
   lightning.destroy();
-  await worker.destroy();
+  await worker.destroy({ immediate: true });
 });
 
 const createDummyWorker = () => {
@@ -653,7 +650,7 @@ test.serial('use local adaptor versions from monorepo', (t) => {
 test.serial("Don't send adaptor logs to stdout", (t) => {
   return new Promise(async (done) => {
     // We have to create a new worker with a different repo for this one
-    await worker.destroy();
+    await worker.destroy({ immediate: true });
     ({ worker, engineLogger } = await createDummyWorker());
 
     const message = 've have been expecting you meester bond';
@@ -692,7 +689,7 @@ test.serial(
   (t) => {
     return new Promise(async (done) => {
       // We want to create our own special worker here
-      await worker.destroy();
+      await worker.destroy({ immediate: true });
       ({ worker, engineLogger } = await createDummyWorker());
 
       const attempt1 = {
@@ -739,7 +736,7 @@ test.serial(
 test.serial('worker should exit if it has an invalid key', (t) => {
   return new Promise(async (done) => {
     if (!worker.destroyed) {
-      await worker.destroy();
+      await worker.destroy({ immediate: true });
     }
 
     // generate a new, invalid, public key
@@ -792,7 +789,7 @@ test.serial('worker should exit if it has an invalid key', (t) => {
 test.serial('set a default timeout on the worker', (t) => {
   return new Promise(async (done) => {
     if (!worker.destroyed) {
-      await worker.destroy();
+      await worker.destroy({ immediate: true });
     }
 
     ({ worker } = await initWorker(lightningPort, {
@@ -842,7 +839,7 @@ test.serial('set a default timeout on the worker', (t) => {
 test.serial('set a timeout on a run', (t) => {
   return new Promise(async (done) => {
     if (!worker.destroyed) {
-      await worker.destroy();
+      await worker.destroy({ immediate: true });
     }
 
     ({ worker } = await initWorker(lightningPort, {
@@ -882,7 +879,7 @@ test.serial('set a timeout on a run', (t) => {
 test.serial('set a default payload limit on the worker', (t) => {
   return new Promise(async (done) => {
     if (!worker.destroyed) {
-      await worker.destroy();
+      await worker.destroy({ immediate: true });
     }
 
     ({ worker } = await initWorker(
@@ -923,7 +920,7 @@ test.serial('set a default payload limit on the worker', (t) => {
 test.serial('override the worker payload through run options', (t) => {
   return new Promise(async (done) => {
     if (!worker.destroyed) {
-      await worker.destroy();
+      await worker.destroy({ immediate: true });
     }
 
     ({ worker } = await initWorker(
@@ -965,7 +962,7 @@ test.serial('override the worker payload through run options', (t) => {
 test.serial('Redact logs which exceed the payload limit', (t) => {
   return new Promise(async (done) => {
     if (!worker.destroyed) {
-      await worker.destroy();
+      await worker.destroy({ immediate: true });
     }
 
     ({ worker } = await initWorker(lightningPort, {
@@ -1006,7 +1003,7 @@ test.serial('Redact logs which exceed the payload limit', (t) => {
 test.serial("Don't return dataclips which exceed the payload limit", (t) => {
   return new Promise(async (done) => {
     if (!worker.destroyed) {
-      await worker.destroy();
+      await worker.destroy({ immediate: true });
     }
 
     ({ worker } = await initWorker(lightningPort, {
@@ -1046,7 +1043,7 @@ test.serial(
   "Don't send job logs to stdout when job_log_level is set to none",
   (t) => {
     return new Promise(async (done) => {
-      await worker.destroy();
+      await worker.destroy({ immediate: true });
       ({ worker, engineLogger } = await createDummyWorker());
 
       const message = 'log that will never exist';

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

@@ -3,8 +3,20 @@ import { INTERNAL_CLAIM_COMPLETE, INTERNAL_RUN_COMPLETE } from '../events';
 
 import type { Logger } from '@openfn/logger';
 
-const destroy = async (app: ServerApp, logger: Logger) => {
-  logger.info('Closing server...');
+export type DestroyOptions = {
+  // Tear everything down now: don't wait for active runs or claims to complete
+  immediate?: boolean;
+};
+
+const destroy = async (
+  app: ServerApp,
+  logger: Logger,
+  options: DestroyOptions = {}
+) => {
+  const { immediate = false } = options;
+  logger.info(
+    immediate ? 'Closing server immediately...' : 'Closing server...'
+  );
 
   // Close the server AND wait for runs to complete in parallel
   // The http server can take a moment to close and the two
@@ -22,14 +34,21 @@ const destroy = async (app: ServerApp, logger: Logger) => {
       app.server.close(async () => {
         resolve();
       });
+
+      if (immediate) {
+        // Don't wait for open connections to drain
+        app.server.closeAllConnections?.();
+      }
     }),
     (async () => {
-      // Let any active runs complete
-      await waitForRunsAndClaims(app, logger);
+      if (!immediate) {
+        // Let any active runs complete
+        await waitForRunsAndClaims(app, logger);
+      }
 
       app.queueChannel?.leave();
       // Kill the engine and socket
-      await app.engine.destroy();
+      await app.engine.destroy(immediate);
       app.socket?.disconnect();
     })(),
   ]);

packages/ws-worker/src/server.ts

@@ -17,6 +17,7 @@ import {
   WORK_AVAILABLE,
 } from './events';
 import destroy from './api/destroy';
+import type { DestroyOptions } from './api/destroy';
 import { Workloop } from './api/workloop';
 import claim from './api/claim';
 import { Context, execute } from './api/execute';
@@ -87,7 +88,7 @@ export interface ServerApp extends Koa {
   runWorkloopMap: Record<string, Workloop>;
 
   execute: ({ id, token }: ClaimRun) => Promise<void>;
-  destroy: () => void;
+  destroy: (options?: DestroyOptions) => Promise<void>;
   resumeWorkloop: (workloop?: Workloop) => void;
   pendingClaims: () => number;
 
@@ -469,7 +470,7 @@ function createServer(engine: RuntimeEngine, options: ServerOptions = {}) {
   app.pendingClaims = () =>
     app.workloops.reduce((sum, w) => sum + Object.keys(w.openClaims).length, 0);
 
-  app.destroy = () => destroy(app, logger);
+  app.destroy = (options?: DestroyOptions) => destroy(app, logger, options);
 
   app.use(router.routes());
 

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

@@ -337,3 +337,72 @@ test("don't claim after destroy", async (t) => {
     lightning.enqueueRun(createRun());
   });
 });
+
+test.serial(
+  'immediate destroy does not wait for an outstanding claim',
+  async (t) => {
+    // Delay the claim response so it is still outstanding when destroy returns
+    initLightning({ socketDelay: 200 });
+    await initWorker({ noLoop: true });
+
+    return new Promise<void>((done) => {
+      const doDestroy = async () => {
+        t.is(worker.pendingClaims(), 1);
+
+        await destroy(worker, logger, { immediate: true });
+
+        // A graceful destroy would block until the claim completed
+        t.is(worker.pendingClaims(), 1);
+        t.true(worker.destroyed);
+        t.false(await ping());
+
+        done();
+      };
+
+      worker.events.once(INTERNAL_CLAIM_START, () => {
+        doDestroy();
+      });
+
+      lightning.enqueueRun(createRun());
+      worker.claim().catch(() => {});
+    });
+  }
+);
+
+test.serial(
+  'immediate destroy tells the engine to destroy instantly',
+  async (t) => {
+    initLightning();
+    await initWorker();
+
+    let instant: boolean | undefined;
+    const originalDestroy = worker.engine.destroy;
+    worker.engine.destroy = async (arg?: boolean) => {
+      instant = arg;
+      return originalDestroy(arg);
+    };
+
+    await worker.destroy({ immediate: true });
+
+    t.true(instant);
+  }
+);
+
+test.serial(
+  'a normal destroy does not ask the engine to destroy instantly',
+  async (t) => {
+    initLightning();
+    await initWorker();
+
+    let instant: boolean | undefined = true;
+    const originalDestroy = worker.engine.destroy;
+    worker.engine.destroy = async (arg?: boolean) => {
+      instant = arg;
+      return originalDestroy(arg);
+    };
+
+    await worker.destroy();
+
+    t.falsy(instant);
+  }
+);