feat: ws-worker respond to wakeup call (#877)

doc-han · josephjclark · web-flow · commit 36dc3dbaeb73 · 2025-02-21T18:09:39.000Z
* feat: trigger instant claim on work-available event

* feat: update lightning mock to send work-available event via ws

* fix: whitelist timestamp key in lighning mock

* chore: fix type issues

* tests: should recieve worker:queue message events from lightning

* tests: add onMessage mock to socket mock

* tests: manually trigger claim on connected workers

* feat: emit messages for worker:queue directly

* minor tweaks
- make ?wakeup easier in the mock
- catch claim errors in the worker after work-available
- fix a typo in the readme

* changesets

* lint

* version bumps

* simplify wake-up tests

* simplify wakup in mock

* worker@1.11.0

---------

Co-authored-by: Joe Clark &lt;jclark@openfn.org&gt;
diff --git a/integration-tests/worker/CHANGELOG.md b/integration-tests/worker/CHANGELOG.md
@@ -1,5 +1,14 @@
 # @openfn/integration-tests-worker
 
+## 1.0.77
+
+### Patch Changes
+
+- Updated dependencies [87f10f7]
+- Updated dependencies [87f10f7]
+  - @openfn/lightning-mock@2.1.0
+  - @openfn/ws-worker@1.11.0
+
 ## 1.0.76
 
 ### 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.76",
+  "version": "1.0.77",
   "description": "Lightning WOrker integration tests",
   "author": "Open Function Group <admin@openfn.org>",
   "license": "ISC",
diff --git a/packages/lightning-mock/CHANGELOG.md b/packages/lightning-mock/CHANGELOG.md
@@ -1,5 +1,11 @@
 # @openfn/lightning-mock
 
+## 2.1.0
+
+### Minor Changes
+
+- 87f10f7: Add ?wakeup paramter to trigger work-available events
+
 ## 2.0.31
 
 ### Patch Changes
diff --git a/packages/lightning-mock/README.md b/packages/lightning-mock/README.md
@@ -33,7 +33,7 @@ The server exposes a small dev API allowing you to post an Run.
 You can add an run (`{ jobs, triggers, edges }`) to the queue with:
 
 ```
-curl http://localhost:8888/run --json @tmp/run.json
+curl -X POST http://localhost:8888/run --json @tmp/run.json
 ```
 
 Here's an example run:
diff --git a/packages/lightning-mock/package.json b/packages/lightning-mock/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@openfn/lightning-mock",
-  "version": "2.0.31",
+  "version": "2.1.0",
   "private": true,
   "description": "A mock Lightning server",
   "main": "dist/index.js",
diff --git a/packages/lightning-mock/src/api-dev.ts b/packages/lightning-mock/src/api-dev.ts
@@ -14,9 +14,11 @@ import type {
 import { ServerState } from './server';
 import { RUN_COMPLETE } from './events';
 import type { DevServer, LightningEvents } from './types';
+import { PhoenixEvent } from './socket-server';
 
 type Api = {
   startRun(runId: string): void;
+  messageClients(message: PhoenixEvent): void;
 };
 
 const setupDevAPI = (
@@ -40,6 +42,10 @@ const setupDevAPI = (
 
   app.getDataclip = (id: string) => state.dataclips[id];
 
+  app.messageSocketClients = (message: PhoenixEvent) => {
+    api.messageClients(message);
+  };
+
   app.enqueueRun = (run: LightningPlan, workerId = 'rte') => {
     state.runs[run.id] = run;
     state.results[run.id] = {
@@ -169,6 +175,19 @@ const setupRestAPI = (
 
     app.enqueueRun(run);
 
+    // triggering wakeup in all connected workers
+    if ('wakeup' in ctx.query) {
+      logger.info(
+        'WAKE UP! Sending work-available event to all listening workers'
+      );
+      app.messageSocketClients({
+        topic: 'worker:queue',
+        event: 'work-available',
+        payload: {},
+        join_ref: '',
+        ref: '',
+      });
+    }
     ctx.response.status = 200;
   });
 
diff --git a/packages/lightning-mock/src/api-sockets.ts b/packages/lightning-mock/src/api-sockets.ts
@@ -159,6 +159,7 @@ const createSocketAPI = (
 
   return {
     startRun,
+    messageClients: wss.sendToClients.bind(this),
     close: () => {
       server.close();
       (wss as any).close();
@@ -352,8 +353,14 @@ const createSocketAPI = (
     runId: string
   ) {
     const { ref, join_ref, topic } = evt;
-    const { final_dataclip_id, reason, error_type, error_message, ...rest } =
-      evt.payload;
+    const {
+      final_dataclip_id,
+      reason,
+      error_type,
+      error_message,
+      timestamp, // whitelist timestamp
+      ...rest
+    } = evt.payload;
 
     logger?.info('Completed run ', runId);
     logger?.debug(final_dataclip_id);
diff --git a/packages/lightning-mock/src/socket-server.ts b/packages/lightning-mock/src/socket-server.ts
@@ -80,6 +80,7 @@ type MockSocketServer = typeof WebSocketServer & {
     topic: Topic,
     events: Record<string, EventHandler>
   ) => { unsubscribe: () => void };
+  sendToClients: (message: PhoenixEvent) => void;
 };
 
 function createServer({
@@ -95,6 +96,8 @@ function createServer({
     phoenix: new Set([() => null]),
   };
 
+  const clients: Set<DevSocket> = new Set();
+
   const wsServer =
     server ||
     new WebSocketServer({
@@ -202,6 +205,9 @@ function createServer({
 
         logger?.debug(`>> [${topic}] ${event} ${ref} :: ${stringify(payload)}`);
 
+        // tracking connected worker:queue workers
+        if (topic === 'worker:queue' && event === 'phx_join') clients.add(ws);
+
         if (event in events) {
           // handle system/phoenix events
           // @ts-ignore
@@ -227,10 +233,19 @@ function createServer({
         }
       }
     });
+
+    ws.on('close', () => clients.delete(ws));
   });
 
   const mockServer = wsServer as MockSocketServer;
 
+  mockServer.sendToClients = async (message) => {
+    clients.forEach((client) => {
+      // @ts-ignore
+      client.sendJSON(message);
+    });
+  };
+
   // debug API
   // TODO should this in fact be (topic, event, fn)?
   mockServer.listenToChannel = (topic: Topic, fn: EventHandler) => {
diff --git a/packages/lightning-mock/src/types.ts b/packages/lightning-mock/src/types.ts
@@ -5,8 +5,9 @@ import type {
   Credential,
 } from '@openfn/lexicon/lightning';
 import type { ServerState } from './server';
+import { PhoenixEvent } from './socket-server';
 
-export type LightningEvents = 'log' | 'run-complete';
+export type LightningEvents = 'log' | 'run-complete' | string; // not complete!
 
 export type DevServer = Koa & {
   state: ServerState;
@@ -20,13 +21,15 @@ export type DevServer = Koa & {
   getQueueLength(): number;
   getResult(runId: string): any;
   getState(): ServerState;
+  messageSocketClients(message: PhoenixEvent): void;
   on(event: LightningEvents, fn: (evt: any) => void): void;
   once(event: LightningEvents, fn: (evt: any) => void): void;
   onSocketEvent(
     event: LightningEvents,
     runId: string,
-    fn: (evt: any) => void
-  ): void;
+    fn: (evt: any) => void,
+    once?: boolean
+  ): () => void;
   registerRun(run: LightningPlan): void;
   removeAllListeners(): void;
   reset(): void;
diff --git a/packages/ws-worker/CHANGELOG.md b/packages/ws-worker/CHANGELOG.md
@@ -1,5 +1,15 @@
 # ws-worker
 
+## 1.11.0
+
+### Minor Changes
+
+- 87f10f7: Respond to `work:available` events.
+
+  When the worker receives `work:available` in the worker queue, it'll instantly trigger a claim event.
+
+  This claim is independent of the workloop and does not affect backoff in any way.
+
 ## 1.10.0
 
 ### Minor 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.10.0",
+  "version": "1.11.0",
   "description": "A Websocket Worker to connect Lightning to a Runtime Engine",
   "main": "dist/index.js",
   "type": "module",
diff --git a/packages/ws-worker/src/channels/worker-queue.ts b/packages/ws-worker/src/channels/worker-queue.ts
@@ -39,6 +39,11 @@ const connectToWorkerQueue = (
 
       const channel = socket.channel('worker:queue') as Channel;
 
+      channel.onMessage = (ev, load) => {
+        events.emit('message', ev, load);
+        return load;
+      };
+
       channel
         .join()
         .receive('ok', () => {
diff --git a/packages/ws-worker/src/events.ts b/packages/ws-worker/src/events.ts
@@ -1,5 +1,8 @@
 import * as l from '@openfn/lexicon/lightning';
 
+// events from lightning to workers
+export const WORK_AVAILABLE = 'work-available';
+
 // These are worker-lightning events, used in the websocket
 export const CLAIM = 'claim';
 export const GET_PLAN = 'fetch:plan';
diff --git a/packages/ws-worker/src/mock/sockets.ts b/packages/ws-worker/src/mock/sockets.ts
@@ -97,6 +97,10 @@ export class MockSocket {
     this.callbacks.onError = callback;
   }
 
+  onMessage(callback: EventHandler): void {
+    this.callbacks.onMessage = callback;
+  }
+
   onClose(callback: EventHandler): void {
     // TODO this isn't actually hooked up right now
     this.callbacks.onClose = callback;
diff --git a/packages/ws-worker/src/server.ts b/packages/ws-worker/src/server.ts
@@ -10,7 +10,7 @@ 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 } from './events';
+import { INTERNAL_RUN_COMPLETE, WORK_AVAILABLE } from './events';
 import destroy from './api/destroy';
 import startWorkloop, { Workloop } from './api/workloop';
 import claim from './api/claim';
@@ -135,6 +135,15 @@ function connect(app: ServerApp, logger: Logger, options: ServerOptions = {}) {
     logger.debug(e);
   };
 
+  // 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
+      });
+    }
+  };
+
   connectToWorkerQueue(
     options.lightning!,
     app.id,
@@ -144,7 +153,8 @@ function connect(app: ServerApp, logger: Logger, options: ServerOptions = {}) {
   )
     .on('connect', onConnect)
     .on('disconnect', onDisconnect)
-    .on('error', onError);
+    .on('error', onError)
+    .on('message', onMessage);
 }
 
 async function setupCollections(options: ServerOptions, logger: Logger) {
diff --git a/packages/ws-worker/test/lightning.test.ts b/packages/ws-worker/test/lightning.test.ts
diff --git a/packages/ws-worker/test/util.ts b/packages/ws-worker/test/util.ts
diff --git a/packages/ws-worker/test/wake-up.test.ts b/packages/ws-worker/test/wake-up.test.ts

Original file line number	Diff line number	Diff line change
`@@ -1,7 +1,7 @@`
`1`	`1`	`{`
`2`	`2`	`"name": "@openfn/integration-tests-worker",`
`3`	`3`	`"private": true,`
`4`		`- "version": "1.0.76",`
	`4`	`+ "version": "1.0.77",`
`5`	`5`	`"description": "Lightning WOrker integration tests",`
`6`	`6`	`"author": "Open Function Group <[email protected]>",`
`7`	`7`	`"license": "ISC",`
Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`{`
`2`	`2`	`"name": "@openfn/lightning-mock",`
`3`		`- "version": "2.0.31",`
	`3`	`+ "version": "2.1.0",`
`4`	`4`	`"private": true,`
`5`	`5`	`"description": "A mock Lightning server",`
`6`	`6`	`"main": "dist/index.js",`