@drift-labs/common
Version:
Common functions for Drift
416 lines • 19.9 kB
JavaScript
"use strict";
var __classPrivateFieldGet = (this && this.__classPrivateFieldGet) || function (receiver, state, kind, f) {
if (kind === "a" && !f) throw new TypeError("Private accessor was defined without a getter");
if (typeof state === "function" ? receiver !== state || !f : !state.has(receiver)) throw new TypeError("Cannot read private member from an object whose class did not declare it");
return kind === "m" ? f : kind === "a" ? f.call(receiver) : f ? f.value : state.get(receiver);
};
var __classPrivateFieldSet = (this && this.__classPrivateFieldSet) || function (receiver, state, value, kind, f) {
if (kind === "m") throw new TypeError("Private method is not writable");
if (kind === "a" && !f) throw new TypeError("Private accessor was defined without a setter");
if (typeof state === "function" ? receiver !== state || !f : !state.has(receiver)) throw new TypeError("Cannot write private member to an object whose class did not declare it");
return (kind === "a" ? f.call(receiver, value) : f ? f.value = value : state.set(receiver, value)), value;
};
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
var _MultiplexWebSocket_webSocket;
Object.defineProperty(exports, "__esModule", { value: true });
exports.MultiplexWebSocket = void 0;
const rxjs_1 = require("rxjs");
const isomorphic_ws_1 = __importDefault(require("isomorphic-ws"));
const DEFAULT_MAX_RECONNECT_ATTEMPTS = 5;
const DEFAULT_MAX_RECONNECT_WINDOW_MS = 60 * 1000;
const DEFAULT_CONNECTION_CLOSE_DELAY_MS = 2 * 1000; // 2 seconds delay before closing connection
const DEFAULT_HEARTBEAT_TIMEOUT_MS = 60 * 1000; // Consider connection dead if no heartbeat within 60 seconds
/**
* Manages reconnection logic for WebSocket connections with exponential backoff and rate limiting.
*
* Features:
* - Tracks reconnection attempts within a time window
* - Implements exponential backoff (1s, 2s, 4s, 8s max)
* - Resets attempt counter after configurable time window
* - Throws error when max attempts exceeded
* - Provides configurable limits for attempts and time window
*
* @example
* ```ts
* const manager = new ReconnectionManager(5, 60000); // 5 attempts in 60s
* const { shouldReconnect, delay } = manager.attemptReconnection('ws://example.com');
* if (shouldReconnect) {
* setTimeout(() => reconnect(), delay);
* }
* ```
*/
class ReconnectionManager {
constructor(maxAttemptsCount = DEFAULT_MAX_RECONNECT_ATTEMPTS, maxAttemptsWindowMs = DEFAULT_MAX_RECONNECT_WINDOW_MS) {
this.reconnectAttempts = 0;
this.lastReconnectWindow = Date.now();
this.maxAttemptsCount = maxAttemptsCount;
this.maxAttemptsWindowMs = maxAttemptsWindowMs;
}
attemptReconnection(wsUrl) {
const now = Date.now();
// Reset reconnect attempts if more than a minute has passed
if (now - this.lastReconnectWindow > this.maxAttemptsWindowMs) {
this.reconnectAttempts = 0;
this.lastReconnectWindow = now;
}
// Check if we've exceeded the maximum reconnect attempts
if (this.reconnectAttempts >= this.maxAttemptsCount) {
throw new Error(`WebSocket reconnection failed: Maximum reconnect attempts (${this.maxAttemptsCount}) exceeded within ${this.maxAttemptsWindowMs}ms for ${wsUrl}`);
}
this.reconnectAttempts++;
// Calculate exponential backoff delay: 1s, 2s, 4s, 8s, etc. (capped at 8s)
const backoffDelay = Math.min(1000 * Math.pow(2, this.reconnectAttempts - 1), 8000);
return {
shouldReconnect: true,
delay: backoffDelay,
};
}
reset() {
this.reconnectAttempts = 0;
this.lastReconnectWindow = Date.now();
}
}
var WebSocketConnectionState;
(function (WebSocketConnectionState) {
WebSocketConnectionState[WebSocketConnectionState["CONNECTING"] = 0] = "CONNECTING";
WebSocketConnectionState[WebSocketConnectionState["CONNECTED"] = 1] = "CONNECTED";
WebSocketConnectionState[WebSocketConnectionState["DISCONNECTING"] = 2] = "DISCONNECTING";
WebSocketConnectionState[WebSocketConnectionState["DISCONNECTED"] = 3] = "DISCONNECTED";
})(WebSocketConnectionState || (WebSocketConnectionState = {}));
/**
* MultiplexWebSocket allows for multiple subscriptions to a single websocket of the same URL, improving efficiency and reducing the number of open connections.
*
* This implementation assumes the following:
* - All websocket streams are treated equally - reconnection attempts are performed at the same standards
* - All messages returned are in the `WebSocketMessage` format
* - A single instance of the websocket manager is created for each websocket URL - this means all subscriptions to the same websocket URL will share the same websocket instance
*
* Internal implementation details:
* - The websocket is closed when the number of subscriptions is 0
* - The websocket will be refreshed (new instance) when it disconnects unexpectedly or errors, until it reaches the maximum number of reconnect attempts
*/
class MultiplexWebSocket {
constructor(wsUrl, enableHeartbeatMonitoring = false) {
_MultiplexWebSocket_webSocket.set(this, void 0);
this.closeTimeout = null;
this.heartbeatTimeout = null;
this.heartbeatMonitoringEnabled = false;
if (MultiplexWebSocket.URL_TO_WEBSOCKETS_LOOKUP.has(wsUrl)) {
throw new Error(`Attempting to create a new websocket for ${wsUrl}, but it already exists`);
}
this.wsUrl = wsUrl;
this.customConnectionState = WebSocketConnectionState.CONNECTING;
this.subject = new rxjs_1.Subject();
this.subscriptions = new Map();
this.reconnectionManager = new ReconnectionManager();
this.heartbeatMonitoringEnabled = enableHeartbeatMonitoring;
this.webSocket = new isomorphic_ws_1.default(wsUrl);
MultiplexWebSocket.URL_TO_WEBSOCKETS_LOOKUP.set(wsUrl, this);
}
/**
* Creates a new virtual websocket subscription. If an existing websocket for the given URL exists, the subscription will be added to the existing websocket.
* Returns a function that can be called to unsubscribe from the subscription.
*/
static createWebSocketSubscription(props) {
const { wsUrl } = props;
const doesWebSocketForWsUrlExist = MultiplexWebSocket.URL_TO_WEBSOCKETS_LOOKUP.has(wsUrl);
if (doesWebSocketForWsUrlExist) {
return this.handleNewSubForExistingWsUrl(props);
}
else {
// Create new websocket for new URL or reopen previously closed websocket
return this.handleNewSubForNewWsUrl(props);
}
}
static handleNewSubForNewWsUrl(newSubscriptionProps) {
var _a;
const newMWS = new MultiplexWebSocket(newSubscriptionProps.wsUrl, (_a = newSubscriptionProps.enableHeartbeatMonitoring) !== null && _a !== void 0 ? _a : false);
newMWS.subscribe(newSubscriptionProps);
return {
unsubscribe: () => {
newMWS.unsubscribe(newSubscriptionProps.subscriptionId);
},
};
}
static handleNewSubForExistingWsUrl(newSubscriptionProps) {
var _a;
const { wsUrl, subscriptionId } = newSubscriptionProps;
if (!MultiplexWebSocket.URL_TO_WEBSOCKETS_LOOKUP.has(wsUrl)) {
throw new Error(`Attempting to subscribe to ${subscriptionId} on websocket ${wsUrl}, but websocket does not exist yet`);
}
const existingMWS = MultiplexWebSocket.URL_TO_WEBSOCKETS_LOOKUP.get(wsUrl);
// Check if heartbeat monitoring settings match
const requestedHeartbeat = (_a = newSubscriptionProps.enableHeartbeatMonitoring) !== null && _a !== void 0 ? _a : false;
if (existingMWS.heartbeatMonitoringEnabled !== requestedHeartbeat) {
console.warn(`WebSocket for ${wsUrl} already exists with heartbeat monitoring ${existingMWS.heartbeatMonitoringEnabled ? 'enabled' : 'disabled'}, ` +
`but new subscription requests ${requestedHeartbeat ? 'enabled' : 'disabled'}. Using existing setting.`);
}
// Track new subscription for existing websocket
existingMWS.subscribe(newSubscriptionProps);
return {
unsubscribe: () => {
existingMWS.unsubscribe(newSubscriptionProps.subscriptionId);
},
};
}
get webSocket() {
return __classPrivateFieldGet(this, _MultiplexWebSocket_webSocket, "f");
}
/**
* Setting the WebSocket instance will automatically add event handlers to the WebSocket instance.
* When the WebSocket is connected, all existing subscriptions will be subscribed to.
*/
set webSocket(webSocket) {
webSocket.onopen = () => {
this.customConnectionState = WebSocketConnectionState.CONNECTED;
this.reconnectionManager.reset(); // Reset reconnection attempts on successful connection
// Start heartbeat monitoring if enabled
if (this.heartbeatMonitoringEnabled) {
this.startHeartbeatMonitoring();
}
// sends subscription message for each subscription for those that are added before the websocket is connected
for (const [subscriptionId] of this.subscriptions.entries()) {
this.subscribeToWebSocket(subscriptionId);
}
};
webSocket.onmessage = (messageEvent) => {
const message = JSON.parse(messageEvent.data);
// Check for heartbeat message from server (only if heartbeat monitoring is enabled)
if (this.heartbeatMonitoringEnabled && this.isHeartbeatMessage(message)) {
this.handleHeartbeat(message);
return; // Don't forward heartbeat messages to subscriptions
}
this.subject.next(message);
};
webSocket.onclose = (event) => {
this.customConnectionState = WebSocketConnectionState.DISCONNECTED;
// Stop heartbeat monitoring when connection closes (if enabled)
if (this.heartbeatMonitoringEnabled) {
this.stopHeartbeatMonitoring();
}
// Restart websocket if it was closed unexpectedly (not by us)
if (!event.wasClean && this.subscriptions.size > 0) {
console.log('WebSocket closed unexpectedly, restarting...', event);
this.refreshWebSocket();
}
};
webSocket.onerror = (error) => {
console.error('MultiplexWebSocket Error', { error, webSocket });
// Forward error to all subscriptions for this websocket URL
const subscriptionIds = MultiplexWebSocket.URL_TO_SUBSCRIPTION_IDS_LOOKUP.get(this.wsUrl);
if (subscriptionIds) {
for (const subscriptionId of subscriptionIds) {
const subscription = this.subscriptions.get(subscriptionId);
if (subscription) {
subscription.onError(error);
}
}
}
// Restart the websocket connection on error
this.refreshWebSocket();
};
__classPrivateFieldSet(this, _MultiplexWebSocket_webSocket, webSocket, "f");
}
subscribeToWebSocket(subscriptionId) {
const subscriptionState = this.subscriptions.get(subscriptionId);
const { subscribeMessage, onError, onMessage, messageFilter, errorMessageFilter, onClose, } = subscriptionState;
this.webSocket.send(subscribeMessage);
if (subscriptionState) {
subscriptionState.hasSentSubscribeMessage = true;
}
const subjectSubscription = this.subject
.pipe((0, rxjs_1.catchError)((err) => {
console.error('Caught websocket error', err);
onError();
return [];
}))
.subscribe({
next: (message) => {
try {
if (!messageFilter(message))
return;
if (errorMessageFilter(message)) {
onError();
return;
}
onMessage(message);
}
catch (err) {
console.error('Error parsing websocket message', err);
onError();
}
},
error: (err) => {
console.error('Error subscribing to websocket', err);
onError();
},
complete: () => {
if (onClose) {
onClose();
}
},
});
subscriptionState.subjectSubscription = subjectSubscription;
}
subscribe(props) {
const { subscriptionId, subscribeMessage, unsubscribeMessage, onError, onMessage, messageFilter, errorMessageFilter, onClose, } = props;
if (this.subscriptions.get(subscriptionId)) {
throw new Error(`Attempting to subscribe to ${subscriptionId} on websocket ${this.wsUrl}, but subscription already exists`);
}
// Cancel any pending delayed close since we're adding a new subscription
this.cancelDelayedClose();
this.subscriptions.set(subscriptionId, {
subscribeMessage,
unsubscribeMessage,
onError,
onMessage,
messageFilter,
errorMessageFilter,
onClose,
hasSentSubscribeMessage: false,
});
// Update URL to subscription IDs lookup
if (!MultiplexWebSocket.URL_TO_SUBSCRIPTION_IDS_LOOKUP.has(this.wsUrl)) {
MultiplexWebSocket.URL_TO_SUBSCRIPTION_IDS_LOOKUP.set(this.wsUrl, new Set());
}
MultiplexWebSocket.URL_TO_SUBSCRIPTION_IDS_LOOKUP.get(this.wsUrl).add(subscriptionId);
if (this.customConnectionState === WebSocketConnectionState.CONNECTED) {
this.subscribeToWebSocket(subscriptionId);
}
else if (this.customConnectionState === WebSocketConnectionState.CONNECTING) {
// do nothing, subscription will automatically start when websocket is connected
}
else {
// handle case where websocket is disconnecting/disconnected
this.refreshWebSocket();
}
}
unsubscribe(subscriptionId) {
var _a;
const subscriptionState = this.subscriptions.get(subscriptionId);
if (subscriptionState) {
(_a = subscriptionState.subjectSubscription) === null || _a === void 0 ? void 0 : _a.unsubscribe();
// Only send unsubscribe message if websocket is connected and ready to send.
//// Otherwise, when the websocket DOES connect we don't have to worry about this subscription because we are deleting it from the subscriptions map. (Which only trigger their connections once the websocket becomes connected)
if (this.customConnectionState === WebSocketConnectionState.CONNECTED &&
this.webSocket.readyState === isomorphic_ws_1.default.OPEN) {
this.webSocket.send(subscriptionState.unsubscribeMessage);
}
this.subscriptions.delete(subscriptionId);
// Update URL to subscription IDs lookup
const subscriptionIds = MultiplexWebSocket.URL_TO_SUBSCRIPTION_IDS_LOOKUP.get(this.wsUrl);
if (subscriptionIds) {
subscriptionIds.delete(subscriptionId);
if (subscriptionIds.size === 0) {
MultiplexWebSocket.URL_TO_SUBSCRIPTION_IDS_LOOKUP.delete(this.wsUrl);
// Schedule delayed close when last subscriber unsubscribes
this.scheduleDelayedClose();
}
}
}
}
scheduleDelayedClose() {
// Cancel any existing delayed close timeout
this.cancelDelayedClose();
// Schedule new delayed close
this.closeTimeout = setTimeout(() => {
this.close();
}, DEFAULT_CONNECTION_CLOSE_DELAY_MS);
}
cancelDelayedClose() {
if (this.closeTimeout) {
clearTimeout(this.closeTimeout);
this.closeTimeout = null;
}
}
close() {
// Cancel any pending delayed close
this.cancelDelayedClose();
// Stop heartbeat monitoring (if enabled)
if (this.heartbeatMonitoringEnabled) {
this.stopHeartbeatMonitoring();
}
for (const [subscriptionId] of this.subscriptions.entries()) {
this.unsubscribe(subscriptionId);
}
this.subscriptions.clear();
this.webSocket.close();
MultiplexWebSocket.URL_TO_WEBSOCKETS_LOOKUP.delete(this.wsUrl);
MultiplexWebSocket.URL_TO_SUBSCRIPTION_IDS_LOOKUP.delete(this.wsUrl);
this.reconnectionManager.reset();
}
startHeartbeatMonitoring() {
// Start the heartbeat timeout - if we don't receive a heartbeat message within the timeout, refresh connection
this.resetHeartbeatTimeout();
}
stopHeartbeatMonitoring() {
if (this.heartbeatTimeout) {
clearTimeout(this.heartbeatTimeout);
this.heartbeatTimeout = null;
}
}
resetHeartbeatTimeout() {
// Clear existing timeout
if (this.heartbeatTimeout) {
clearTimeout(this.heartbeatTimeout);
}
// Set new timeout
this.heartbeatTimeout = setTimeout(() => {
console.warn(`No heartbeat received within ${DEFAULT_HEARTBEAT_TIMEOUT_MS}ms - connection appears dead`);
this.refreshWebSocket();
}, DEFAULT_HEARTBEAT_TIMEOUT_MS);
}
isHeartbeatMessage(message) {
// Check if message is a heartbeat message from server
return (message === null || message === void 0 ? void 0 : message.channel) === 'heartbeat';
}
handleHeartbeat(_message) {
// Reset the heartbeat timeout
this.resetHeartbeatTimeout();
}
refreshWebSocket() {
// Cancel any pending delayed close since we're refreshing
this.cancelDelayedClose();
// Reset heartbeat monitoring during refresh (if enabled) - it will restart when the new connection opens
if (this.heartbeatMonitoringEnabled) {
this.stopHeartbeatMonitoring();
}
const { shouldReconnect, delay } = this.reconnectionManager.attemptReconnection(this.wsUrl);
if (!shouldReconnect) {
return;
}
// Clean up current websocket
const currentWebSocket = this.webSocket;
if (currentWebSocket) {
currentWebSocket.onerror = () => { };
currentWebSocket.onclose = () => { };
currentWebSocket.onmessage = () => { };
currentWebSocket.onopen = () => {
// in the event where the websocket has yet to connect, we close the connection after it is connected
currentWebSocket.close();
};
currentWebSocket.close();
}
// Reset subscription states
this.subscriptions.forEach((subscription) => {
subscription.hasSentSubscribeMessage = false;
});
// Use exponential backoff before attempting to reconnect
setTimeout(() => {
this.webSocket = new isomorphic_ws_1.default(this.wsUrl);
}, delay);
}
}
exports.MultiplexWebSocket = MultiplexWebSocket;
_MultiplexWebSocket_webSocket = new WeakMap();
/**
* A lookup of all websockets by their URL.
*/
MultiplexWebSocket.URL_TO_WEBSOCKETS_LOOKUP = new Map();
/**
* A lookup from websocket URL to all subscription IDs for that URL.
*/
MultiplexWebSocket.URL_TO_SUBSCRIPTION_IDS_LOOKUP = new Map();
//# sourceMappingURL=MultiplexWebSocket.js.map