@dudousxd/nestjs-durable-core 0.74.0 → 0.75.1

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,10 @@ 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;
+  /** Re-arms the deadline of every `drain()` currently waiting, when more wait is reported. */
+  drainRearms = /* @__PURE__ */ new Set();
   constructor(deps) {
     this.store = deps.store;
     this.pool = new TransportPool(deps.transports ?? (deps.transport ? [
@@ -3916,7 +3922,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,23 +4548,59 @@ 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;
+    if (Number.isFinite(spentMs) && spentMs > 0) {
+      this.drainSpentMs += spentMs;
+      for (const rearm of this.drainRearms) rearm();
+    }
+  }
+  /**
   * 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();
     if (this.inflight.size === 0) return;
+    const start = Date.now();
+    let rearm = /* @__PURE__ */ __name(() => {
+    }, "rearm");
+    let timeout;
     const timer = new Promise((resolve) => {
-      const t = setTimeout(resolve, timeoutMs);
-      t.unref?.();
+      rearm = /* @__PURE__ */ __name(() => {
+        clearTimeout(timeout);
+        const remaining = Math.max(0, start + timeoutMs - this.drainSpentMs - Date.now());
+        timeout = setTimeout(resolve, remaining);
+        timeout.unref?.();
+      }, "rearm");
+      rearm();
     });
-    await Promise.race([
-      Promise.allSettled([
-        ...this.inflight
-      ]),
-      timer
-    ]);
+    this.drainRearms.add(rearm);
+    try {
+      await Promise.race([
+        Promise.allSettled([
+          ...this.inflight
+        ]),
+        timer
+      ]);
+    } finally {
+      this.drainRearms.delete(rearm);
+      clearTimeout(timeout);
+    }
   }
   /**
   * Cancel in-flight runs that have outlived their workflow's `executionTimeout`. Call it from the
@@ -4591,6 +4634,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 +4650,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 +4680,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 +5116,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));