Worker: ensure workers are released after an error on run:complete (#685

) * worker: added logging Noisy but useful * worker: ensure run:complete returns even if the socket throws * changeset * worker: add a bit more logging on the channel * version: [email protected] * worker: fix mock socket
OpenFn · May 10, 2024 · 75f9087 · 75f9087
1 parent 24d2014
commit 75f9087
Show file tree

Hide file tree

Showing 10 changed files with 145 additions and 7 deletions.
diff --git a/integration-tests/worker/CHANGELOG.md b/integration-tests/worker/CHANGELOG.md
@@ -1,5 +1,12 @@
 # @openfn/integration-tests-worker
 
+## 1.0.42
+
+### Patch Changes
+
+- Updated dependencies [bdff4b2]
+  - @openfn/ws-worker@1.1.7
+
 ## 1.0.41
 
 ### Patch Changes

diff --git a/integration-tests/worker/package.json b/integration-tests/worker/package.json
@@ -1,7 +1,7 @@
 {
   "name": "@openfn/integration-tests-worker",
   "private": true,
-  "version": "1.0.41",
+  "version": "1.0.42",
   "description": "Lightning WOrker integration tests",
   "author": "Open Function Group <[email protected]>",
   "license": "ISC",

diff --git a/packages/ws-worker/CHANGELOG.md b/packages/ws-worker/CHANGELOG.md
@@ -1,5 +1,11 @@
 # ws-worker
 
+## 1.1.7
+
+### Patch Changes
+
+- bdff4b2: Fix an issue where workers may not be returned to the pool if run:complete throws
+
 ## 1.1.6
 
 ### Patch Changes

diff --git a/packages/ws-worker/package.json b/packages/ws-worker/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@openfn/ws-worker",
-  "version": "1.1.6",
+  "version": "1.1.7",
   "description": "A Websocket Worker to connect Lightning to a Runtime Engine",
   "main": "dist/index.js",
   "type": "module",

diff --git a/packages/ws-worker/src/api/claim.ts b/packages/ws-worker/src/api/claim.ts
@@ -35,10 +35,12 @@ const claim = (
 
     const activeWorkers = Object.keys(app.workflows).length;
     if (activeWorkers >= maxWorkers) {
+      logger.debug('skipping claim attempt: server at capacity');
       return reject(new Error('Server at capacity'));
     }
 
     if (!app.queueChannel) {
+      logger.debug('skipping claim attempt: websocket unavailable');
       return reject(new Error('No websocket available'));
     }
 

diff --git a/packages/ws-worker/src/channels/run.ts b/packages/ws-worker/src/channels/run.ts
@@ -46,7 +46,21 @@ const joinRunChannel = (
       .receive('error', (err: any) => {
         logger.error(`error connecting to ${channelName}`, err);
         reject(err);
+      })
+      .receive('timeout', (err: any) => {
+        logger.error(`Timeout for ${channelName}`, err);
+        reject(err);
       });
+    channel.onClose(() => {
+      // channel was explicitly closed by the client or server
+      logger.debug(`Leaving ${channelName}`);
+    });
+    channel.onError((e: any) => {
+      // Error occurred on the channel
+      // (the socket will try to reconnect with backoff)
+      logger.debug(`Error in ${channelName}`);
+      logger.debug(e);
+    });
   });
 };
 

diff --git a/packages/ws-worker/src/events/run-complete.ts b/packages/ws-worker/src/events/run-complete.ts
@@ -10,7 +10,7 @@ export default async function onWorkflowComplete(
   context: Context,
   _event: WorkflowCompletePayload
 ) {
-  const { state, channel, onFinish } = context;
+  const { state, channel, onFinish, logger } = context;
 
   // TODO I dont think the run final dataclip IS the last job dataclip
   // Especially not in parallelisation
@@ -19,10 +19,17 @@ export default async function onWorkflowComplete(
   const reason = calculateRunExitReason(state);
   await logFinalReason(context, reason);
 
-  await sendEvent<RunCompletePayload>(channel, RUN_COMPLETE, {
-    final_dataclip_id: state.lastDataclipId!,
-    ...reason,
-  });
+  try {
+    await sendEvent<RunCompletePayload>(channel, RUN_COMPLETE, {
+      final_dataclip_id: state.lastDataclipId!,
+      ...reason,
+    });
+  } catch (e) {
+    logger.error(
+      `${state.plan.id} failed to send ${RUN_COMPLETE} event. This run will be lost!`
+    );
+    logger.error(e);
+  }
 
   onFinish({ reason, state: result });
 }
diff --git a/packages/ws-worker/src/mock/sockets.ts b/packages/ws-worker/src/mock/sockets.ts
@@ -22,6 +22,8 @@ export const mockChannel = (
           } catch (e) {
             responses.error?.(e);
           }
+        } else {
+          responses.timeout?.('timeout');
         }
       }, 1);
 
@@ -64,6 +66,8 @@ export const mockChannel = (
       return receive;
     },
     leave: () => {},
+    onClose: () => {},
+    onError: () => {},
   };
   return c;
 };

diff --git a/packages/ws-worker/src/server.ts b/packages/ws-worker/src/server.ts
@@ -180,6 +180,7 @@ function createServer(engine: RuntimeEngine, options: ServerOptions = {}) {
 
       // Callback to be triggered when the work is done (including errors)
       const onFinish = () => {
+        logger.debug(`workflow ${id} complete: releasing worker`);
         delete app.workflows[id];
         runChannel.leave();
 

diff --git a/packages/ws-worker/test/events/run-complete.test.ts b/packages/ws-worker/test/events/run-complete.test.ts
@@ -1,4 +1,6 @@
 import test from 'ava';
+import { createMockLogger } from '@openfn/logger';
+
 import handleRunComplete from '../../src/events/run-complete';
 
 import { mockChannel } from '../../src/mock/sockets';
@@ -143,3 +145,98 @@ test('should send a reason log and return reason for fail', async (t) => {
   t.is(completeEvent.error_type, 'TEST');
   t.is(completeEvent.error_message, 'err');
 });
+
+test('should call onFinish even if the lightning event throws', async (t) => {
+  const plan = createPlan();
+
+  const state = createRunState(plan);
+  state.dataclips = {
+    x: {},
+  };
+  state.lastDataclipId = 'x';
+
+  const channel = mockChannel({
+    [RUN_LOG]: () => true,
+    [RUN_COMPLETE]: () => {
+      throw new Error('they came from... behind!');
+    },
+  });
+
+  const event: any = {};
+
+  const logger = createMockLogger();
+
+  const context: any = {
+    channel,
+    state,
+    onFinish: () => {
+      t.pass('On finish called');
+    },
+    logger,
+  };
+  await handleRunComplete(context, event);
+});
+
+test('should log if the lightning event throws', async (t) => {
+  const plan = createPlan();
+
+  const state = createRunState(plan);
+  state.dataclips = {
+    x: {},
+  };
+  state.lastDataclipId = 'x';
+
+  const channel = mockChannel({
+    [RUN_LOG]: () => true,
+    [RUN_COMPLETE]: () => {
+      throw new Error('they came from... behind!');
+    },
+  });
+
+  const event: any = {};
+
+  const logger = createMockLogger();
+
+  const context: any = {
+    channel,
+    state,
+    onFinish: () => {
+      const message = logger._find(
+        'error',
+        /failed to send run:complete event/
+      );
+      t.truthy(message);
+    },
+    logger,
+  };
+  await handleRunComplete(context, event);
+});
+
+test('should call onFinish even if the lightning event timesout', async (t) => {
+  const plan = createPlan();
+
+  const state = createRunState(plan);
+  state.dataclips = {
+    x: {},
+  };
+  state.lastDataclipId = 'x';
+
+  const channel = mockChannel({
+    [RUN_LOG]: () => true,
+    // no event handler is registered, so the mock will throw a timeout
+  });
+
+  const event: any = {};
+
+  const logger = createMockLogger();
+
+  const context: any = {
+    channel,
+    state,
+    onFinish: () => {
+      t.pass('On finish called');
+    },
+    logger,
+  };
+  await handleRunComplete(context, event);
+});