Ensure that pending claims count towards worker capacity

.changeset/rude-squids-mix.md

@@ -0,0 +1,5 @@
+---
+'@openfn/ws-worker': patch
+---
+
+Count outstanding claim requests as capacity. This fixes an issue where `work-available` messages can cause a worker to over-claim, particularly during periods of high load on the Lightning database.

.changeset/tough-ties-tease.md

@@ -0,0 +1,5 @@
+---
+'@openfn/lightning-mock': patch
+---
+
+Ensure json payloads respect sockey delay

packages/lightning-mock/src/socket-server.ts

@@ -190,8 +190,10 @@ function createServer({
         topic,
         payload: stringify(payload), // TODO do we stringify this? All of it?
       });
-      // @ts-ignore
-      ws.send(evt);
+      setTimeout(() => {
+        // @ts-ignore
+        ws.send(evt);
+      }, socketDelay);
     };
 
     ws.on('message', async function (data: string) {

packages/lightning-mock/test.js

@@ -0,0 +1,27 @@
+let id = 0;
+
+const wf = {
+  id: 'my-attempt',
+  triggers: [],
+  edges: [],
+  jobs: [
+    {
+      id: 'job1',
+      adaptor: '@openfn/language-common@1.7.7',
+      body: 'fn(() => new Promise((resolve) => setTimeout(resolve, 5000)) )',
+    },
+  ],
+};
+
+// setInterval(() => {
+fetch('http://localhost:8888/run', {
+  method: 'POST',
+  body: JSON.stringify({
+    ...wf,
+    id: `${++id}`,
+  }),
+  headers: {
+    'content-type': 'application/json',
+  },
+});
+// }, 500);

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

@@ -24,6 +24,7 @@ export const verifyToken = async (token: string, publicKey: string) => {
 
 type ClaimOptions = {
   maxWorkers?: number;
+  demand?: number;
 };
 
 // used to report the pod name in logging, for tracking
@@ -38,20 +39,41 @@ class ClaimError extends Error {
   }
 }
 
+let claimIdGen = 0;
+
+export const resetClaimIdGen = () => {
+  claimIdGen = 0;
+};
+
 const claim = (
   app: ServerApp,
   logger: Logger = mockLogger,
   options: ClaimOptions = {}
 ) => {
   return new Promise<void>((resolve, reject) => {
-    const { maxWorkers = 5 } = options;
+    app.openClaims ??= {};
+
+    const { maxWorkers = 5, demand = 1 } = options;
     const podName = NAME ? `[${NAME}] ` : '';
 
     const activeWorkers = Object.keys(app.workflows).length;
+
+    const pendingClaims = Object.values(app.openClaims).reduce(
+      (a, b) => a + b,
+      0
+    );
+
     if (activeWorkers >= maxWorkers) {
       // Important: stop the workloop so that we don't try and claim any more
       app.workloop?.stop(`server at capacity (${activeWorkers}/${maxWorkers})`);
       return reject(new ClaimError('Server at capacity'));
+    } else if (activeWorkers + pendingClaims >= maxWorkers) {
+      // There are active claims which haven't yet been fulfilled
+      // This can happen in response to the work-available event
+      app.workloop?.stop(
+        `server at capacity (${activeWorkers}/${maxWorkers}, ${pendingClaims} pending)`
+      );
+      return reject(new ClaimError('Server at capacity'));
     }
 
     if (!app.queueChannel) {
@@ -67,15 +89,20 @@ const claim = (
       return reject(e);
     }
 
+    const claimId = ++claimIdGen;
+
+    app.openClaims[claimId] = demand;
+
     logger.debug(`requesting run (capacity ${activeWorkers}/${maxWorkers})`);
 
     const start = Date.now();
     app.queueChannel
       .push<ClaimPayload>(CLAIM, {
-        demand: 1,
+        demand,
         worker_name: NAME || null,
       })
       .receive('ok', ({ runs }: ClaimReply) => {
+        delete app.openClaims[claimId];
         const duration = Date.now() - start;
         logger.debug(
           `${podName}claimed ${runs.length} runs in ${duration}ms (${
@@ -114,9 +141,12 @@ const claim = (
       // TODO need implementations for both of these really
       // What do we do if we fail to join the worker channel?
       .receive('error', (e) => {
+        delete app.openClaims[claimId];
         logger.error('Error on claim', e);
+        reject(new Error('claim error'));
       })
       .receive('timeout', () => {
+        delete app.openClaims[claimId];
         logger.error('TIMEOUT on claim. Runs may be lost.');
         reject(new Error('timeout'));
       });

packages/ws-worker/src/server.ts

@@ -60,6 +60,7 @@ export interface ServerApp extends Koa {
   socket?: any;
   queueChannel?: Channel;
   workflows: Record<string, true | Context>;
+  openClaims: Record<string, number>;
   destroyed: boolean;
   events: EventEmitter;
   server: Server;
@@ -221,6 +222,7 @@ function createServer(engine: RuntimeEngine, options: ServerOptions = {}) {
     })
   );
 
+  app.openClaims = {};
   app.workflows = {};
   app.destroyed = false;