refactor(WebSocketShard): throttling error handling (#9701)

* refactor(WebSocketShard): handle unknown identify errors

* chore: use better abort check

---------

Co-authored-by: kodiakhq[bot] <49736102+kodiakhq[bot]@users.noreply.github.com>

packages/ws/src/strategies/context/IContextFetchingStrategy.ts

@@ -25,7 +25,9 @@ export interface IContextFetchingStrategy {
 	retrieveSessionInfo(shardId: number): Awaitable<SessionInfo | null>;
 	updateSessionInfo(shardId: number, sessionInfo: SessionInfo | null): Awaitable<void>;
 	/**
-	 * Resolves once the given shard should be allowed to identify, or rejects if the operation was aborted
+	 * Resolves once the given shard should be allowed to identify
+	 * This should correctly handle the signal and reject with an abort error if the operation is aborted.
+	 * Other errors will cause the shard to reconnect.
 	 */
 	waitForIdentify(shardId: number, signal: AbortSignal): Promise<void>;
 }

packages/ws/src/strategies/context/WorkerContextFetchingStrategy.ts

@@ -9,17 +9,13 @@ import {
 } from '../sharding/WorkerShardingStrategy.js';
 import type { FetchingStrategyOptions, IContextFetchingStrategy } from './IContextFetchingStrategy.js';
 
-// Because the global types are incomplete for whatever reason
-interface PolyFillAbortSignal {
-	readonly aborted: boolean;
-	addEventListener(type: 'abort', listener: () => void): void;
-	removeEventListener(type: 'abort', listener: () => void): void;
-}
-
 export class WorkerContextFetchingStrategy implements IContextFetchingStrategy {
 	private readonly sessionPromises = new Collection<number, (session: SessionInfo | null) => void>();
 
-	private readonly waitForIdentifyPromises = new Collection<number, { reject(): void; resolve(): void }>();
+	private readonly waitForIdentifyPromises = new Collection<
+		number,
+		{ reject(error: unknown): void; resolve(): void; signal: AbortSignal }
+	>();
 
 	public constructor(public readonly options: FetchingStrategyOptions) {
 		if (isMainThread) {
@@ -37,7 +33,8 @@ export class WorkerContextFetchingStrategy implements IContextFetchingStrategy {
 				if (payload.ok) {
 					promise?.resolve();
 				} else {
-					promise?.reject();
+					// We need to make sure we reject with an abort error
+					promise?.reject(promise.signal.reason);
 				}
 
 				this.waitForIdentifyPromises.delete(payload.nonce);
@@ -77,7 +74,7 @@ export class WorkerContextFetchingStrategy implements IContextFetchingStrategy {
 		};
 		const promise = new Promise<void>((resolve, reject) =>
 			// eslint-disable-next-line no-promise-executor-return
-			this.waitForIdentifyPromises.set(nonce, { resolve, reject }),
+			this.waitForIdentifyPromises.set(nonce, { signal, resolve, reject }),
 		);
 
 		parentPort!.postMessage(payload);
@@ -91,12 +88,12 @@ export class WorkerContextFetchingStrategy implements IContextFetchingStrategy {
 			parentPort!.postMessage(payload);
 		};
 
-		(signal as unknown as PolyFillAbortSignal).addEventListener('abort', listener);
+		signal.addEventListener('abort', listener);
 
 		try {
 			await promise;
 		} finally {
-			(signal as unknown as PolyFillAbortSignal).removeEventListener('abort', listener);
+			signal.removeEventListener('abort', listener);
 		}
 	}
 }

packages/ws/src/ws/WebSocketShard.ts

@@ -387,8 +387,21 @@ export class WebSocketShard extends AsyncEventEmitter<WebSocketShardEventsMap> {
 		try {
 			await this.strategy.waitForIdentify(this.id, controller.signal);
 		} catch {
-			this.debug(['Was waiting for an identify, but the shard closed in the meantime']);
-			return;
+			if (controller.signal.aborted) {
+				this.debug(['Was waiting for an identify, but the shard closed in the meantime']);
+				return;
+			}
+
+			this.debug([
+				'IContextFetchingStrategy#waitForIdentify threw an unknown error.',
+				"If you're using a custom strategy, this is probably nothing to worry about.",
+				"If you're not, please open an issue on GitHub.",
+			]);
+
+			await this.destroy({
+				reason: 'Identify throttling logic failed',
+				recover: WebSocketShardDestroyRecovery.Resume,
+			});
 		} finally {
 			this.off(WebSocketShardEvents.Closed, closeHandler);
 		}