@dudousxd/nestjs-durable-core 0.75.0 → 0.75.1

package/dist/index.cjs

@@ -3872,6 +3872,8 @@ var WorkflowEngine = class {
   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 ? [
@@ -4557,7 +4559,10 @@ var WorkflowEngine = class {
   */
   beginDrain(spentMs = 0) {
     this.draining = true;
-    this.drainSpentMs += Math.max(0, spentMs);
+    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
@@ -4570,18 +4575,32 @@ var WorkflowEngine = class {
   */
   async drain(timeoutMs = 1e4) {
     this.beginDrain();
-    const budget = Math.max(0, timeoutMs - this.drainSpentMs);
     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, budget);
-      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