From 65fa1f9a7bd19bdd625dabd434e0edc3920166ac Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Mon, 3 Aug 2026 14:58:13 -0700 Subject: [PATCH] fix(signal): drain accepted websocket messages before shutdown (#118975) * fix(signal): drain accepted websocket messages before shutdown * test(signal): avoid returning timer handles from Promise executors --- .../signal/src/client-container.test.ts | 170 +++++++++++++++++- extensions/signal/src/client-container.ts | 20 ++- 2 files changed, 187 insertions(+), 3 deletions(-) diff --git a/extensions/signal/src/client-container.test.ts b/extensions/signal/src/client-container.test.ts index f39bffd940db..d409fa78af4e 100644 --- a/extensions/signal/src/client-container.test.ts +++ b/extensions/signal/src/client-container.test.ts @@ -216,9 +216,17 @@ function delayedBodyStream( }; } const wsMockState = vi.hoisted(() => ({ - behavior: "close" as "close" | "open" | "error" | "message" | "unexpected-response", + behavior: "close" as + | "close" + | "open" + | "error" + | "message" + | "buffered-message" + | "pending" + | "unexpected-response", urls: [] as string[], options: [] as Array<{ maxPayload?: number; handshakeTimeout?: number } | undefined>, + terminations: 0, })); beforeEach(() => { @@ -226,6 +234,7 @@ beforeEach(() => { wsMockState.behavior = "close"; wsMockState.urls = []; wsMockState.options = []; + wsMockState.terminations = 0; }); function requireFetchCall(index = 0): [RequestInfo | URL, RequestInit] { @@ -266,6 +275,7 @@ function expectMockLogNotContains(mock: ReturnType, expected: stri vi.mock("ws", () => ({ default: class MockWebSocket { private handlers = new Map void>>(); + private bufferedMessageFlushed = false; constructor(url: string | URL, options?: { maxPayload?: number; handshakeTimeout?: number }) { wsMockState.urls.push(String(url)); @@ -281,6 +291,11 @@ vi.mock("ws", () => ({ } else if (wsMockState.behavior === "message") { this.emit("message", Buffer.from('{"envelope":{"timestamp":1}}')); this.emit("close", 1000, Buffer.from("done")); + } else if (wsMockState.behavior === "buffered-message") { + this.emit("open"); + this.emit("message", Buffer.from('{"envelope":{"timestamp":1}}')); + } else if (wsMockState.behavior === "pending") { + // Keep the opening handshake unresolved until shutdown closes it. } else { this.emit("close", 1000, Buffer.from("done")); } @@ -306,10 +321,17 @@ vi.mock("ws", () => ({ } close() { + if (wsMockState.behavior === "buffered-message" && !this.bufferedMessageFlushed) { + this.bufferedMessageFlushed = true; + // ws flushes already-buffered receiver frames before its final close event. + this.emit("message", Buffer.from('{"envelope":{"timestamp":2}}')); + } this.emit("close", 1000, Buffer.from("done")); } - terminate() {} + terminate() { + wsMockState.terminations += 1; + } private emit(event: string, ...args: unknown[]) { for (const handler of this.handlers.get(event) ?? []) { @@ -1625,6 +1647,150 @@ describe("streamContainerEvents", () => { expect(removeEventListener).toHaveBeenCalledWith("abort", abortHandler); }); + it("drains accepted and socket-buffered receive events before resolving shutdown", async () => { + wsMockState.behavior = "buffered-message"; + const abort = new AbortController(); + const removeEventListener = vi.spyOn(abort.signal, "removeEventListener"); + let releaseFirst!: () => void; + const firstDelivery = new Promise((resolve) => { + releaseFirst = resolve; + }); + let markFirstStarted!: () => void; + const firstStarted = new Promise((resolve) => { + markFirstStarted = resolve; + }); + const timestamps: number[] = []; + const stream = streamContainerEvents({ + baseUrl: "http://localhost:8080", + abortSignal: abort.signal, + timeoutMs: 0, + onEvent: async (event) => { + timestamps.push(event.envelope?.timestamp ?? 0); + if (timestamps.length === 1) { + markFirstStarted(); + await firstDelivery; + } + }, + }); + let settled = false; + void stream.then(() => { + settled = true; + }); + + await firstStarted; + abort.abort(); + const settledBeforeDrain = await Promise.race([ + stream.then(() => true), + new Promise((resolve) => { + setTimeout(() => resolve(false), 10); + }), + ]); + expect(settledBeforeDrain).toBe(false); + expect(settled).toBe(false); + + releaseFirst(); + await expect(stream).resolves.toBeUndefined(); + expect(timestamps).toEqual([1, 2]); + expect(removeEventListener).toHaveBeenCalledWith("abort", expect.any(Function)); + expect(wsMockState.terminations).toBe(0); + }); + + it("propagates a receive-handler failure that settles during shutdown", async () => { + wsMockState.behavior = "buffered-message"; + const abort = new AbortController(); + const appendError = new Error("durable append failed during shutdown"); + let rejectDelivery!: (error: Error) => void; + const delivery = new Promise((_resolve, reject) => { + rejectDelivery = reject; + }); + let markFirstStarted!: () => void; + const firstStarted = new Promise((resolve) => { + markFirstStarted = resolve; + }); + const stream = streamContainerEvents({ + baseUrl: "http://localhost:8080", + abortSignal: abort.signal, + onEvent: async () => { + markFirstStarted(); + await delivery; + }, + }); + + await firstStarted; + abort.abort(); + rejectDelivery(appendError); + await expect(stream).rejects.toBe(appendError); + expect(wsMockState.terminations).toBe(0); + }); + + it("bounds a stalled shutdown drain and records the accepted-message loss risk", async () => { + vi.useFakeTimers(); + try { + wsMockState.behavior = "buffered-message"; + const abort = new AbortController(); + const error = vi.fn(); + let settled = false; + const stream = streamContainerEvents({ + baseUrl: "http://localhost:8080", + abortSignal: abort.signal, + timeoutMs: 0, + onEvent: async () => await new Promise(() => {}), + logger: { error }, + }).then(() => { + settled = true; + }); + + await vi.advanceTimersByTimeAsync(0); + abort.abort(); + await vi.advanceTimersByTimeAsync(1_499); + expect(settled).toBe(false); + expect(error).not.toHaveBeenCalled(); + + await vi.advanceTimersByTimeAsync(1); + await expect(stream).resolves.toBeUndefined(); + expect(error).toHaveBeenCalledWith(expect.stringMatching(/receive events.*may be lost/i)); + expect(wsMockState.terminations).toBe(1); + expect(vi.getTimerCount()).toBe(0); + } finally { + vi.useRealTimers(); + } + }); + + it("closes immediately when aborted before the opening handshake completes", async () => { + wsMockState.behavior = "pending"; + const abort = new AbortController(); + const removeEventListener = vi.spyOn(abort.signal, "removeEventListener"); + const stream = streamContainerEvents({ + baseUrl: "http://localhost:8080", + abortSignal: abort.signal, + onEvent: vi.fn(), + }); + + abort.abort(); + await expect(stream).resolves.toBeUndefined(); + expect(removeEventListener).toHaveBeenCalledWith("abort", expect.any(Function)); + expect(wsMockState.terminations).toBe(0); + }); + + it("handles an already-aborted signal without leaving its connection pending", async () => { + wsMockState.behavior = "pending"; + const abort = new AbortController(); + abort.abort(); + const result = await Promise.race([ + streamContainerEvents({ + baseUrl: "http://localhost:8080", + abortSignal: abort.signal, + onEvent: vi.fn(), + }).then(() => "closed"), + new Promise((resolve) => { + setTimeout(() => resolve("still pending"), 25); + }), + ]); + + expect(result).toBe("closed"); + expect(wsMockState.terminations).toBe(0); + }); + it("propagates receive-handler failures to the stream", async () => { wsMockState.behavior = "message"; const appendError = new Error("durable append failed"); diff --git a/extensions/signal/src/client-container.ts b/extensions/signal/src/client-container.ts index 3ea54fe7cf76..dbf81d234591 100644 --- a/extensions/signal/src/client-container.ts +++ b/extensions/signal/src/client-container.ts @@ -63,6 +63,7 @@ const SIGNAL_REST_SUCCESS_RESPONSE_MAX_BYTES = 16 * 1024 * 1024; // Receive envelopes contain metadata only; cap frames, and do not let upgrades block reconnect. const WS_MAX_PAYLOAD = 1024 * 1024; const WS_HANDSHAKE_MS = 30_000; +const WS_SHUTDOWN_DRAIN_TIMEOUT_MS = 1_500; // Outbound file paths are converted to base64 before posting to the container. Cap // reads to the same default the native signal send path uses (8 MiB) so a path to a // huge or symlinked file cannot OOM the gateway before encoding. @@ -445,8 +446,13 @@ export async function streamContainerEvents(params: { let settled = false; let eventChain = Promise.resolve(); let abortHandler: (() => void) | undefined; + let shutdownDrainTimer: ReturnType | undefined; const cleanup = () => { + if (shutdownDrainTimer) { + clearTimeout(shutdownDrainTimer); + shutdownDrainTimer = undefined; + } if (abortHandler) { params.abortSignal?.removeEventListener("abort", abortHandler); abortHandler = undefined; @@ -532,10 +538,22 @@ export async function streamContainerEvents(params: { if (params.abortSignal) { abortHandler = () => { log("[signal-ws] aborted, closing connection"); + // Arm before close: ws can synchronously flush buffered messages and emit + // close, whose final eventChain owns every accepted durable admission. + shutdownDrainTimer = setTimeout(() => { + logError( + "[signal-ws] shutdown timed out draining accepted receive events; messages may be lost", + ); + ws.terminate(); + resolveOnce(); + }, WS_SHUTDOWN_DRAIN_TIMEOUT_MS); + shutdownDrainTimer.unref?.(); ws.close(); - resolveOnce(); }; params.abortSignal.addEventListener("abort", abortHandler, { once: true }); + if (params.abortSignal.aborted) { + abortHandler(); + } } }); }