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

* 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 <jclark@openfn.org>

Author: doc-han

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

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

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

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:

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

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;
   });
 

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

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) => {

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;

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