worker: don't destroy until all outstanding claims have come home

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<void>(async (resolve) => {
+      api.close();
+      server.close(() => {
+        resolve();
+      });
+    });
   return app;
 };
 

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<void>;
   getRun(id: string): LightningPlan;
   getCredential(id: string): Credential;
   getDataclip(id: string): DataClip;

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

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",

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<ClaimPayload>(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?

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<void>(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<void>((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();

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;

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<void>;
   destroy: () => void;
   resumeWorkloop: () => void;
+
+  // debug API
+  claim: () => Promise<any>;
 }
 
 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,11 @@ 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
-      });
+      if (!app.destroyed) {
+        claim(app, logger, { maxWorkers: options.maxWorkflows }).catch(() => {
+          // do nothing - it's fine if  claim throws here
+        });
+      }
     }
   };
 
@@ -351,6 +361,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());

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: '.' };
 
@@ -144,10 +145,11 @@ const createMockApp = (opts: any) => {
     openClaims: {},
     workflows,
     queueChannel: channel,
-    execute: (...args) => {
+    execute: (...args: any) => {
       onExecute(...args);
     },
-  } as ServerApp;
+    events: new EventEmitter(),
+  } as unknown as ServerApp;
 };
 const logger = createMockLogger();