refactor(WebSocketShard): identify throttling (#8888)

* refactor(WebSocketShard): identify throttling

* chore: add worker handling

* refactor: worker handling

* chore: update tests

* chore: use satisfies where applicable

* chore: add informative comment

* chore: apply suggestions

* refactor(SimpleContextFetchingStrategy): support multiple managers

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

@@ -1,6 +1,5 @@
 import { Collection } from '@discordjs/collection';
 import type { GatewaySendPayload } from 'discord-api-types/v10';
-import { IdentifyThrottler } from '../../utils/IdentifyThrottler.js';
 import type { WebSocketManager } from '../../ws/WebSocketManager';
 import { WebSocketShard, WebSocketShardEvents, type WebSocketShardDestroyOptions } from '../../ws/WebSocketShard.js';
 import { managerToFetchingStrategyOptions } from '../context/IContextFetchingStrategy.js';
@@ -15,11 +14,8 @@ export class SimpleShardingStrategy implements IShardingStrategy {
 
 	private readonly shards = new Collection<number, WebSocketShard>();
 
-	private readonly throttler: IdentifyThrottler;
-
 	public constructor(manager: WebSocketManager) {
 		this.manager = manager;
-		this.throttler = new IdentifyThrottler(manager);
 	}
 
 	/**
@@ -46,7 +42,6 @@ export class SimpleShardingStrategy implements IShardingStrategy {
 		const promises = [];
 
 		for (const shard of this.shards.values()) {
-			await this.throttler.waitForIdentify();
 			promises.push(shard.connect());
 		}
 

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

@@ -18,10 +18,12 @@ export enum WorkerSendPayloadOp {
 	Destroy,
 	Send,
 	SessionInfoResponse,
+	ShardCanIdentify,
 }
 
 export type WorkerSendPayload =
 	| { 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 };
@@ -32,12 +34,14 @@ export enum WorkerRecievePayloadOp {
 	Event,
 	RetrieveSessionInfo,
 	UpdateSessionInfo,
+	WaitForIdentify,
 }
 
 export type WorkerRecievePayload =
 	// 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: WorkerRecievePayloadOp.Event; shardId: number }
 	| { nonce: number; op: WorkerRecievePayloadOp.RetrieveSessionInfo; shardId: number }
+	| { nonce: number; op: WorkerRecievePayloadOp.WaitForIdentify }
 	| { op: WorkerRecievePayloadOp.Connected; shardId: number }
 	| { op: WorkerRecievePayloadOp.Destroyed; shardId: number }
 	| { op: WorkerRecievePayloadOp.UpdateSessionInfo; session: SessionInfo | null; shardId: number };
@@ -118,12 +122,10 @@ export class WorkerShardingStrategy implements IShardingStrategy {
 		const promises = [];
 
 		for (const [shardId, worker] of this.#workerByShardId.entries()) {
-			await this.throttler.waitForIdentify();
-
-			const payload: WorkerSendPayload = {
+			const payload = {
 				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));
@@ -141,11 +143,11 @@ export class WorkerShardingStrategy implements IShardingStrategy {
 		const promises = [];
 
 		for (const [shardId, worker] of this.#workerByShardId.entries()) {
-			const payload: WorkerSendPayload = {
+			const payload = {
 				op: WorkerSendPayloadOp.Destroy,
 				shardId,
 				options,
-			};
+			} satisfies WorkerSendPayload;
 
 			promises.push(
 				// eslint-disable-next-line no-promise-executor-return, promise/prefer-await-to-then
@@ -169,11 +171,11 @@ export class WorkerShardingStrategy implements IShardingStrategy {
 			throw new Error(`No worker found for shard ${shardId}`);
 		}
 
-		const payload: WorkerSendPayload = {
+		const payload = {
 			op: WorkerSendPayloadOp.Send,
 			shardId,
 			payload: data,
-		};
+		} satisfies WorkerSendPayload;
 		worker.postMessage(payload);
 	}
 
@@ -213,6 +215,16 @@ export class WorkerShardingStrategy implements IShardingStrategy {
 				await this.manager.options.updateSessionInfo(payload.shardId, payload.session);
 				break;
 			}
+
+			case WorkerRecievePayloadOp.WaitForIdentify: {
+				await this.throttler.waitForIdentify();
+				const response: WorkerSendPayload = {
+					op: WorkerSendPayloadOp.ShardCanIdentify,
+					nonce: payload.nonce,
+				};
+				worker.postMessage(response);
+				break;
+			}
 		}
 	}
 }

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

@@ -40,12 +40,12 @@ for (const shardId of data.shardIds) {
 	for (const event of Object.values(WebSocketShardEvents)) {
 		// @ts-expect-error: Event types incompatible
 		shard.on(event, (data) => {
-			const payload: WorkerRecievePayload = {
+			const payload = {
 				op: WorkerRecievePayloadOp.Event,
 				event,
 				data,
 				shardId,
-			};
+			} satisfies WorkerRecievePayload;
 			parentPort!.postMessage(payload);
 		});
 	}
@@ -93,5 +93,9 @@ parentPort!
 			case WorkerSendPayloadOp.SessionInfoResponse: {
 				break;
 			}
+
+			case WorkerSendPayloadOp.ShardCanIdentify: {
+				break;
+			}
 		}
 	});