@openfn/ws-worker 1.21.5 → 1.22.0

package/dist/index.js

@@ -33,7 +33,9 @@ var destroy = async (app, logger) => {
   await Promise.all([
     new Promise((resolve) => {
       app.destroyed = true;
-      app.workloop?.stop("server closed");
+      for (const w of app.workloops) {
+        w.stop("server closed");
+      }
       app.server.close(async () => {
         resolve();
       });
@@ -51,11 +53,11 @@ var destroy = async (app, logger) => {
 var waitForRunsAndClaims = (app, logger) => new Promise((resolve) => {
   const log = () => {
     logger.debug(
-      `Waiting for ${Object.keys(app.workflows).length} runs and ${Object.keys(app.openClaims).length} claims to complete...`
+      `Waiting for ${Object.keys(app.workflows).length} runs and ${app.pendingClaims()} claims to complete...`
     );
   };
   const checkAllClear = () => {
-    if (Object.keys(app.workflows).length + Object.keys(app.openClaims).length === 0) {
+    if (Object.keys(app.workflows).length + app.pendingClaims() === 0) {
       logger.debug("All runs completed!");
       app.events.off(INTERNAL_RUN_COMPLETE, checkAllClear);
       app.events.off(INTERNAL_CLAIM_COMPLETE, checkAllClear);
@@ -64,7 +66,7 @@ var waitForRunsAndClaims = (app, logger) => new Promise((resolve) => {
       log();
     }
   };
-  if (Object.keys(app.workflows).length || Object.keys(app.openClaims).length) {
+  if (Object.keys(app.workflows).length || app.pendingClaims()) {
     log();
     app.events.on(INTERNAL_RUN_COMPLETE, checkAllClear);
     app.events.on(INTERNAL_CLAIM_COMPLETE, checkAllClear);
@@ -75,51 +77,6 @@ var waitForRunsAndClaims = (app, logger) => new Promise((resolve) => {
 });
 var destroy_default = destroy;
 
-// src/util/try-with-backoff.ts
-var BACKOFF_MULTIPLIER = 1.15;
-var tryWithBackoff = (fn, opts = {}) => {
-  const { min = 1e3, max = 1e4, maxRuns, runs = 1 } = opts;
-  let cancelled = false;
-  if (!opts.isCancelled) {
-    opts.isCancelled = () => cancelled;
-  }
-  const promise = new Promise(async (resolve, reject) => {
-    try {
-      await fn();
-      resolve();
-    } catch (e) {
-      if (e?.abort) {
-        cancelled = true;
-        return reject();
-      }
-      if (opts.isCancelled()) {
-        return resolve();
-      }
-      if (!isNaN(maxRuns) && runs >= maxRuns) {
-        return reject(new Error("max runs exceeded"));
-      }
-      setTimeout(() => {
-        if (opts.isCancelled()) {
-          return resolve();
-        }
-        const nextOpts = {
-          maxRuns,
-          runs: runs + 1,
-          min: Math.min(max, min * BACKOFF_MULTIPLIER),
-          max,
-          isCancelled: opts.isCancelled
-        };
-        tryWithBackoff(fn, nextOpts).then(resolve).catch(reject);
-      }, min);
-    }
-  });
-  promise.cancel = () => {
-    cancelled = true;
-  };
-  return promise;
-};
-var try_with_backoff_default = tryWithBackoff;
-
 // src/api/claim.ts
 import v8 from "node:v8";
 import * as Sentry from "@sentry/node";
@@ -148,24 +105,26 @@ var ClaimError = class extends Error {
   }
 };
 var claimIdGen = 0;
-var claim = (app, logger = mockLogger, options = {}) => {
+var claim = (app, workloop, logger = mockLogger, options = {}) => {
   return new Promise((resolve, reject) => {
-    app.openClaims ??= {};
-    const { maxWorkers = 5, demand = 1 } = options;
+    const { demand = 1 } = options;
     const podName = NAME ? `[${NAME}] ` : "";
-    const activeWorkers = Object.keys(app.workflows).length;
-    const pendingClaims = Object.values(app.openClaims).reduce(
+    const activeInWorkloop = workloop.activeRuns.size;
+    const capacity = workloop.capacity;
+    const pendingWorkloopClaims = Object.values(workloop.openClaims).reduce(
       (a, b) => a + b,
       0
     );
-    if (activeWorkers >= maxWorkers) {
-      app.workloop?.stop(`server at capacity (${activeWorkers}/${maxWorkers})`);
-      return reject(new ClaimError("Server at capacity"));
-    } else if (activeWorkers + pendingClaims >= maxWorkers) {
-      app.workloop?.stop(
-        `server at capacity (${activeWorkers}/${maxWorkers}, ${pendingClaims} pending)`
+    if (activeInWorkloop >= capacity) {
+      workloop.stop(
+        `workloop ${workloop.id} at capacity (${activeInWorkloop}/${capacity})`
       );
-      return reject(new ClaimError("Server at capacity"));
+      return reject(new ClaimError("Workloop at capacity"));
+    } else if (activeInWorkloop + pendingWorkloopClaims >= capacity) {
+      workloop.stop(
+        `workloop ${workloop.id} at capacity (${activeInWorkloop}/${capacity}, ${pendingWorkloopClaims} pending)`
+      );
+      return reject(new ClaimError("Workloop at capacity"));
     }
     if (!app.queueChannel) {
       logger.warn("skipping claim attempt: websocket unavailable");
@@ -178,21 +137,22 @@ var claim = (app, logger = mockLogger, options = {}) => {
       return reject(e);
     }
     const claimId = ++claimIdGen;
-    app.openClaims[claimId] = demand;
+    workloop.openClaims[claimId] = demand;
     const { used_heap_size, heap_size_limit } = v8.getHeapStatistics();
     const usedHeapMb = Math.round(used_heap_size / 1024 / 1024);
     const totalHeapMb = Math.round(heap_size_limit / 1024 / 1024);
     const memPercent = Math.round(usedHeapMb / totalHeapMb * 100);
     logger.debug(
-      `Claiming runs :: demand ${demand} | capacity ${activeWorkers}/${maxWorkers} | memory ${memPercent}% (${usedHeapMb}/${totalHeapMb}mb)`
+      `Claiming runs [${workloop.id}] :: demand ${demand} | capacity ${activeInWorkloop}/${capacity} | memory ${memPercent}% (${usedHeapMb}/${totalHeapMb}mb)`
     );
     app.events.emit(INTERNAL_CLAIM_START);
     const start = Date.now();
     app.queueChannel.push(CLAIM, {
       demand,
-      worker_name: NAME || null
+      worker_name: NAME || null,
+      queues: workloop.queues
     }).receive("ok", async ({ runs }) => {
-      delete app.openClaims[claimId];
+      delete workloop.openClaims[claimId];
       const duration = Date.now() - start;
       logger.debug(
         `${podName}claimed ${runs.length} runs in ${duration}ms (${runs.length ? runs.map((r) => r.id).join(",") : "-"})`
@@ -216,17 +176,19 @@ var claim = (app, logger = mockLogger, options = {}) => {
         } else {
           logger.debug("skipping run token validation for", run.id);
         }
+        workloop.activeRuns.add(run.id);
+        app.runWorkloopMap[run.id] = workloop;
         logger.debug(`${podName} starting run ${run.id}`);
         app.execute(run);
       }
       resolve();
       app.events.emit(INTERNAL_CLAIM_COMPLETE, { runs });
     }).receive("error", (e) => {
-      delete app.openClaims[claimId];
+      delete workloop.openClaims[claimId];
       logger.error("Error on claim", e);
       reject(new Error("claim error"));
     }).receive("timeout", () => {
-      delete app.openClaims[claimId];
+      delete workloop.openClaims[claimId];
       logger.error("TIMEOUT on claim. Runs may be lost.");
       reject(new Error("timeout"));
     });
@@ -234,43 +196,6 @@ var claim = (app, logger = mockLogger, options = {}) => {
 };
 var claim_default = claim;
 
-// src/api/workloop.ts
-var startWorkloop = (app, logger, minBackoff, maxBackoff, maxWorkers) => {
-  let promise;
-  let cancelled = false;
-  const workLoop = () => {
-    if (!cancelled) {
-      promise = try_with_backoff_default(
-        () => claim_default(app, logger, {
-          maxWorkers
-        }),
-        {
-          min: minBackoff,
-          max: maxBackoff
-        }
-      );
-      promise.then(() => {
-        if (!cancelled) {
-          setTimeout(workLoop, minBackoff);
-        }
-      }).catch(() => {
-      });
-    }
-  };
-  workLoop();
-  return {
-    stop: (reason = "reason unknown") => {
-      if (!cancelled) {
-        logger.info(`cancelling workloop: ${reason}`);
-        cancelled = true;
-        promise.cancel();
-      }
-    },
-    isStopped: () => cancelled
-  };
-};
-var workloop_default = startWorkloop;
-
 // src/api/execute.ts
 import * as Sentry4 from "@sentry/node";
 import {
@@ -575,6 +500,51 @@ var stringify_default = (obj) => stringify(obj, (_key, value) => {
   return value;
 });
 
+// src/util/try-with-backoff.ts
+var BACKOFF_MULTIPLIER = 1.15;
+var tryWithBackoff = (fn, opts = {}) => {
+  const { min = 1e3, max = 1e4, maxRuns, runs = 1 } = opts;
+  let cancelled = false;
+  if (!opts.isCancelled) {
+    opts.isCancelled = () => cancelled;
+  }
+  const promise = new Promise(async (resolve, reject) => {
+    try {
+      await fn();
+      resolve();
+    } catch (e) {
+      if (e?.abort) {
+        cancelled = true;
+        return reject();
+      }
+      if (opts.isCancelled()) {
+        return resolve();
+      }
+      if (!isNaN(maxRuns) && runs >= maxRuns) {
+        return reject(new Error("max runs exceeded"));
+      }
+      setTimeout(() => {
+        if (opts.isCancelled()) {
+          return resolve();
+        }
+        const nextOpts = {
+          maxRuns,
+          runs: runs + 1,
+          min: Math.min(max, min * BACKOFF_MULTIPLIER),
+          max,
+          isCancelled: opts.isCancelled
+        };
+        tryWithBackoff(fn, nextOpts).then(resolve).catch(reject);
+      }, min);
+    }
+  });
+  promise.cancel = () => {
+    cancelled = true;
+  };
+  return promise;
+};
+var try_with_backoff_default = tryWithBackoff;
+
 // src/util/timestamp.ts
 var timeInMicroseconds = (time) => time && (BigInt(time) / BigInt(1e3)).toString();
 
@@ -1235,6 +1205,7 @@ var connectToWorkerQueue = (endpoint, serverId, secret, logger, options) => {
     messageTimeout = DEFAULT_MESSAGE_TIMEOUT_SECONDS,
     claimTimeout = DEFAULT_CLAIM_TIMEOUT_SECONDS,
     capacity,
+    queues,
     SocketConstructor = PhxSocket
   } = options;
   const events = new EventEmitter();
@@ -1271,6 +1242,9 @@ var connectToWorkerQueue = (endpoint, serverId, secret, logger, options) => {
       didOpen = true;
       shouldReportConnectionError = true;
       const joinPayload = { capacity };
+      if (queues) {
+        joinPayload.queues = queues;
+      }
       const channel = socket.channel("worker:queue", joinPayload);
       channel.onMessage = (ev, load) => {
         events.emit("message", ev, load);
@@ -1311,6 +1285,140 @@ var connectToWorkerQueue = (endpoint, serverId, secret, logger, options) => {
 };
 var worker_queue_default = connectToWorkerQueue;
 
+// src/api/workloop.ts
+var Workloop = class {
+  constructor({
+    id,
+    queues,
+    capacity
+  }) {
+    this.activeRuns = /* @__PURE__ */ new Set();
+    this.openClaims = {};
+    this.cancelled = true;
+    this.id = id;
+    this.queues = queues;
+    this.capacity = capacity;
+  }
+  hasCapacity() {
+    const pendingClaims = Object.values(this.openClaims).reduce(
+      (a, b) => a + b,
+      0
+    );
+    return this.activeRuns.size + pendingClaims < this.capacity;
+  }
+  start(app, logger, minBackoff, maxBackoff) {
+    this.logger = logger;
+    this.cancelled = false;
+    const loop = () => {
+      if (!this.cancelled) {
+        this.promise = try_with_backoff_default(() => claim_default(app, this, logger), {
+          min: minBackoff,
+          max: maxBackoff
+        });
+        this.promise.then(() => {
+          if (!this.cancelled) {
+            setTimeout(loop, minBackoff);
+          }
+        }).catch(() => {
+        });
+      }
+    };
+    loop();
+  }
+  stop(reason = "reason unknown") {
+    if (!this.cancelled) {
+      this.logger?.info(`cancelling workloop: ${reason}`);
+      this.cancelled = true;
+      this.promise?.cancel();
+    }
+  }
+  isStopped() {
+    return this.cancelled;
+  }
+};
+
+// src/util/parse-workloops.ts
+var WorkloopValidationError = class extends Error {
+  constructor(message) {
+    super(message);
+    this.name = "WorkloopValidationError";
+  }
+};
+var VALID_NAME = /^[a-zA-Z0-9_]+$/;
+function parseWorkloops(input) {
+  const trimmed = input.trim();
+  if (!trimmed) {
+    throw new WorkloopValidationError("Workloop configuration cannot be empty");
+  }
+  const tokens = trimmed.split(/\s+/);
+  const configs = tokens.map(parseToken);
+  const seenConfigs = /* @__PURE__ */ new Map();
+  for (let i = 0; i < configs.length; i++) {
+    const key = JSON.stringify(configs[i].queues);
+    if (seenConfigs.has(key)) {
+      const prevIndex = seenConfigs.get(key);
+      console.warn(
+        `Warning: workloops at positions ${prevIndex} and ${i} have identical queue configurations: ${tokens[prevIndex]} and ${tokens[i]}`
+      );
+    } else {
+      seenConfigs.set(key, i);
+    }
+  }
+  return configs;
+}
+function parseToken(token) {
+  const lastColon = token.lastIndexOf(":");
+  if (lastColon === -1) {
+    throw new WorkloopValidationError(
+      `Invalid token "${token}": missing :<count> suffix`
+    );
+  }
+  const prefStr = token.slice(0, lastColon);
+  const countStr = token.slice(lastColon + 1);
+  const count = Number(countStr);
+  if (!Number.isInteger(count) || countStr !== String(Math.floor(count))) {
+    throw new WorkloopValidationError(
+      `Invalid count "${countStr}" in token "${token}": must be a positive integer`
+    );
+  }
+  if (count < 1) {
+    throw new WorkloopValidationError(
+      `Invalid count "${countStr}" in token "${token}": must be >= 1`
+    );
+  }
+  const names = prefStr.split(">");
+  for (const name of names) {
+    if (name === "") {
+      throw new WorkloopValidationError(`Empty queue name in token "${token}"`);
+    }
+    if (name !== "*" && !VALID_NAME.test(name)) {
+      throw new WorkloopValidationError(
+        `Invalid queue name "${name}" in token "${token}": must match /^[a-zA-Z0-9_]+$/ or be "*"`
+      );
+    }
+  }
+  const nonWildcardNames = names.filter((n) => n !== "*");
+  const seen = /* @__PURE__ */ new Set();
+  for (const name of nonWildcardNames) {
+    if (seen.has(name)) {
+      console.warn(
+        `Warning: duplicate queue name "${name}" in token "${token}"`
+      );
+    }
+    seen.add(name);
+  }
+  const wildcardIndex = names.indexOf("*");
+  if (wildcardIndex !== -1 && wildcardIndex !== names.length - 1) {
+    throw new WorkloopValidationError(
+      `Wildcard "*" must be the last element in token "${token}"`
+    );
+  }
+  return new Workloop({ id: token, queues: names, capacity: count });
+}
+
+// src/util/get-default-workloop-config.ts
+var get_default_workloop_config_default = (capacity = 5) => `manual>*:${capacity}`;
+
 // src/server.ts
 var exec = promisify(_exec);
 var DEFAULT_PORT = 2222;
@@ -1338,8 +1446,10 @@ function connect(app, logger, options = {}) {
     app.resumeWorkloop();
   };
   const onDisconnect = () => {
-    if (!app.workloop?.isStopped()) {
-      app.workloop?.stop("Socket disconnected unexpectedly");
+    for (const w of app.workloops) {
+      if (!w.isStopped()) {
+        w.stop("Socket disconnected unexpectedly");
+      }
     }
     if (!app.destroyed) {
       logger.info("Connection to lightning lost");
@@ -1361,17 +1471,26 @@ function connect(app, logger, options = {}) {
   const onMessage = (event) => {
     if (event === WORK_AVAILABLE) {
       if (!app.destroyed) {
-        claim_default(app, logger, { maxWorkers: options.maxWorkflows }).catch(() => {
-        });
+        for (const w of app.workloops) {
+          if (w.hasCapacity()) {
+            claim_default(app, w, logger).catch(() => {
+            });
+          }
+        }
       }
     }
   };
+  const queuesMap = {};
+  for (const w of app.workloops) {
+    queuesMap[w.queues.join(">")] = w.capacity;
+  }
   worker_queue_default(options.lightning, app.id, options.secret, logger, {
     // TODO: options.socketTimeoutSeconds wins because this is what USED to be used
     // But it's deprecated and should be removed soon
     messageTimeout: options.socketTimeoutSeconds ?? options.messageTimeoutSeconds,
     claimTimeout: options.claimTimeoutSeconds,
-    capacity: options.maxWorkflows
+    capacity: options.maxWorkflows,
+    queues: queuesMap
   }).on("connect", onConnect).on("disconnect", onDisconnect).on("error", onError).on("message", onMessage);
 }
 async function setupCollections(options, logger) {
@@ -1420,27 +1539,36 @@ function createServer(engine, options = {}) {
       logger.debug(str);
     })
   );
-  app.openClaims = {};
   app.workflows = {};
   app.destroyed = false;
+  app.workloops = parseWorkloops(
+    options.workloopConfigs ?? get_default_workloop_config_default(options.maxWorkflows)
+  );
+  app.runWorkloopMap = {};
   app.server = app.listen(port);
   logger.success(`Worker ${app.id} listening on ${port}`);
   process.send?.("READY");
   router.get("/livez", healthcheck_default);
   router.get("/", healthcheck_default);
   app.options = options;
-  app.resumeWorkloop = () => {
+  app.resumeWorkloop = (workloop) => {
     if (options.noLoop || app.destroyed) {
       return;
     }
-    if (!app.workloop || app.workloop?.isStopped()) {
-      logger.info("Starting workloop");
-      app.workloop = workloop_default(
+    const targets = workloop ? [workloop] : app.workloops;
+    for (const w of targets) {
+      if (!w.hasCapacity()) {
+        continue;
+      }
+      if (!w.isStopped()) {
+        w.stop("restarting");
+      }
+      logger.info(`Starting workloop for ${w.id}`);
+      w.start(
         app,
         logger,
         options.backoff?.min || MIN_BACKOFF,
-        options.backoff?.max || MAX_BACKOFF,
-        options.maxWorkflows
+        options.backoff?.max || MAX_BACKOFF
       );
     }
   };
@@ -1485,8 +1613,16 @@ function createServer(engine, options = {}) {
           );
           delete app.workflows[id];
           runChannel.leave();
-          app.events.emit(INTERNAL_RUN_COMPLETE);
-          app.resumeWorkloop();
+          const owningWorkloop = app.runWorkloopMap[id];
+          if (owningWorkloop) {
+            owningWorkloop.activeRuns.delete(id);
+            delete app.runWorkloopMap[id];
+            app.events.emit(INTERNAL_RUN_COMPLETE);
+            app.resumeWorkloop(owningWorkloop);
+          } else {
+            app.events.emit(INTERNAL_RUN_COMPLETE);
+            app.resumeWorkloop();
+          }
         };
         const context = execute(
           runChannel,
@@ -1500,7 +1636,14 @@ function createServer(engine, options = {}) {
         app.workflows[id] = context;
       } catch (e) {
         delete app.workflows[id];
-        app.resumeWorkloop();
+        const owningWorkloop = app.runWorkloopMap[id];
+        if (owningWorkloop) {
+          owningWorkloop.activeRuns.delete(id);
+          delete app.runWorkloopMap[id];
+          app.resumeWorkloop(owningWorkloop);
+        } else {
+          app.resumeWorkloop();
+        }
         logger.error(`Unexpected error executing ${id}`);
         logger.error(e);
       }
@@ -1510,9 +1653,13 @@ function createServer(engine, options = {}) {
   };
   router.post("/claim", async (ctx) => {
     logger.info("triggering claim from POST request");
-    return claim_default(app, logger, {
-      maxWorkers: options.maxWorkflows
-    }).then(() => {
+    const promises = app.workloops.map((w) => {
+      if (w.hasCapacity()) {
+        return claim_default(app, w, logger);
+      }
+      return Promise.reject(new Error("Workloop at capacity"));
+    });
+    return Promise.any(promises).then(() => {
       logger.info("claim complete: 1 run claimed");
       ctx.body = "complete";
       ctx.status = 200;
@@ -1523,10 +1670,15 @@ function createServer(engine, options = {}) {
     });
   });
   app.claim = () => {
-    return claim_default(app, logger, {
-      maxWorkers: options.maxWorkflows
+    const promises = app.workloops.map((w) => {
+      if (w.hasCapacity()) {
+        return claim_default(app, w, logger);
+      }
+      return Promise.reject(new Error("Workloop at capacity"));
     });
+    return Promise.any(promises);
   };
+  app.pendingClaims = () => app.workloops.reduce((sum, w) => sum + Object.keys(w.openClaims).length, 0);
   app.destroy = () => destroy_default(app, logger);
   app.use(router.routes());
   if (options.lightning) {