fix(relay): tolerate slow connections with coordinated deadlines

apps/relay/src/http-exchange.test.ts

@@ -0,0 +1,87 @@
+/// <reference types="bun" />
+import { afterEach, beforeEach, expect, jest, mock, test } from "bun:test";
+import { HttpExchanges } from "./http-exchange";
+
+beforeEach(() => jest.useFakeTimers());
+afterEach(() => jest.useRealTimers());
+
+const request = {
+	method: "GET",
+	pathWithQuery: "/health",
+	headers: {},
+	body: new Uint8Array(),
+};
+
+test("a slow dial still leaves time for the HTTP response", async () => {
+	const exchanges = new HttpExchanges();
+	const result = exchanges.begin("ticket", request);
+	const dial = { send: mock(() => {}), close: mock(() => {}) };
+	jest.advanceTimersByTime(32_000);
+	expect(exchanges.has("ticket")).toBe(true);
+	exchanges.onDialConnect("ticket", dial);
+	expect(dial.send).toHaveBeenCalled();
+	jest.advanceTimersByTime(25_000);
+	exchanges.onDialMessage(
+		"ticket",
+		dial,
+		JSON.stringify({ type: "http:response", status: 200, headers: {} }),
+	);
+	exchanges.onDialMessage("ticket", dial, JSON.stringify({ type: "http:end" }));
+	expect(await result).toMatchObject({ ok: true, status: 200 });
+	expect(exchanges.has("ticket")).toBe(false);
+	expect(dial.close).toHaveBeenCalledTimes(1);
+});
+
+test("an unanswered HTTP exchange still expires", async () => {
+	const exchanges = new HttpExchanges();
+	const result = exchanges.begin("ticket", request);
+	jest.advanceTimersByTime(65_000);
+	expect(await result).toEqual({ ok: false, reason: "timeout" });
+	expect(exchanges.has("ticket")).toBe(false);
+});
+
+test("a host-reported dial failure fails immediately", async () => {
+	const exchanges = new HttpExchanges();
+	const result = exchanges.begin("ticket", request);
+	exchanges.fail("ticket");
+	expect(await result).toEqual({ ok: false, reason: "dial-failed" });
+	expect(exchanges.has("ticket")).toBe(false);
+});
+
+test("completion removes the timeout and ignores late frames", async () => {
+	const exchanges = new HttpExchanges();
+	const result = exchanges.begin("ticket", request);
+	const dial = { send: mock(() => {}), close: mock(() => {}) };
+	exchanges.onDialConnect("ticket", dial);
+	exchanges.onDialMessage(
+		"ticket",
+		dial,
+		JSON.stringify({ type: "http:response", status: 201, headers: {} }),
+	);
+	exchanges.onDialMessage(
+		"ticket",
+		dial,
+		new TextEncoder().encode("response").buffer,
+	);
+	exchanges.onDialMessage("ticket", dial, JSON.stringify({ type: "http:end" }));
+	expect(await result).toMatchObject({
+		ok: true,
+		status: 201,
+		body: new TextEncoder().encode("response"),
+	});
+	jest.advanceTimersByTime(100_000);
+	exchanges.onDialMessage("ticket", dial, JSON.stringify({ type: "http:end" }));
+	expect(dial.close).toHaveBeenCalledTimes(1);
+});
+
+test("shutdown settles all pending exchanges", async () => {
+	const exchanges = new HttpExchanges();
+	const first = exchanges.begin("first", request);
+	const second = exchanges.begin("second", request);
+	exchanges.abortAll();
+	expect(await first).toEqual({ ok: false, reason: "timeout" });
+	expect(await second).toEqual({ ok: false, reason: "timeout" });
+	expect(exchanges.has("first")).toBe(false);
+	expect(exchanges.has("second")).toBe(false);
+	jest.advanceTimersByTime(100_000);
+});

apps/relay/src/http-exchange.ts

@@ -1,9 +1,10 @@
-import type {
-	HttpDialFrame,
-	HttpResponseHeader,
+import {
+	DIAL_TIMEOUT_MS,
+	type HttpDialFrame,
+	type HttpResponseHeader,
 } from "@superset/shared/tunnel-protocol";
 
-const EXCHANGE_TIMEOUT_MS = 30_000;
+const EXCHANGE_TIMEOUT_MS = DIAL_TIMEOUT_MS + 30_000;
 // Chunked below the Durable Object's per-message ceiling.
 const BODY_CHUNK_BYTES = 256 * 1024;
 

packages/host-service/src/tunnel/tunnel-client.integration.test.ts

@@ -0,0 +1,129 @@
+import { expect, test } from "bun:test";
+import { TunnelClient } from "./tunnel-client";
+
+async function exerciseDial({
+	delayMs,
+	rejectFirst = false,
+}: {
+	delayMs: number;
+	rejectFirst?: boolean;
+}) {
+	let attempts = 0;
+	let requests = 0;
+	let resultResolve: (value: string) => void = () => {};
+	const result = new Promise<string>((resolve) => {
+		resultResolve = resolve;
+	});
+	const local = Bun.serve({
+		port: 0,
+		fetch(request) {
+			requests++;
+			expect(request.headers.get("authorization")).toBe("Bearer test-secret");
+			return new Response("proxied successfully");
+		},
+	});
+	const relay = Bun.serve<{ control: boolean }>({
+		port: 0,
+		async fetch(request, server) {
+			const control = new URL(request.url).pathname === "/v2/control";
+			if (!control) {
+				attempts++;
+				if (rejectFirst && attempts === 1) {
+					return new Response("temporary failure", { status: 503 });
+				}
+				await Bun.sleep(delayMs);
+			}
+			if (server.upgrade(request, { data: { control } })) return;
+			return new Response("upgrade failed", { status: 400 });
+		},
+		websocket: {
+			open(socket) {
+				if (socket.data.control) {
+					socket.send(
+						JSON.stringify({
+							type: "stream:dial",
+							ticket: "test-ticket",
+							kind: "http",
+							path: "/health",
+						}),
+					);
+				} else {
+					socket.send(
+						JSON.stringify({
+							type: "http:request",
+							method: "GET",
+							path: "/health",
+							headers: {},
+						}),
+					);
+					socket.send(JSON.stringify({ type: "http:end" }));
+				}
+			},
+			message(socket, data) {
+				if (typeof data !== "string") return;
+				const frame = JSON.parse(data);
+				if (frame.type === "ping")
+					socket.send(JSON.stringify({ type: "pong" }));
+				if (frame.type === "stream:dial-failed") resultResolve("failed");
+				if (frame.type === "http:response") expect(frame.status).toBe(200);
+				if (frame.type === "http:end") resultResolve("success");
+			},
+		},
+	});
+	const client = new TunnelClient({
+		relayUrl: `http://127.0.0.1:${relay.port}`,
+		hostId: "org:host",
+		getAuthToken: async () => "test-token",
+		localPort: local.port ?? 0,
+		hostServiceSecret: "test-secret",
+	});
+	const start = performance.now();
+	let timer: ReturnType<typeof setTimeout> | undefined;
+	try {
+		await client.connect();
+		const outcome = await Promise.race([
+			result,
+			new Promise<never>((_, reject) => {
+				timer = setTimeout(
+					() => reject(new Error("Dial did not settle")),
+					40_000,
+				);
+			}),
+		]);
+		return {
+			outcome,
+			attempts,
+			requests,
+			elapsedMs: performance.now() - start,
+		};
+	} finally {
+		clearTimeout(timer);
+		client.close();
+		relay.stop(true);
+		local.stop(true);
+	}
+}
+
+test("real sockets proxy HTTP after a 25-second opening handshake", async () => {
+	const result = await exerciseDial({ delayMs: 25_000 });
+	expect(result.outcome).toBe("success");
+	expect(result.attempts).toBe(1);
+	expect(result.requests).toBe(1);
+	expect(result.elapsedMs).toBeGreaterThanOrEqual(25_000);
+}, 45_000);
+
+test("real sockets recover from a rejected upgrade and a slow retry", async () => {
+	const result = await exerciseDial({ delayMs: 4_500, rejectFirst: true });
+	expect(result.outcome).toBe("success");
+	expect(result.attempts).toBe(2);
+	expect(result.requests).toBe(1);
+}, 45_000);
+
+test("real sockets report a stalled upgrade without forwarding the request", async () => {
+	const result = await exerciseDial({ delayMs: 35_000 });
+	expect(result.outcome).toBe("failed");
+	expect(result.attempts).toBe(1);
+	expect(result.requests).toBe(0);
+	expect(result.elapsedMs).toBeGreaterThanOrEqual(29_900);
+	expect(result.elapsedMs).toBeLessThan(35_000);
+}, 45_000);

packages/host-service/src/tunnel/tunnel-client.test.ts

@@ -0,0 +1,161 @@
+import {
+	afterEach,
+	beforeEach,
+	expect,
+	jest,
+	mock,
+	spyOn,
+	test,
+} from "bun:test";
+import { TunnelClient } from "./tunnel-client";
+
+class FakeSocket extends EventTarget {
+	static instances: FakeSocket[] = [];
+	static OPEN = 1;
+	static CONNECTING = 0;
+	readyState = 0;
+	binaryType = "blob";
+	close = mock(() => {
+		this.readyState = 3;
+		this.dispatchEvent(new Event("close"));
+	});
+
+	constructor(readonly url: string) {
+		super();
+		FakeSocket.instances.push(this);
+	}
+
+	static at(index: number) {
+		const socket = FakeSocket.instances[index];
+		if (!socket) throw new Error(`Missing socket ${index}`);
+		return socket;
+	}
+
+	open() {
+		this.readyState = 1;
+		this.dispatchEvent(new Event("open"));
+	}
+}
+
+const originalWebSocket = globalThis.WebSocket;
+let elapsed = 0;
+
+beforeEach(() => {
+	jest.useFakeTimers();
+	elapsed = 0;
+	spyOn(performance, "now").mockImplementation(() => elapsed);
+	FakeSocket.instances = [];
+	globalThis.WebSocket = FakeSocket as unknown as typeof WebSocket;
+});
+
+afterEach(() => {
+	globalThis.WebSocket = originalWebSocket;
+	mock.restore();
+	jest.useRealTimers();
+});
+
+function advance(ms: number) {
+	elapsed += ms;
+	jest.advanceTimersByTime(ms);
+}
+
+function startDial() {
+	const client = new TunnelClient({
+		relayUrl: "https://relay.example",
+		hostId: "org:host",
+		getAuthToken: async () => null,
+		localPort: 1234,
+		hostServiceSecret: "test",
+	});
+	const attach = mock(() => {});
+	const failed = spyOn(
+		client as unknown as { reportDialFailed(ticket: string): void },
+		"reportDialFailed",
+	).mockImplementation(() => {});
+	// biome-ignore lint/complexity/useLiteralKeys: Bracket access permits testing the private dial without a control connection.
+	client["dialRelay"]("ticket", attach);
+	return { attach, failed };
+}
+
+test("a slow dial opens after the old timeout without being restarted", () => {
+	const { attach, failed } = startDial();
+	advance(25_000);
+	expect(FakeSocket.instances).toHaveLength(1);
+	expect(FakeSocket.at(0).close).not.toHaveBeenCalled();
+	FakeSocket.at(0).open();
+	advance(60_000);
+	expect(attach).toHaveBeenCalledTimes(1);
+	expect(failed).not.toHaveBeenCalled();
+});
+
+test("an explicit failure retries within the original total budget", () => {
+	const { attach, failed } = startDial();
+	advance(20_000);
+	FakeSocket.at(0).close();
+	expect(FakeSocket.instances).toHaveLength(2);
+	advance(9_999);
+	expect(failed).not.toHaveBeenCalled();
+	advance(1);
+	expect(failed).toHaveBeenCalledTimes(1);
+	expect(FakeSocket.at(1).close).toHaveBeenCalledTimes(1);
+	FakeSocket.at(1).open();
+	expect(attach).not.toHaveBeenCalled();
+});
+
+test("a stalled dial stops at 30 seconds without starting a doomed retry", () => {
+	const { failed } = startDial();
+	advance(30_000);
+	expect(failed).toHaveBeenCalledTimes(1);
+	expect(FakeSocket.instances).toHaveLength(1);
+	expect(FakeSocket.at(0).close).toHaveBeenCalledTimes(1);
+});
+
+test("an opened ticket is never retried when its stream closes", () => {
+	const { attach, failed } = startDial();
+	FakeSocket.at(0).open();
+	FakeSocket.at(0).close();
+	advance(60_000);
+	expect(FakeSocket.instances).toHaveLength(1);
+	expect(attach).toHaveBeenCalledTimes(1);
+	expect(failed).not.toHaveBeenCalled();
+});
+
+test("a successful second attempt cancels the shared deadline", () => {
+	const { attach, failed } = startDial();
+	advance(10_000);
+	FakeSocket.at(0).close();
+	advance(19_999);
+	FakeSocket.at(1).open();
+	advance(60_000);
+	expect(attach).toHaveBeenCalledTimes(1);
+	expect(failed).not.toHaveBeenCalled();
+	expect(FakeSocket.at(1).close).not.toHaveBeenCalled();
+});
+
+test("two explicit failures stop without a third attempt", () => {
+	const { attach, failed } = startDial();
+	FakeSocket.at(0).close();
+	FakeSocket.at(1).close();
+	advance(60_000);
+	expect(FakeSocket.instances).toHaveLength(2);
+	expect(failed).toHaveBeenCalledTimes(1);
+	expect(attach).not.toHaveBeenCalled();
+});
+
+test("wall-clock changes cannot extend the dial budget", () => {
+	const { failed } = startDial();
+	advance(20_000);
+	spyOn(Date, "now").mockReturnValue(-1_000_000);
+	FakeSocket.at(0).close();
+	advance(10_000);
+	expect(failed).toHaveBeenCalledTimes(1);
+	expect(FakeSocket.instances).toHaveLength(2);
+});
+
+test("timer rounding cannot start a retry with a fractional budget", () => {
+	const { failed } = startDial();
+	elapsed = 29_999.5;
+	jest.advanceTimersByTime(30_000);
+	expect(failed).toHaveBeenCalledTimes(1);
+	expect(FakeSocket.instances).toHaveLength(1);
+});