diff --git a/packages/ws/src/strategies/context/IContextFetchingStrategy.ts b/packages/ws/src/strategies/context/IContextFetchingStrategy.ts index b57f0ff82..6abe25171 100644 --- a/packages/ws/src/strategies/context/IContextFetchingStrategy.ts +++ b/packages/ws/src/strategies/context/IContextFetchingStrategy.ts @@ -25,7 +25,9 @@ export interface IContextFetchingStrategy { retrieveSessionInfo(shardId: number): Awaitable; updateSessionInfo(shardId: number, sessionInfo: SessionInfo | null): Awaitable; /** - * 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; } diff --git a/packages/ws/src/strategies/context/WorkerContextFetchingStrategy.ts b/packages/ws/src/strategies/context/WorkerContextFetchingStrategy.ts index 5a79eb88f..474b0ddaa 100644 --- a/packages/ws/src/strategies/context/WorkerContextFetchingStrategy.ts +++ b/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 void>(); - private readonly waitForIdentifyPromises = new Collection(); + 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((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); } } } diff --git a/packages/ws/src/ws/WebSocketShard.ts b/packages/ws/src/ws/WebSocketShard.ts index 5671a69b2..7b34f6a99 100644 --- a/packages/ws/src/ws/WebSocketShard.ts +++ b/packages/ws/src/ws/WebSocketShard.ts @@ -387,8 +387,21 @@ export class WebSocketShard extends AsyncEventEmitter { 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); }