silberengel
/
gc-alexandria
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity
Browse Source

Introduce WebSocket handles and simplify 

- The `acquire`, `release`, and `drain` method signatures remain unchanged.
- `WebSocketHandle` objects are introduced to tie each socket's state to the socket itself.
- Initial connection logic is simplified.
- Some minor improvements and validations are introduced based on AI-assisted code review.
master
buttercat1791 8 months ago
parent
fee33e8f92
commit
2a57122436
1 changed files with 152 additions and 124 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Unified View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 274
      src/lib/data_structures/websocket_pool.ts

274
src/lib/data_structures/websocket_pool.ts
Unescape View File

@ -1,6 +1,16 @@
/**
* A representation of a WebSocket connection used internally by the WebSocketPool. Attaches
* information about the connection's state within the resource pool to the connection itself.
*/
interface WebSocketHandle {
ws: WebSocket;
refCount: number;
idleTimer?: ReturnType<typeof setTimeout>;
}
interface WebSocketPoolWaitingQueueItem { interface WebSocketPoolWaitingQueueItem {
url: string; url: string;
resolve: (ws: WebSocket) => void; resolve: (handle: WebSocketHandle) => void;
reject: (reason?: any) => void; reject: (reason?: any) => void;
} }
@ -16,10 +26,11 @@ interface WebSocketPoolWaitingQueueItem {
export class WebSocketPool { export class WebSocketPool {
static #shared: WebSocketPool; static #shared: WebSocketPool;
#pool: Map<string, WebSocket> = new Map(); /**
#connecting: Map<string, Promise<WebSocket>> = new Map(); * A map of WebSocket URLs to the handles containing the connection and its state.
#refCounts: Map<WebSocket, number> = new Map(); */
#idleTimers: Map<WebSocket, ReturnType<typeof setTimeout>> = new Map(); #pool: Map<string, WebSocketHandle> = new Map();
#idleTimeoutMs: number; #idleTimeoutMs: number;
#maxConnections: number; #maxConnections: number;
#waitingQueue: WebSocketPoolWaitingQueueItem[] = []; #waitingQueue: WebSocketPoolWaitingQueueItem[] = [];
@ -50,71 +61,38 @@ export class WebSocketPool {
// #region Resource Management Interface // #region Resource Management Interface
/** /**
* Acquires a WebSocket connection for the specified URL. If a connection is available in the * Sets the maximum number of simultaneous WebSocket connections.
* pool, that connection is returned. If no connection is available but the pool is not full, a
* new connection is created and returned. If the pool is full, the request is queued until a
* connection is released, at which point the newly-available connection is returned to the
* caller.
* *
* @param url - The URL to connect to. * @param limit - The new connection limit.
* @returns A promise that resolves with a WebSocket connection.
*/ */
public async acquire(url: string): Promise<WebSocket> { public set maxConnections(limit: number) {
const normalizedUrl = this.#normalizeUrl(url); if (limit === this.#maxConnections) {
const existingSocket = this.#pool.get(normalizedUrl); return;
if (existingSocket && existingSocket.readyState < WebSocket.CLOSING) {
this.#checkOutSocket(existingSocket);
return Promise.resolve(existingSocket);
} }
const connectingPromise = this.#connecting.get(normalizedUrl); if (limit == null || isNaN(limit)) {
if (connectingPromise) { throw new Error('[WebSocketPool] Connection limit must be a number.');
const ws = await connectingPromise;
if (ws.readyState === WebSocket.OPEN) {
this.#checkOutSocket(ws);
return ws;
}
throw new Error(`[WebSocketPool] WebSocket connection failed for ${normalizedUrl}`);
} }
if (this.#pool.size + this.#connecting.size >= this.#maxConnections) { if (limit <= 0) {
return new Promise((resolve, reject) => { throw new Error('[WebSocketPool] Connection limit must be greater than 0.');
this.#waitingQueue.push({ url: normalizedUrl, resolve, reject });
});
} }
const newConnectionPromise = this.#createSocket(normalizedUrl); if (!Number.isInteger(limit)) {
this.#connecting.set(normalizedUrl, newConnectionPromise); throw new Error('[WebSocketPool] Connection limit must be an integer.');
}
newConnectionPromise.finally(() => {
this.#connecting.delete(normalizedUrl);
});
return newConnectionPromise; this.#maxConnections = limit;
this.#processWaitingQueue();
} }
/** /**
* Releases a WebSocket connection back to the pool. If there are pending requests for the same * Gets the current maximum number of simultaneous WebSocket connections.
* URL, the connection is passed to the requestor in the queue. Otherwise, the connection is
* marked as available.
* *
* @param ws - The WebSocket connection to release. * @returns The current connection limit.
*/ */
public release(ws: WebSocket): void { public get maxConnections(): number {
const currentCount = this.#refCounts.get(ws); return this.#maxConnections;
if (!currentCount) {
throw new Error('[WebSocketPool] Attempted to release an unmanaged WebSocket connection.');
}
if (currentCount > 0) {
const newCount = currentCount - 1;
this.#refCounts.set(ws, newCount);
if (newCount === 0) {
this.#startIdleTimer(ws);
}
}
} }
/** /**
@ -122,7 +100,23 @@ export class WebSocketPool {
* *
* @param timeoutMs - The timeout in milliseconds after which idle connections will be closed. * @param timeoutMs - The timeout in milliseconds after which idle connections will be closed.
*/ */
public setIdleTimeout(timeoutMs: number): void { public set idleTimeoutMs(timeoutMs: number) {
if (timeoutMs === this.#idleTimeoutMs) {
return;
}
if (timeoutMs == null || isNaN(timeoutMs)) {
throw new Error('[WebSocketPool] Idle timeout must be a number.');
}
if (timeoutMs <= 0) {
throw new Error('[WebSocketPool] Idle timeout must be greater than 0.');
}
if (!Number.isInteger(timeoutMs)) {
throw new Error('[WebSocketPool] Idle timeout must be an integer.');
}
this.#idleTimeoutMs = timeoutMs; this.#idleTimeoutMs = timeoutMs;
} }
@ -131,27 +125,66 @@ export class WebSocketPool {
* *
* @returns The current idle timeout in milliseconds. * @returns The current idle timeout in milliseconds.
*/ */
public getIdleTimeout(): number { public get idleTimeoutMs(): number {
return this.#idleTimeoutMs; return this.#idleTimeoutMs;
} }
/** /**
* Sets the maximum number of simultaneous WebSocket connections. * Acquires a WebSocket connection for the specified URL. If a connection is available in the
* pool, that connection is returned. If no connection is available but the pool is not full, a
* new connection is created and returned. If the pool is full, the request is queued until a
* connection is released, at which point the newly-available connection is returned to the
* caller.
* *
* @param limit - The new connection limit. * @param url - The URL to connect to.
* @returns A promise that resolves with a WebSocket connection.
*/ */
public set maxConnections(limit: number) { public async acquire(url: string): Promise<WebSocket> {
this.#maxConnections = limit; const normalizedUrl = this.#normalizeUrl(url);
this.#processWaitingQueue(); const handle = this.#pool.get(normalizedUrl);
try {
if (handle && handle.ws.readyState === WebSocket.OPEN) {
this.#checkOut(handle);
return handle.ws;
}
if (this.#pool.size >= this.#maxConnections) {
return new Promise((resolve, reject) => {
this.#waitingQueue.push({
url: normalizedUrl,
resolve: (handle) => resolve(handle.ws),
reject,
});
});
}
const newHandle = await this.#createSocket(normalizedUrl);
return newHandle.ws;
} catch (error) {
throw new Error(
`[WebSocketPool] Failed to acquire connection for ${normalizedUrl}: ${error}`
);
}
} }
/** /**
* Gets the current maximum number of simultaneous WebSocket connections. * Releases a WebSocket connection back to the pool. If there are pending requests for the same
* URL, the connection is passed to the requestor in the queue. Otherwise, the connection is
* marked as available.
* *
* @returns The current connection limit. * @param handle - The WebSocketHandle to release.
*/ */
public get maxConnections(): number { public release(ws: WebSocket): void {
return this.#maxConnections; const normalizedUrl = this.#normalizeUrl(ws.url);
const handle = this.#pool.get(normalizedUrl);
if (!handle) {
throw new Error('[WebSocketPool] Attempted to release an unmanaged WebSocket connection.');
}
if (--handle.refCount === 0) {
this.#startIdleTimer(handle);
}
} }
/** /**
@ -159,52 +192,47 @@ export class WebSocketPool {
*/ */
public drain(): void { public drain(): void {
// Clear all idle timers first // Clear all idle timers first
for (const timer of this.#idleTimers.values()) { for (const handle of this.#pool.values()) {
clearTimeout(timer); this.#clearIdleTimer(handle);
} }
this.#idleTimers.clear();
for (const { reject } of this.#waitingQueue) { for (const { reject } of this.#waitingQueue) {
reject(new Error('[WebSocketPool] Draining pool.')); reject(new Error('[WebSocketPool] Draining pool.'));
} }
this.#waitingQueue = []; this.#waitingQueue = [];
for (const promise of this.#connecting.values()) { for (const handle of this.#pool.values()) {
// While we can't cancel the connection attempt, we can prevent callers from using it. handle.ws.close();
promise.catch(() => {
/* ignore connection errors during drain */
});
}
this.#connecting.clear();
for (const ws of this.#pool.values()) {
ws.close();
} }
this.#pool.clear(); this.#pool.clear();
this.#refCounts.clear();
} }
// #endregion // #endregion
// #region Private Helper Methods // #region Private Helper Methods
#createSocket(url: string): Promise<WebSocket> { #createSocket(url: string): Promise<WebSocketHandle> {
return new Promise((resolve, reject) => { return new Promise((resolve, reject) => {
try { try {
const ws = new WebSocket(url); const handle: WebSocketHandle = {
ws.onopen = () => { ws: new WebSocket(url),
this.#pool.set(url, ws); refCount: 1,
this.#refCounts.set(ws, 1); };
handle.ws.onopen = () => {
this.#pool.set(url, handle);
// Remove the socket from the pool when it is closed. The socket may be closed by // Remove the socket from the pool when it is closed. The socket may be closed by
// either the client or the server. // either the client or the server.
ws.onclose = () => this.#removeSocket(ws); handle.ws.onclose = () => this.#removeSocket(handle);
resolve(ws); resolve(handle);
}; };
ws.onerror = (event) => { handle.ws.onerror = (event) => {
this.#removeSocket(handle);
this.#processWaitingQueue(); this.#processWaitingQueue();
reject(new Error(`[WebSocketPool] WebSocket connection failed for ${url}: ${event.type}`)); reject(
new Error(`[WebSocketPool] WebSocket connection failed for ${url}: ${event.type}`)
);
}; };
} catch (error) { } catch (error) {
this.#processWaitingQueue(); this.#processWaitingQueue();
@ -213,11 +241,10 @@ export class WebSocketPool {
}); });
} }
#removeSocket(ws: WebSocket): void { #removeSocket(handle: WebSocketHandle): void {
const url = ws.url; this.#clearIdleTimer(handle);
this.#pool.delete(url); handle.ws.onopen = handle.ws.onerror = handle.ws.onclose = null;
this.#refCounts.delete(ws); this.#pool.delete(this.#normalizeUrl(handle.ws.url));
this.#clearIdleTimer(ws);
this.#processWaitingQueue(); this.#processWaitingQueue();
} }
@ -227,64 +254,65 @@ export class WebSocketPool {
* *
* @param ws - The WebSocket for which to start the idle timer. * @param ws - The WebSocket for which to start the idle timer.
*/ */
#startIdleTimer(ws: WebSocket): void { #startIdleTimer(handle: WebSocketHandle): void {
// Clear any existing timer first // Clear any existing timer first
this.#clearIdleTimer(ws); this.#clearIdleTimer(handle);
const timer = setTimeout(() => { handle.idleTimer = setTimeout(() => {
const refCount = this.#refCounts.get(ws); const refCount = handle.refCount;
if ((!refCount || refCount === 0) && ws.readyState === WebSocket.OPEN) { if (refCount === 0 && handle.ws.readyState === WebSocket.OPEN) {
ws.close(); handle.ws.close();
this.#removeSocket(ws); this.#removeSocket(handle);
} }
}, this.#idleTimeoutMs); }, this.#idleTimeoutMs);
this.#idleTimers.set(ws, timer);
} }
/** /**
* Clears the idle timer for the specified WebSocket. * Clears the idle timer for the specified WebSocket.
* *
* @param ws - The WebSocket for which to clear the idle timer. * @param handle - The WebSocketHandle for which to clear the idle timer.
*/ */
#clearIdleTimer(ws: WebSocket): void { #clearIdleTimer(handle: WebSocketHandle): void {
const timer = this.#idleTimers.get(ws); const timer = handle.idleTimer;
if (timer) { if (timer) {
clearTimeout(timer); clearTimeout(timer);
this.#idleTimers.delete(ws); handle.idleTimer = undefined;
} }
} }
/**
* Processes pending requests to acquire a connection. Reuses existing connections when possible.
*/
#processWaitingQueue(): void { #processWaitingQueue(): void {
while ( while (
this.#waitingQueue.length > 0 && this.#waitingQueue.length > 0 &&
this.#pool.size + this.#connecting.size < this.#maxConnections this.#pool.size < this.#maxConnections
) { ) {
const nextInQueue = this.#waitingQueue.shift(); const nextInQueue = this.#waitingQueue.shift();
if (!nextInQueue) { if (!nextInQueue) {
continue; continue;
} }
const { url, resolve, reject } = nextInQueue; const { url, resolve, reject } = nextInQueue;
// Re-check if a connection for this URL was created while this request was in the queue
const existingSocket = this.#pool.get(url); const existingHandle = this.#pool.get(url);
if (existingSocket && existingSocket.readyState < WebSocket.CLOSING) { if (existingHandle && existingHandle.ws.readyState === WebSocket.OPEN) {
this.#checkOutSocket(existingSocket); this.#checkOut(existingHandle);
resolve(existingSocket); resolve(existingHandle);
} else { return;
const connectingPromise = this.#connecting.get(url);
if (connectingPromise) {
connectingPromise.then(resolve, reject);
} }
this.#createSocket(url).then(resolve, reject);
} }
} }
#checkOut(handle: WebSocketHandle): void {
if (handle.refCount == null) {
throw new Error('[WebSocketPool] Handle refCount unexpectedly null.');
} }
#checkOutSocket(ws: WebSocket): void { ++handle.refCount;
const count = (this.#refCounts.get(ws) || 0) + 1; this.#clearIdleTimer(handle);
this.#refCounts.set(ws, count);
this.#clearIdleTimer(ws);
} }
#normalizeUrl(url: string): string { #normalizeUrl(url: string): string {

Write Preview
Loading…
Cancel
Save