@dudousxd/nestjs-durable-core 0.74.0 → 0.75.0

package/dist/index.cjs

@@ -2350,12 +2350,14 @@ var ScheduleClient = class {
     const store = this.requireStore();
     const started = [];
     for (let round = 0; round < 50; round++) {
+      if (this.host.isDraining?.()) break;
       const due = await store.listSchedules({
         dueBy: nowMs,
         namespace: this.options.namespace,
         limit: TICK_BATCH
       });
       for (const record of due) {
+        if (this.host.isDraining?.()) return started;
         const runId = await this.fire(record, nowMs);
         if (runId) started.push(runId);
       }
@@ -3868,6 +3870,8 @@ var WorkflowEngine = class {
   /** Executions currently in flight, so a graceful shutdown can wait for them to settle. */
   inflight = /* @__PURE__ */ new Set();
   draining = false;
+  /** Wait already spent on this shutdown before `drain()` (see `beginDrain`). */
+  drainSpentMs = 0;
   constructor(deps) {
     this.store = deps.store;
     this.pool = new TransportPool(deps.transports ?? (deps.transport ? [
@@ -3916,7 +3920,8 @@ var WorkflowEngine = class {
     });
     this.schedules = new ScheduleClient(this.store, {
       start: /* @__PURE__ */ __name((workflow, input, runId, opts) => this.start(workflow, input, runId, opts), "start"),
-      getRun: /* @__PURE__ */ __name((runId) => this.store.getRun(runId), "getRun")
+      getRun: /* @__PURE__ */ __name((runId) => this.store.getRun(runId), "getRun"),
+      isDraining: /* @__PURE__ */ __name(() => this.draining, "isDraining")
     }, {
       namespace: this.namespace,
       clock: this.clock
@@ -4541,15 +4546,34 @@ var WorkflowEngine = class {
     return p;
   }
   /**
+  * Stop picking up new work without waiting: recovery, timer, schedule and timeout sweeps (and
+  * their per-run lock loops) start nothing once this is called, and {@link runOne} returns null. A
+  * poller calls it on its own shutdown so an in-flight sweep stops locking runs before it is awaited.
+  *
+  * `spentMs` is wait time ALREADY spent on this shutdown (the poller awaiting its in-flight sweep).
+  * It is deducted from the budget of a later {@link drain}, so that wait plus the drain stay within
+  * one `shutdownTimeoutMs`. Only time actually spent waiting is charged: time that merely passes
+  * between this call and `drain()` (other hooks, the HTTP server closing) never is. Calls add up.
+  */
+  beginDrain(spentMs = 0) {
+    this.draining = true;
+    this.drainSpentMs += Math.max(0, spentMs);
+  }
+  /**
   * Graceful shutdown: stop picking up new runs (recovery/timer become no-ops) and wait for
   * in-flight executions to settle, up to `timeoutMs`. Call from your app's shutdown hook so a
   * deploy hands off cleanly instead of leaving runs to the lease timeout.
+  *
+  * `timeoutMs` is the total budget for waiting on in-flight runs; any wait already reported through
+  * {@link beginDrain} (`spentMs`) is deducted, so a shutdown that already waited on the timer poll
+  * for part of it waits only for the remainder. With nothing reported, the full `timeoutMs` applies.
   */
   async drain(timeoutMs = 1e4) {
-    this.draining = true;
+    this.beginDrain();
+    const budget = Math.max(0, timeoutMs - this.drainSpentMs);
     if (this.inflight.size === 0) return;
     const timer = new Promise((resolve) => {
-      const t = setTimeout(resolve, timeoutMs);
+      const t = setTimeout(resolve, budget);
       t.unref?.();
     });
     await Promise.race([
@@ -4591,6 +4615,7 @@ var WorkflowEngine = class {
   */
   async sweepTimeouts(now = this.clock()) {
     for (const reg of new Set(this.latest.values())) {
+      if (this.draining) return;
       if (reg.executionTimeoutMs == null) continue;
       const deadline = now - reg.executionTimeoutMs;
       const inflight = [
@@ -4606,6 +4631,7 @@ var WorkflowEngine = class {
         })
       ];
       for (const run of inflight) {
+        if (this.draining) return;
         if (run.createdAt.getTime() > deadline) continue;
         const error2 = {
           message: "execution timeout",
@@ -4635,6 +4661,7 @@ var WorkflowEngine = class {
     if (this.draining) return [];
     const results = [];
     for (const listed of await this.store.listIncompleteRuns(this.namespace)) {
+      if (this.draining) break;
       const acquired = await this.store.tryLockRun(listed.id, this.instanceId, nowMs + this.leaseMs, nowMs);
       if (!acquired) continue;
       const run = await this.store.getRun(listed.id);
@@ -5070,8 +5097,13 @@ var WorkflowEngine = class {
     if (this.draining) return [];
     const results = [];
     for (const run of runs) {
+      if (this.draining) break;
       const acquired = await this.store.tryLockRun(run.id, this.instanceId, nowMs + this.leaseMs, nowMs);
       if (!acquired) continue;
+      if (this.draining) {
+        await this.store.releaseRunLock(run.id);
+        break;
+      }
       try {
         const settled = onLocked ? await onLocked(run) : void 0;
         results.push(settled ?? await this.resumeImplicit(run.id));