From 1b3af4565790986ad783f00625b6cb86ec6964a2 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Fri, 10 Jul 2026 22:33:33 -0700 Subject: [PATCH] fix(gateway): admit reconnect delivery drains (#104143) --- src/plugin-sdk/delivery-queue-runtime.test.ts | 105 ++++++++++++++++-- src/plugin-sdk/delivery-queue-runtime.ts | 14 ++- 2 files changed, 107 insertions(+), 12 deletions(-) diff --git a/src/plugin-sdk/delivery-queue-runtime.test.ts b/src/plugin-sdk/delivery-queue-runtime.test.ts index 7217c992bae2..f709cbbd72b9 100644 --- a/src/plugin-sdk/delivery-queue-runtime.test.ts +++ b/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((resolveStarted) => { + mocks.coreDrainPendingDeliveries.mockImplementationOnce( + () => + new Promise((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", diff --git a/src/plugin-sdk/delivery-queue-runtime.ts b/src/plugin-sdk/delivery-queue-runtime.ts index f94343ed2146..033382440a8a 100644 --- a/src/plugin-sdk/delivery-queue-runtime.ts +++ b/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 { - 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, + }); }); }