diff --git a/ts-sdk/src/coordinator/comm-channel.ts b/ts-sdk/src/coordinator/comm-channel.ts index 2c8eeb0a908b3..db9e629baa8a0 100644 --- a/ts-sdk/src/coordinator/comm-channel.ts +++ b/ts-sdk/src/coordinator/comm-channel.ts @@ -51,6 +51,10 @@ export interface RequestOptions { timeoutMs?: number; } +export interface ConnectOptions { + timeoutMs?: number; +} + export const COORDINATOR_REQUEST_TIMEOUT_MS = 30_000; export class CommChannel { @@ -81,11 +85,22 @@ export class CommChannel { sock.on("error", (err) => this.handleClose(err)); } - /** Connect and wait for the supervisor's greeting; rejects if the - * socket dies before it arrives. */ - static async connect(addr: string, logs: LogChannel | null = null): Promise { + /** Connect and wait for the supervisor's greeting; rejects if the socket dies, or no greeting + * arrives within `timeoutMs`, before it arrives. A timeout destroys the socket so the runtime + * never waits on a wedged or misbehaving supervisor forever. */ + static async connect( + addr: string, + logs: LogChannel | null = null, + opts: ConnectOptions = {}, + ): Promise { const sock = await connectTcp(addr); const channel = new CommChannel(sock, logs); + const timeoutMs = opts.timeoutMs ?? COORDINATOR_REQUEST_TIMEOUT_MS; + channel.greeting.rejectAfter(timeoutMs, () => { + const err = new Error(`Timed out waiting for supervisor greeting after ${timeoutMs} ms`); + sock.destroy(err); + return err; + }); const firstFrame = await channel.greeting.promise; return { channel, firstFrame }; } diff --git a/ts-sdk/tests/coordinator/comm-channel.test.ts b/ts-sdk/tests/coordinator/comm-channel.test.ts index ae71e659b7ff1..21444cd292f0d 100644 --- a/ts-sdk/tests/coordinator/comm-channel.test.ts +++ b/ts-sdk/tests/coordinator/comm-channel.test.ts @@ -18,9 +18,15 @@ */ import { EventEmitter } from "node:events"; +import type { Socket } from "node:net"; import { afterEach, describe, expect, it, vi } from "vitest"; import { CommChannel, COORDINATOR_REQUEST_TIMEOUT_MS } from "../../src/coordinator/comm-channel.js"; import { encodeResponse } from "../../src/coordinator/frames.js"; +import { connectTcp } from "../../src/coordinator/tcp-connect.js"; + +vi.mock("../../src/coordinator/tcp-connect.js", () => ({ + connectTcp: vi.fn(), +})); class FakeSocket extends EventEmitter { writeCallback: ((err?: Error) => void) | undefined; @@ -150,6 +156,59 @@ describe("CommChannel", () => { await vi.advanceTimersByTimeAsync(10); }); + it("times out waiting for the supervisor greeting and destroys the socket", async () => { + vi.useFakeTimers(); + const sock = new FakeSocket(); + vi.mocked(connectTcp).mockResolvedValue(sock as unknown as Socket); + + const connection = CommChannel.connect("127.0.0.1:0", null, { timeoutMs: 10 }); + const assertion = expect(connection).rejects.toThrow( + "Timed out waiting for supervisor greeting after 10 ms", + ); + await vi.advanceTimersByTimeAsync(10); + + await assertion; + expect(sock.destroy).toHaveBeenCalledWith(expect.any(Error)); + }); + + it("uses the default request timeout for the greeting when no override is given", async () => { + vi.useFakeTimers(); + const sock = new FakeSocket(); + vi.mocked(connectTcp).mockResolvedValue(sock as unknown as Socket); + + const connection = CommChannel.connect("127.0.0.1:0", null); + const rejection = vi.fn(); + connection.catch(rejection); + const assertion = expect(connection).rejects.toThrow( + `Timed out waiting for supervisor greeting after ${COORDINATOR_REQUEST_TIMEOUT_MS} ms`, + ); + + await vi.advanceTimersByTimeAsync(COORDINATOR_REQUEST_TIMEOUT_MS - 1); + expect(rejection).not.toHaveBeenCalled(); + await vi.advanceTimersByTimeAsync(1); + + await assertion; + expect(rejection).toHaveBeenCalledOnce(); + }); + + it("clears the greeting timeout once the greeting arrives", async () => { + vi.useFakeTimers(); + const sock = new FakeSocket(); + vi.mocked(connectTcp).mockResolvedValue(sock as unknown as Socket); + + const connection = CommChannel.connect("127.0.0.1:0", null, { timeoutMs: 10 }); + // Let the mocked connectTcp() promise settle and the "data" listener attach + // before the greeting frame arrives. + await vi.advanceTimersByTimeAsync(0); + sock.emit("data", encodeResponse(0, { type: "StartupDetails" })); + + await expect(connection).resolves.toMatchObject({ + firstFrame: { body: { type: "StartupDetails" } }, + }); + await vi.advanceTimersByTimeAsync(10); + expect(sock.destroy).not.toHaveBeenCalled(); + }); + it("does not log clean close events", async () => { const logs = { debug: vi.fn(), warning: vi.fn() }; const sock = new FakeSocket();