fix(gateway): admit reconnect delivery drains (#104143)

src/plugin-sdk/delivery-queue-runtime.test.ts

@@ -1,21 +1,31 @@
 /**
  * Tests delivery queue runtime ordering and retry behavior.
  */
-import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
+import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
+import {
+  getActiveGatewayRootWorkCount,
+  resetGatewayWorkAdmission,
+  runWithGatewayRootWorkAdmission,
+  tryBeginGatewaySuspendAdmission,
+} from "../process/gateway-work-admission.js";
 
 const mocks = vi.hoisted(() => ({
   coreDrainPendingDeliveries: vi.fn(async () => {}),
   deliverOutboundPayloads: vi.fn(async () => []),
+  deliverRuntimeModuleLoads: 0,
 }));
 
 vi.mock("../infra/outbound/delivery-queue.js", () => ({
   drainPendingDeliveries: mocks.coreDrainPendingDeliveries,
 }));
 
-vi.mock("../infra/outbound/deliver-runtime.js", () => ({
-  deliverOutboundPayloads: mocks.deliverOutboundPayloads,
-  deliverOutboundPayloadsInternal: mocks.deliverOutboundPayloads,
-}));
+vi.mock("../infra/outbound/deliver-runtime.js", () => {
+  mocks.deliverRuntimeModuleLoads += 1;
+  return {
+    deliverOutboundPayloads: mocks.deliverOutboundPayloads,
+    deliverOutboundPayloadsInternal: mocks.deliverOutboundPayloads,
+  };
+});
 
 type DeliveryQueueRuntimeModule = typeof import("./delivery-queue-runtime.js");
 
@@ -32,14 +42,95 @@ beforeAll(async () => {
 });
 
 beforeEach(() => {
-  mocks.coreDrainPendingDeliveries.mockClear();
-  mocks.deliverOutboundPayloads.mockClear();
+  resetGatewayWorkAdmission();
+  mocks.coreDrainPendingDeliveries.mockReset().mockResolvedValue(undefined);
+  mocks.deliverOutboundPayloads.mockReset().mockResolvedValue([]);
   log.info.mockClear();
   log.warn.mockClear();
   log.error.mockClear();
 });
 
+afterEach(resetGatewayWorkAdmission);
+
 describe("plugin-sdk delivery queue drainPendingDeliveries", () => {
+  it("defers lazy runtime resolution and core draining while suspension is prepared", async () => {
+    const suspension = tryBeginGatewaySuspendAdmission(() => {});
+    expect(suspension?.commit()).toBe(true);
+
+    const pending = drainPendingDeliveries({
+      drainKey: "demo:test",
+      logLabel: "Demo reconnect drain",
+      cfg: {},
+      log,
+      selectEntry: () => ({ match: false }),
+    });
+    await Promise.resolve();
+
+    expect(mocks.deliverRuntimeModuleLoads).toBe(0);
+    expect(mocks.coreDrainPendingDeliveries).not.toHaveBeenCalled();
+
+    expect(suspension?.release()).toBe(true);
+    await pending;
+
+    expect(mocks.deliverRuntimeModuleLoads).toBe(1);
+    expect(mocks.coreDrainPendingDeliveries).toHaveBeenCalledOnce();
+    expect(getActiveGatewayRootWorkCount()).toBe(0);
+  });
+
+  it("counts a reconnect drain independently from its admitted parent", async () => {
+    let finishDrain = () => {};
+    const drainStarted = new Promise<void>((resolveStarted) => {
+      mocks.coreDrainPendingDeliveries.mockImplementationOnce(
+        () =>
+          new Promise<void>((resolveDrain) => {
+            finishDrain = resolveDrain;
+            resolveStarted();
+          }),
+      );
+    });
+    const deliver = vi.fn(async () => []);
+
+    await runWithGatewayRootWorkAdmission(async () => {
+      expect(getActiveGatewayRootWorkCount()).toBe(1);
+      const pending = drainPendingDeliveries({
+        drainKey: "demo:test",
+        logLabel: "Demo reconnect drain",
+        cfg: {},
+        log,
+        deliver,
+        selectEntry: () => ({ match: false }),
+      });
+      await drainStarted;
+
+      expect(getActiveGatewayRootWorkCount()).toBe(2);
+      finishDrain();
+      await pending;
+      expect(getActiveGatewayRootWorkCount()).toBe(1);
+    });
+
+    expect(getActiveGatewayRootWorkCount()).toBe(0);
+  });
+
+  it("releases its independent root when the core drain rejects", async () => {
+    mocks.coreDrainPendingDeliveries.mockImplementationOnce(async () => {
+      expect(getActiveGatewayRootWorkCount()).toBe(1);
+      throw new Error("drain failed");
+    });
+
+    await expect(
+      drainPendingDeliveries({
+        drainKey: "demo:test",
+        logLabel: "Demo reconnect drain",
+        cfg: {},
+        log,
+        deliver: vi.fn(async () => []),
+        selectEntry: () => ({ match: false }),
+      }),
+    ).rejects.toThrow("drain failed");
+
+    expect(getActiveGatewayRootWorkCount()).toBe(0);
+  });
+
   it("injects the lazy outbound deliver runtime when no deliver fn is provided", async () => {
     await drainPendingDeliveries({
       drainKey: "demo:test",

src/plugin-sdk/delivery-queue-runtime.ts

@@ -3,6 +3,7 @@ import {
   drainPendingDeliveries as coreDrainPendingDeliveries,
   type DeliverFn,
 } from "../infra/outbound/delivery-queue.js";
+import { runWithGatewayIndependentRootWorkAdmission } from "../process/gateway-work-admission.js";
 import { createLazyRuntimeModule } from "../shared/lazy-runtime.js";
 
 type DrainPendingDeliveriesOptions = Omit<
@@ -23,10 +24,13 @@ const loadOutboundDeliverRuntime = createLazyRuntimeModule(
  * loaded lazily so importing this SDK subpath does not eagerly bind send internals.
  */
 export async function drainPendingDeliveries(opts: DrainPendingDeliveriesOptions): Promise<void> {
-  const deliver =
-    opts.deliver ?? (await loadOutboundDeliverRuntime()).deliverOutboundPayloadsInternal;
-  await coreDrainPendingDeliveries({
-    ...opts,
-    deliver,
+  await runWithGatewayIndependentRootWorkAdmission(async () => {
+    // Keep lazy resolution and draining in one lease so suspension cannot split the handoff.
+    const deliver =
+      opts.deliver ?? (await loadOutboundDeliverRuntime()).deliverOutboundPayloadsInternal;
+    await coreDrainPendingDeliveries({
+      ...opts,
+      deliver,
+    });
   });
 }