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/utils/IdentifyThrottler.ts

@@ -1,39 +0,0 @@
-import { setTimeout as sleep } from 'node:timers/promises';
-import { AsyncQueue } from '@sapphire/async-queue';
-import type { WebSocketManager } from '../ws/WebSocketManager.js';
-
-export class IdentifyThrottler {
-	private readonly queue = new AsyncQueue();
-
-	private identifyState = {
-		remaining: 0,
-		resetsAt: Number.POSITIVE_INFINITY,
-	};
-
-	public constructor(private readonly manager: WebSocketManager) {}
-
-	public async waitForIdentify(): Promise<void> {
-		await this.queue.wait();
-
-		try {
-			if (this.identifyState.remaining <= 0) {
-				const diff = this.identifyState.resetsAt - Date.now();
-				if (diff <= 5_000) {
-					// To account for the latency the IDENTIFY payload goes through, we add a bit more wait time
-					const time = diff + Math.random() * 1_500;
-					await sleep(time);
-				}
-
-				const info = await this.manager.fetchGatewayInformation();
-				this.identifyState = {
-					remaining: info.session_start_limit.max_concurrency,
-					resetsAt: Date.now() + 5_000,
-				};
-			}
-
-			this.identifyState.remaining--;
-		} finally {
-			this.queue.shift();
-		}
-	}
-}

packages/ws/src/utils/WorkerBootstrapper.ts

@@ -117,7 +117,7 @@ export class WorkerBootstrapper {
 						break;
 					}
 
-					case WorkerSendPayloadOp.ShardCanIdentify: {
+					case WorkerSendPayloadOp.ShardIdentifyResponse: {
 						break;
 					}
 
@@ -127,11 +127,11 @@ export class WorkerBootstrapper {
 							throw new Error(`Shard ${payload.shardId} does not exist`);
 						}
 
-						const response = {
+						const response: WorkerReceivePayload = {
 							op: WorkerReceivePayloadOp.FetchStatusResponse,
 							status: shard.status,
 							nonce: payload.nonce,
-						} satisfies WorkerReceivePayload;
+						};
 
 						parentPort!.postMessage(response);
 						break;
@@ -150,12 +150,12 @@ export class WorkerBootstrapper {
 			for (const event of options.forwardEvents ?? Object.values(WebSocketShardEvents)) {
 				// @ts-expect-error: Event types incompatible
 				shard.on(event, (data) => {
-					const payload = {
+					const payload: WorkerReceivePayload = {
 						op: WorkerReceivePayloadOp.Event,
 						event,
 						data,
 						shardId,
-					} satisfies WorkerReceivePayload;
+					};
 					parentPort!.postMessage(payload);
 				});
 			}
@@ -168,9 +168,9 @@ export class WorkerBootstrapper {
 		// Lastly, start listening to messages from the parent thread
 		this.setupThreadEvents();
 
-		const message = {
+		const message: WorkerReceivePayload = {
 			op: WorkerReceivePayloadOp.WorkerReady,
-		} satisfies WorkerReceivePayload;
+		};
 		parentPort!.postMessage(message);
 	}
 }

packages/ws/src/utils/constants.ts

@@ -3,7 +3,8 @@ import { Collection } from '@discordjs/collection';
 import { lazy } from '@discordjs/util';
 import { APIVersion, GatewayOpcodes } from 'discord-api-types/v10';
 import { SimpleShardingStrategy } from '../strategies/sharding/SimpleShardingStrategy.js';
-import type { SessionInfo, OptionalWebSocketManagerOptions } from '../ws/WebSocketManager.js';
+import { SimpleIdentifyThrottler } from '../throttling/SimpleIdentifyThrottler.js';
+import type { SessionInfo, OptionalWebSocketManagerOptions, WebSocketManager } from '../ws/WebSocketManager.js';
 import type { SendRateLimitState } from '../ws/WebSocketShard.js';
 
 /**
@@ -28,6 +29,10 @@ const getDefaultSessionStore = lazy(() => new Collection<number, SessionInfo | n
  * Default options used by the manager
  */
 export const DefaultWebSocketManagerOptions = {
+	async buildIdentifyThrottler(manager: WebSocketManager) {
+		const info = await manager.fetchGatewayInformation();
+		return new SimpleIdentifyThrottler(info.session_start_limit.max_concurrency);
+	},
 	buildStrategy: (manager) => new SimpleShardingStrategy(manager),
 	shardCount: null,
 	shardIds: null,