This commit is contained in:
Andras Schmelczer 2026-09-27 08:29:17 +01:00
commit 3dab48f58b

View file

@ -4,8 +4,7 @@ import type { WebSocketServerMessage } from "./types/WebSocketServerMessage";
import type { WebSocketClientMessage } from "./types/WebSocketClientMessage";
import type { CursorPositionFromClient } from "./types/CursorPositionFromClient";
import type { ClientCursors } from "./types/ClientCursors";
import { createPromise } from "../utils/create-promise";
import { WEBSOCKET_DISCONNECT_TIMEOUT_IN_S, WEBSOCKET_RECEIVE_TIMEOUT_MS } from "../consts";
import { WEBSOCKET_RECEIVE_TIMEOUT_MS } from "../consts";
import { removeFromArray } from "../utils/remove-from-array";
import { EventListeners } from "../utils/data-structures/event-listeners";
import { awaitAll } from "../utils/await-all";
@ -24,7 +23,6 @@ export class WebSocketManager {
>();
private isStopped = true;
private resolveDisconnectingPromise: null | (() => unknown) = null;
private receiveTimeoutId: ReturnType<typeof setTimeout> | undefined;
private reconnectTimeoutId: ReturnType<typeof setTimeout> | undefined;
@ -69,58 +67,10 @@ export class WebSocketManager {
}
public async stop(): Promise<void> {
const [promise, resolve] = createPromise();
this.resolveDisconnectingPromise = resolve;
this.isStopped = true;
clearTimeout(this.receiveTimeoutId);
if (this.reconnectTimeoutId !== undefined) {
clearTimeout(this.reconnectTimeoutId);
this.reconnectTimeoutId = undefined;
}
const socket = this.webSocket;
try {
socket?.close(1000, "WebSocketManager has been stopped");
} catch (error) {
this.logger.warn(`Failed to close WebSocket: ${String(error)}`);
if (socket) this.reconnect(socket);
}
// eslint-disable-next-line @typescript-eslint/init-declarations
let timeoutId: ReturnType<typeof setTimeout> | undefined;
const timeoutPromise = new Promise<void>((_, reject) => {
timeoutId = setTimeout(() => {
reject(
new Error(
`Timeout waiting for WebSocket to close after ${WEBSOCKET_DISCONNECT_TIMEOUT_IN_S} seconds`
)
);
}, WEBSOCKET_DISCONNECT_TIMEOUT_IN_S * 1000);
});
try {
while (this.isWebSocketConnected) {
await Promise.race([promise, timeoutPromise]);
}
} catch (error) {
this.logger.error(
`Error while waiting for WebSocket to close: ${String(error)}`
);
// Force cleanup even if close didn't work
this.resolveDisconnectingPromise();
this.resolveDisconnectingPromise = null;
} finally {
// Clear timeout to prevent unhandled rejection
if (timeoutId !== undefined) {
clearTimeout(timeoutId);
}
// A throwing, silent or timed-out close must still fence callbacks
// before a stopped client can release its state or restart.
if (socket) this.reconnect(socket);
}
clearTimeout(this.reconnectTimeoutId);
this.reconnectTimeoutId = undefined;
this.retireSocket();
await this.waitUntilFinished();
}
@ -128,27 +78,6 @@ export class WebSocketManager {
await awaitAll(this.outstandingPromises);
}
public sendHandshakeMessage(
message: WebSocketClientMessage & { type: "handshake" }
): void {
const { webSocket } = this;
if (!webSocket) {
throw new Error(
"WebSocket is not connected, cannot send handshake message"
);
}
try {
webSocket.send(JSON.stringify(message));
} catch (error) {
this.logger.error(
`Failed to send handshake message: ${String(error)}`
);
this.reconnect(webSocket);
throw error;
}
}
public updateLocalCursors(cursorPositions: CursorPositionFromClient): void {
if (!this.isWebSocketConnected || !this.webSocket) {
// A missing cursor update is fine, we can just skip it if needed
@ -190,30 +119,31 @@ export class WebSocketManager {
this.receiveTimeoutId.unref?.();
}
private reconnect(socket: WebSocket): void {
if (this.webSocket !== socket) return;
private retireSocket(): void {
clearTimeout(this.receiveTimeoutId);
const socket = this.webSocket;
if (!socket) return;
this.webSocket = undefined;
socket.onopen = null;
socket.onclose = null;
socket.onmessage = null;
socket.onerror = (): void => {
/* Retired transports may still emit errors. */
};
this.webSocket = undefined;
this.onWebSocketStatusChanged.trigger(false);
try {
socket.close();
} catch {
/* The transport may already be gone. */
}
if (this.isStopped) {
this.resolveDisconnectingPromise?.();
this.resolveDisconnectingPromise = null;
} else {
this.scheduleReconnect();
socket.close(1000, "WebSocket retired");
} catch (error) {
this.logger.warn(`Failed to close WebSocket: ${String(error)}`);
}
}
private reconnect(socket: WebSocket): void {
if (this.webSocket !== socket) return;
this.retireSocket();
this.scheduleReconnect();
}
private scheduleReconnect(): void {
if (this.isStopped) return;
clearTimeout(this.reconnectTimeoutId);
@ -237,23 +167,7 @@ export class WebSocketManager {
}
private openWebSocket(): void {
// Clean up old WebSocket handlers to prevent race conditions
if (this.webSocket) {
try {
// Remove handlers to prevent them from firing after new connection
this.webSocket.onopen = null;
this.webSocket.onclose = null;
this.webSocket.onmessage = null;
this.webSocket.onerror = (): void => {
/* Retired transports may still emit errors. */
};
this.webSocket.close();
} catch (e) {
this.logger.error(
`Failed to close previous WebSocket connection: ${e}`
);
}
}
this.retireSocket();
const { remoteUri, vaultName } = this.settings.getSettings();
const wsUri = new URL(remoteUri);
@ -274,19 +188,17 @@ export class WebSocketManager {
};
this.webSocket.onopen = (): void => {
// Check if we've been stopped while connecting
if (this.isStopped) {
this.webSocket?.close(
1000,
"WebSocketManager was stopped during connection"
);
return;
}
this.logger.info("WebSocket connection opened");
if (this.webSocket !== socket || this.isStopped) return;
try {
const handshake: WebSocketClientMessage = {
type: "handshake",
token: this.settings.getSettings().token,
deviceId: this.deviceId
};
socket.send(JSON.stringify(handshake));
this.logger.info("WebSocket connection opened");
this.onWebSocketStatusChanged.trigger(true);
// A handshake failure can retire this socket from inside an
// observer. Finish delivery with its actual connection state.
// An observer may stop or replace the connection during delivery.
if (this.webSocket !== socket)
this.onWebSocketStatusChanged.trigger(false);
} catch (error) {
@ -298,6 +210,7 @@ export class WebSocketManager {
};
this.webSocket.onmessage = (event): void => {
if (this.webSocket !== socket || this.isStopped) return;
try {
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
const message = JSON.parse(
@ -322,7 +235,7 @@ export class WebSocketManager {
);
});
void this.outstandingPromises.push(messageHandlingPromise); // ignore the returned promise
this.outstandingPromises.push(messageHandlingPromise);
} catch (error) {
this.logger.error(
`Error parsing WebSocket message: ${String(error)}`