refactor: abstract identify throttling and correct max_concurrency handling (#9375)

* refactor: properly support max_concurrency ratelimit keys

* fix: properly block for same key

* chore: export session state

* chore: throttler no longer requires manager

* refactor: abstract throttlers

* chore: proper member order

* chore: remove leftover debug log

* chore: use @link tag in doc comment

Co-authored-by: Jiralite <33201955+Jiralite@users.noreply.github.com>

* chore: suggested changes

* fix(WebSocketShard): cancel identify if the shard closed in the meantime

* refactor(throttlers): support abort signals

* fix: memory leak

* chore: remove leftover

---------

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

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

@@ -5,7 +5,13 @@ import type { SessionInfo, WebSocketManager, WebSocketManagerOptions } from '../
 export interface FetchingStrategyOptions
 	extends Omit<
 		WebSocketManagerOptions,
-		'buildStrategy' | 'rest' | 'retrieveSessionInfo' | 'shardCount' | 'shardIds' | 'updateSessionInfo'
+		| 'buildIdentifyThrottler'
+		| 'buildStrategy'
+		| 'rest'
+		| 'retrieveSessionInfo'
+		| 'shardCount'
+		| 'shardIds'
+		| 'updateSessionInfo'
 	> {
 	readonly gatewayInformation: APIGatewayBotInfo;
 	readonly shardCount: number;
@@ -18,13 +24,25 @@ export interface IContextFetchingStrategy {
 	readonly options: FetchingStrategyOptions;
 	retrieveSessionInfo(shardId: number): Awaitable<SessionInfo | null>;
 	updateSessionInfo(shardId: number, sessionInfo: SessionInfo | null): Awaitable<void>;
-	waitForIdentify(): Promise<void>;
+	/**
+	 * Resolves once the given shard should be allowed to identify, or rejects if the operation was aborted
+	 */
+	waitForIdentify(shardId: number, signal: AbortSignal): Promise<void>;
 }
 
 export async function managerToFetchingStrategyOptions(manager: WebSocketManager): Promise<FetchingStrategyOptions> {
-	// eslint-disable-next-line @typescript-eslint/unbound-method
-	const { buildStrategy, retrieveSessionInfo, updateSessionInfo, shardCount, shardIds, rest, ...managerOptions } =
-		manager.options;
+	/* eslint-disable @typescript-eslint/unbound-method */
+	const {
+		buildIdentifyThrottler,
+		buildStrategy,
+		retrieveSessionInfo,
+		updateSessionInfo,
+		shardCount,
+		shardIds,
+		rest,
+		...managerOptions
+	} = manager.options;
+	/* eslint-enable @typescript-eslint/unbound-method */
 
 	return {
 		...managerOptions,

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

@@ -1,28 +1,25 @@
-import { IdentifyThrottler } from '../../utils/IdentifyThrottler.js';
+import type { IIdentifyThrottler } from '../../throttling/IIdentifyThrottler.js';
 import type { SessionInfo, WebSocketManager } from '../../ws/WebSocketManager.js';
 import type { FetchingStrategyOptions, IContextFetchingStrategy } from './IContextFetchingStrategy.js';
 
 export class SimpleContextFetchingStrategy implements IContextFetchingStrategy {
 	// This strategy assumes every shard is running under the same process - therefore we need a single
 	// IdentifyThrottler per manager.
-	private static throttlerCache = new WeakMap<WebSocketManager, IdentifyThrottler>();
+	private static throttlerCache = new WeakMap<WebSocketManager, IIdentifyThrottler>();
 
-	private static ensureThrottler(manager: WebSocketManager): IdentifyThrottler {
-		const existing = SimpleContextFetchingStrategy.throttlerCache.get(manager);
-		if (existing) {
-			return existing;
+	private static async ensureThrottler(manager: WebSocketManager): Promise<IIdentifyThrottler> {
+		const throttler = SimpleContextFetchingStrategy.throttlerCache.get(manager);
+		if (throttler) {
+			return throttler;
 		}
 
-		const throttler = new IdentifyThrottler(manager);
-		SimpleContextFetchingStrategy.throttlerCache.set(manager, throttler);
-		return throttler;
+		const newThrottler = await manager.options.buildIdentifyThrottler(manager);
+		SimpleContextFetchingStrategy.throttlerCache.set(manager, newThrottler);
+
+		return newThrottler;
 	}
 
-	private readonly throttler: IdentifyThrottler;
-
-	public constructor(private readonly manager: WebSocketManager, public readonly options: FetchingStrategyOptions) {
-		this.throttler = SimpleContextFetchingStrategy.ensureThrottler(manager);
-	}
+	public constructor(private readonly manager: WebSocketManager, public readonly options: FetchingStrategyOptions) {}
 
 	public async retrieveSessionInfo(shardId: number): Promise<SessionInfo | null> {
 		return this.manager.options.retrieveSessionInfo(shardId);
@@ -32,7 +29,8 @@ export class SimpleContextFetchingStrategy implements IContextFetchingStrategy {
 		return this.manager.options.updateSessionInfo(shardId, sessionInfo);
 	}
 
-	public async waitForIdentify(): Promise<void> {
-		await this.throttler.waitForIdentify();
+	public async waitForIdentify(shardId: number, signal: AbortSignal): Promise<void> {
+		const throttler = await SimpleContextFetchingStrategy.ensureThrottler(this.manager);
+		await throttler.waitForIdentify(shardId, signal);
 	}
 }

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

@@ -9,10 +9,17 @@ 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, () => void>();
+	private readonly waitForIdentifyPromises = new Collection<number, { reject(): void; resolve(): void }>();
 
 	public constructor(public readonly options: FetchingStrategyOptions) {
 		if (isMainThread) {
@@ -25,8 +32,14 @@ export class WorkerContextFetchingStrategy implements IContextFetchingStrategy {
 				this.sessionPromises.delete(payload.nonce);
 			}
 
-			if (payload.op === WorkerSendPayloadOp.ShardCanIdentify) {
-				this.waitForIdentifyPromises.get(payload.nonce)?.();
+			if (payload.op === WorkerSendPayloadOp.ShardIdentifyResponse) {
+				const promise = this.waitForIdentifyPromises.get(payload.nonce);
+				if (payload.ok) {
+					promise?.resolve();
+				} else {
+					promise?.reject();
+				}
+
 				this.waitForIdentifyPromises.delete(payload.nonce);
 			}
 		});
@@ -34,11 +47,11 @@ export class WorkerContextFetchingStrategy implements IContextFetchingStrategy {
 
 	public async retrieveSessionInfo(shardId: number): Promise<SessionInfo | null> {
 		const nonce = Math.random();
-		const payload = {
+		const payload: WorkerReceivePayload = {
 			op: WorkerReceivePayloadOp.RetrieveSessionInfo,
 			shardId,
 			nonce,
-		} satisfies WorkerReceivePayload;
+		};
 		// eslint-disable-next-line no-promise-executor-return
 		const promise = new Promise<SessionInfo | null>((resolve) => this.sessionPromises.set(nonce, resolve));
 		parentPort!.postMessage(payload);
@@ -46,23 +59,44 @@ export class WorkerContextFetchingStrategy implements IContextFetchingStrategy {
 	}
 
 	public updateSessionInfo(shardId: number, sessionInfo: SessionInfo | null) {
-		const payload = {
+		const payload: WorkerReceivePayload = {
 			op: WorkerReceivePayloadOp.UpdateSessionInfo,
 			shardId,
 			session: sessionInfo,
-		} satisfies WorkerReceivePayload;
+		};
 		parentPort!.postMessage(payload);
 	}
 
-	public async waitForIdentify(): Promise<void> {
+	public async waitForIdentify(shardId: number, signal: AbortSignal): Promise<void> {
 		const nonce = Math.random();
-		const payload = {
+
+		const payload: WorkerReceivePayload = {
 			op: WorkerReceivePayloadOp.WaitForIdentify,
 			nonce,
-		} satisfies WorkerReceivePayload;
-		// eslint-disable-next-line no-promise-executor-return
-		const promise = new Promise<void>((resolve) => this.waitForIdentifyPromises.set(nonce, resolve));
+			shardId,
+		};
+		const promise = new Promise<void>((resolve, reject) =>
+			// eslint-disable-next-line no-promise-executor-return
+			this.waitForIdentifyPromises.set(nonce, { resolve, reject }),
+		);
+
 		parentPort!.postMessage(payload);
-		return promise;
+
+		const listener = () => {
+			const payload: WorkerReceivePayload = {
+				op: WorkerReceivePayloadOp.CancelIdentify,
+				nonce,
+			};
+
+			parentPort!.postMessage(payload);
+		};
+
+		(signal as unknown as PolyFillAbortSignal).addEventListener('abort', listener);
+
+		try {
+			await promise;
+		} finally {
+			(signal as unknown as PolyFillAbortSignal).removeEventListener('abort', listener);
+		}
 	}
 }

packages/ws/src/strategies/sharding/SimpleShardingStrategy.ts

@@ -23,6 +23,7 @@ export class SimpleShardingStrategy implements IShardingStrategy {
 	 */
 	public async spawn(shardIds: number[]) {
 		const strategyOptions = await managerToFetchingStrategyOptions(this.manager);
+
 		for (const shardId of shardIds) {
 			const strategy = new SimpleContextFetchingStrategy(this.manager, strategyOptions);
 			const shard = new WebSocketShard(strategy, shardId);

packages/ws/src/strategies/sharding/WorkerShardingStrategy.ts

@@ -3,7 +3,7 @@ import { join, isAbsolute, resolve } from 'node:path';
 import { Worker } from 'node:worker_threads';
 import { Collection } from '@discordjs/collection';
 import type { GatewaySendPayload } from 'discord-api-types/v10';
-import { IdentifyThrottler } from '../../utils/IdentifyThrottler.js';
+import type { IIdentifyThrottler } from '../../throttling/IIdentifyThrottler';
 import type { SessionInfo, WebSocketManager } from '../../ws/WebSocketManager';
 import type { WebSocketShardDestroyOptions, WebSocketShardEvents, WebSocketShardStatus } from '../../ws/WebSocketShard';
 import { managerToFetchingStrategyOptions, type FetchingStrategyOptions } from '../context/IContextFetchingStrategy.js';
@@ -18,14 +18,14 @@ export enum WorkerSendPayloadOp {
 	Destroy,
 	Send,
 	SessionInfoResponse,
-	ShardCanIdentify,
+	ShardIdentifyResponse,
 	FetchStatus,
 }
 
 export type WorkerSendPayload =
+	| { nonce: number; ok: boolean; op: WorkerSendPayloadOp.ShardIdentifyResponse }
 	| { nonce: number; op: WorkerSendPayloadOp.FetchStatus; shardId: number }
 	| { nonce: number; op: WorkerSendPayloadOp.SessionInfoResponse; session: SessionInfo | null }
-	| { nonce: number; op: WorkerSendPayloadOp.ShardCanIdentify }
 	| { op: WorkerSendPayloadOp.Connect; shardId: number }
 	| { op: WorkerSendPayloadOp.Destroy; options?: WebSocketShardDestroyOptions; shardId: number }
 	| { op: WorkerSendPayloadOp.Send; payload: GatewaySendPayload; shardId: number };
@@ -39,14 +39,16 @@ export enum WorkerReceivePayloadOp {
 	WaitForIdentify,
 	FetchStatusResponse,
 	WorkerReady,
+	CancelIdentify,
 }
 
 export type WorkerReceivePayload =
 	// Can't seem to get a type-safe union based off of the event, so I'm sadly leaving data as any for now
 	| { data: any; event: WebSocketShardEvents; op: WorkerReceivePayloadOp.Event; shardId: number }
+	| { nonce: number; op: WorkerReceivePayloadOp.CancelIdentify }
 	| { nonce: number; op: WorkerReceivePayloadOp.FetchStatusResponse; status: WebSocketShardStatus }
 	| { nonce: number; op: WorkerReceivePayloadOp.RetrieveSessionInfo; shardId: number }
-	| { nonce: number; op: WorkerReceivePayloadOp.WaitForIdentify }
+	| { nonce: number; op: WorkerReceivePayloadOp.WaitForIdentify; shardId: number }
 	| { op: WorkerReceivePayloadOp.Connected; shardId: number }
 	| { op: WorkerReceivePayloadOp.Destroyed; shardId: number }
 	| { op: WorkerReceivePayloadOp.UpdateSessionInfo; session: SessionInfo | null; shardId: number }
@@ -84,11 +86,12 @@ export class WorkerShardingStrategy implements IShardingStrategy {
 
 	private readonly fetchStatusPromises = new Collection<number, (status: WebSocketShardStatus) => void>();
 
-	private readonly throttler: IdentifyThrottler;
+	private readonly waitForIdentifyControllers = new Collection<number, AbortController>();
+
+	private throttler?: IIdentifyThrottler;
 
 	public constructor(manager: WebSocketManager, options: WorkerShardingStrategyOptions) {
 		this.manager = manager;
-		this.throttler = new IdentifyThrottler(manager);
 		this.options = options;
 	}
 
@@ -122,10 +125,10 @@ export class WorkerShardingStrategy implements IShardingStrategy {
 		const promises = [];
 
 		for (const [shardId, worker] of this.#workerByShardId.entries()) {
-			const payload = {
+			const payload: WorkerSendPayload = {
 				op: WorkerSendPayloadOp.Connect,
 				shardId,
-			} satisfies WorkerSendPayload;
+			};
 
 			// eslint-disable-next-line no-promise-executor-return
 			const promise = new Promise<void>((resolve) => this.connectPromises.set(shardId, resolve));
@@ -143,11 +146,11 @@ export class WorkerShardingStrategy implements IShardingStrategy {
 		const promises = [];
 
 		for (const [shardId, worker] of this.#workerByShardId.entries()) {
-			const payload = {
+			const payload: WorkerSendPayload = {
 				op: WorkerSendPayloadOp.Destroy,
 				shardId,
 				options,
-			} satisfies WorkerSendPayload;
+			};
 
 			promises.push(
 				// eslint-disable-next-line no-promise-executor-return, promise/prefer-await-to-then
@@ -171,11 +174,11 @@ export class WorkerShardingStrategy implements IShardingStrategy {
 			throw new Error(`No worker found for shard ${shardId}`);
 		}
 
-		const payload = {
+		const payload: WorkerSendPayload = {
 			op: WorkerSendPayloadOp.Send,
 			shardId,
 			payload: data,
-		} satisfies WorkerSendPayload;
+		};
 		worker.postMessage(payload);
 	}
 
@@ -187,11 +190,11 @@ export class WorkerShardingStrategy implements IShardingStrategy {
 
 		for (const [shardId, worker] of this.#workerByShardId.entries()) {
 			const nonce = Math.random();
-			const payload = {
+			const payload: WorkerSendPayload = {
 				op: WorkerSendPayloadOp.FetchStatus,
 				shardId,
 				nonce,
-			} satisfies WorkerSendPayload;
+			};
 
 			// eslint-disable-next-line no-promise-executor-return
 			const promise = new Promise<WebSocketShardStatus>((resolve) => this.fetchStatusPromises.set(nonce, resolve));
@@ -297,10 +300,21 @@ export class WorkerShardingStrategy implements IShardingStrategy {
 			}
 
 			case WorkerReceivePayloadOp.WaitForIdentify: {
-				await this.throttler.waitForIdentify();
+				const throttler = await this.ensureThrottler();
+
+				// If this rejects it means we aborted, in which case we reply elsewhere.
+				try {
+					const controller = new AbortController();
+					this.waitForIdentifyControllers.set(payload.nonce, controller);
+					await throttler.waitForIdentify(payload.shardId, controller.signal);
+				} catch {
+					return;
+				}
+
 				const response: WorkerSendPayload = {
-					op: WorkerSendPayloadOp.ShardCanIdentify,
+					op: WorkerSendPayloadOp.ShardIdentifyResponse,
 					nonce: payload.nonce,
+					ok: true,
 				};
 				worker.postMessage(response);
 				break;
@@ -315,6 +329,25 @@ export class WorkerShardingStrategy implements IShardingStrategy {
 			case WorkerReceivePayloadOp.WorkerReady: {
 				break;
 			}
+
+			case WorkerReceivePayloadOp.CancelIdentify: {
+				this.waitForIdentifyControllers.get(payload.nonce)?.abort();
+				this.waitForIdentifyControllers.delete(payload.nonce);
+
+				const response: WorkerSendPayload = {
+					op: WorkerSendPayloadOp.ShardIdentifyResponse,
+					nonce: payload.nonce,
+					ok: false,
+				};
+				worker.postMessage(response);
+
+				break;
+			}
 		}
 	}
+
+	private async ensureThrottler(): Promise<IIdentifyThrottler> {
+		this.throttler ??= await this.manager.options.buildIdentifyThrottler(this.manager);
+		return this.throttler;
+	}
 }