Various improvements to reconnecting websockets + More tests

Author: EhabY

src/api/coderApi.ts

@@ -240,6 +240,11 @@ export class CoderApi extends Api {
 				socketFactory,
 				this.output,
 				configs.apiRoute,
+				undefined,
+				() =>
+					this.reconnectingSockets.delete(
+						reconnectingSocket as ReconnectingWebSocket<unknown>,
+					),
 			);
 
 			this.reconnectingSockets.add(

src/websocket/reconnectingWebSocket.ts

@@ -1,4 +1,5 @@
 import type { WebSocketEventType } from "coder/site/src/utils/OneWayWebSocket";
+import type { CloseEvent } from "ws";
 
 import type { Logger } from "../logging/logger";
 
@@ -18,9 +19,6 @@ export type ReconnectingWebSocketOptions = {
 // 403 Forbidden, 410 Gone, 426 Upgrade Required, 1002/1003 Protocol errors
 const UNRECOVERABLE_CLOSE_CODES = new Set([403, 410, 426, 1002, 1003]);
 
-// Custom close code for intentional reconnection (4000-4999 range is for private use)
-const CLOSE_CODE_RECONNECTING = 4000;
-
 export class ReconnectingWebSocket<TData = unknown>
 	implements UnidirectionalStream<TData>
 {
@@ -40,12 +38,15 @@ export class ReconnectingWebSocket<TData = unknown>
 	#reconnectTimeoutId: NodeJS.Timeout | null = null;
 	#isDisposed = false;
 	#isConnecting = false;
+	#pendingReconnect = false;
+	readonly #onDispose?: () => void;
 
 	private constructor(
 		socketFactory: SocketFactory<TData>,
 		logger: Logger,
 		apiRoute: string,
 		options: ReconnectingWebSocketOptions = {},
+		onDispose?: () => void,
 	) {
 		this.#socketFactory = socketFactory;
 		this.#logger = logger;
@@ -56,21 +57,24 @@ export class ReconnectingWebSocket<TData = unknown>
 			jitterFactor: options.jitterFactor ?? 0.1,
 		};
 		this.#backoffMs = this.#options.initialBackoffMs;
+		this.#onDispose = onDispose;
 	}
 
 	static async create<TData>(
 		socketFactory: SocketFactory<TData>,
 		logger: Logger,
 		apiRoute: string,
 		options: ReconnectingWebSocketOptions = {},
+		onDispose?: () => void,
 	): Promise<ReconnectingWebSocket<TData>> {
 		const instance = new ReconnectingWebSocket<TData>(
 			socketFactory,
 			logger,
 			apiRoute,
 			options,
+			onDispose,
 		);
-		await instance.#connect();
+		await instance.connect();
 		return instance;
 	}
 
@@ -85,10 +89,6 @@ export class ReconnectingWebSocket<TData = unknown>
 		(this.#eventHandlers[event] as Set<EventHandler<TData, TEvent>>).add(
 			callback,
 		);
-
-		if (this.#currentSocket) {
-			this.#currentSocket.addEventListener(event, callback);
-		}
 	}
 
 	removeEventListener<TEvent extends WebSocketEventType>(
@@ -98,95 +98,95 @@ export class ReconnectingWebSocket<TData = unknown>
 		(this.#eventHandlers[event] as Set<EventHandler<TData, TEvent>>).delete(
 			callback,
 		);
-
-		if (this.#currentSocket) {
-			this.#currentSocket.removeEventListener(event, callback);
-		}
 	}
 
-	close(code?: number, reason?: string): void {
+	reconnect(): void {
 		if (this.#isDisposed) {
 			return;
 		}
 
-		this.#isDisposed = true;
-
 		if (this.#reconnectTimeoutId !== null) {
 			clearTimeout(this.#reconnectTimeoutId);
 			this.#reconnectTimeoutId = null;
 		}
 
-		if (this.#currentSocket) {
-			this.#currentSocket.close(code, reason);
-			this.#currentSocket = null;
+		// If already connecting, schedule reconnect after current attempt
+		if (this.#isConnecting) {
+			this.#pendingReconnect = true;
+			return;
 		}
 
-		for (const set of Object.values(this.#eventHandlers)) {
-			set.clear();
-		}
+		// connect() will close any existing socket
+		this.connect().catch((error) => {
+			if (!this.#isDisposed) {
+				this.#logger.warn(
+					`Manual reconnection failed for ${this.#apiRoute}: ${error instanceof Error ? error.message : String(error)}`,
+				);
+				this.scheduleReconnect();
+			}
+		});
 	}
 
-	reconnect(): void {
+	close(code?: number, reason?: string): void {
 		if (this.#isDisposed) {
 			return;
 		}
 
-		if (this.#reconnectTimeoutId !== null) {
-			clearTimeout(this.#reconnectTimeoutId);
-			this.#reconnectTimeoutId = null;
-		}
-
+		// Fire close handlers synchronously before disposing
 		if (this.#currentSocket) {
-			this.#currentSocket.close(CLOSE_CODE_RECONNECTING, "Reconnecting");
+			this.executeHandlers("close", {
+				code: code ?? 1000,
+				reason: reason ?? "",
+				wasClean: true,
+				type: "close",
+				target: this.#currentSocket,
+			} as CloseEvent);
 		}
+
+		this.dispose(code, reason);
 	}
 
-	async #connect(): Promise<void> {
+	private async connect(): Promise<void> {
 		if (this.#isDisposed || this.#isConnecting) {
 			return;
 		}
 
 		this.#isConnecting = true;
 		try {
+			// Close any existing socket before creating a new one
+			if (this.#currentSocket) {
+				this.#currentSocket.close(1000, "Replacing connection");
+				this.#currentSocket = null;
+			}
+
 			const socket = await this.#socketFactory();
 			this.#currentSocket = socket;
 
-			socket.addEventListener("open", () => {
+			socket.addEventListener("open", (event) => {
 				this.#backoffMs = this.#options.initialBackoffMs;
+				this.executeHandlers("open", event);
 			});
 
-			for (const handler of this.#eventHandlers.open) {
-				socket.addEventListener("open", handler);
-			}
-
-			for (const handler of this.#eventHandlers.message) {
-				socket.addEventListener("message", handler);
-			}
+			socket.addEventListener("message", (event) => {
+				this.executeHandlers("message", event);
+			});
 
-			for (const handler of this.#eventHandlers.error) {
-				socket.addEventListener("error", handler);
-			}
+			socket.addEventListener("error", (event) => {
+				this.executeHandlers("error", event);
+			});
 
 			socket.addEventListener("close", (event) => {
-				for (const handler of this.#eventHandlers.close) {
-					handler(event);
-				}
-
 				if (this.#isDisposed) {
 					return;
 				}
 
+				this.executeHandlers("close", event);
+
 				if (UNRECOVERABLE_CLOSE_CODES.has(event.code)) {
 					this.#logger.error(
 						`WebSocket connection closed with unrecoverable error code ${event.code}`,
 					);
-					this.#isDisposed = true;
-					return;
-				}
-
-				// Reconnect if this was an intentional close for reconnection
-				if (event.code === CLOSE_CODE_RECONNECTING) {
-					this.#scheduleReconnect();
+					this.dispose();
 					return;
 				}
 
@@ -196,14 +196,19 @@ export class ReconnectingWebSocket<TData = unknown>
 				}
 
 				// Reconnect on abnormal closures (e.g., 1006) or other unexpected codes
-				this.#scheduleReconnect();
+				this.scheduleReconnect();
 			});
 		} finally {
 			this.#isConnecting = false;
+
+			if (this.#pendingReconnect) {
+				this.#pendingReconnect = false;
+				this.reconnect();
+			}
 		}
 	}
 
-	#scheduleReconnect(): void {
+	private scheduleReconnect(): void {
 		if (this.#isDisposed || this.#reconnectTimeoutId !== null) {
 			return;
 		}
@@ -218,21 +223,58 @@ export class ReconnectingWebSocket<TData = unknown>
 
 		this.#reconnectTimeoutId = setTimeout(() => {
 			this.#reconnectTimeoutId = null;
-			// Errors already handled in #connect
-			this.#connect().catch((error) => {
+			this.connect().catch((error) => {
 				if (!this.#isDisposed) {
 					this.#logger.warn(
 						`WebSocket connection failed for ${this.#apiRoute}: ${error instanceof Error ? error.message : String(error)}`,
 					);
-					this.#scheduleReconnect();
+					this.scheduleReconnect();
 				}
 			});
 		}, delayMs);
 
 		this.#backoffMs = Math.min(this.#backoffMs * 2, this.#options.maxBackoffMs);
 	}
 
-	isDisposed(): boolean {
-		return this.#isDisposed;
+	private executeHandlers<TEvent extends WebSocketEventType>(
+		event: TEvent,
+		eventData: Parameters<EventHandler<TData, TEvent>>[0],
+	): void {
+		const handlers = this.#eventHandlers[event] as Set<
+			EventHandler<TData, TEvent>
+		>;
+		for (const handler of handlers) {
+			try {
+				handler(eventData);
+			} catch (error) {
+				this.#logger.error(
+					`Error in ${event} handler for ${this.#apiRoute}: ${error instanceof Error ? error.message : String(error)}`,
+				);
+			}
+		}
+	}
+
+	private dispose(code?: number, reason?: string): void {
+		if (this.#isDisposed) {
+			return;
+		}
+
+		this.#isDisposed = true;
+
+		if (this.#reconnectTimeoutId !== null) {
+			clearTimeout(this.#reconnectTimeoutId);
+			this.#reconnectTimeoutId = null;
+		}
+
+		if (this.#currentSocket) {
+			this.#currentSocket.close(code, reason);
+			this.#currentSocket = null;
+		}
+
+		for (const set of Object.values(this.#eventHandlers)) {
+			set.clear();
+		}
+
+		this.#onDispose?.();
 	}
 }