feat: @discordjs/ws (#8260)

Co-authored-by: Parbez <imranbarbhuiya.fsd@gmail.com>

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

@@ -0,0 +1,31 @@
+import type { Awaitable } from '@vladfrangu/async_event_emitter';
+import type { APIGatewayBotInfo } from 'discord-api-types/v10';
+import type { SessionInfo, WebSocketManager, WebSocketManagerOptions } from '../../ws/WebSocketManager';
+
+export interface FetchingStrategyOptions
+	extends Omit<
+		WebSocketManagerOptions,
+		'retrieveSessionInfo' | 'updateSessionInfo' | 'shardCount' | 'shardIds' | 'rest'
+	> {
+	readonly gatewayInformation: APIGatewayBotInfo;
+	readonly shardCount: number;
+}
+
+/**
+ * Strategies responsible solely for making manager information accessible
+ */
+export interface IContextFetchingStrategy {
+	readonly options: FetchingStrategyOptions;
+	retrieveSessionInfo: (shardId: number) => Awaitable<SessionInfo | null>;
+	updateSessionInfo: (shardId: number, sessionInfo: SessionInfo | null) => Awaitable<void>;
+}
+
+export async function managerToFetchingStrategyOptions(manager: WebSocketManager): Promise<FetchingStrategyOptions> {
+	const { retrieveSessionInfo, updateSessionInfo, shardCount, shardIds, rest, ...managerOptions } = manager.options;
+
+	return {
+		...managerOptions,
+		gatewayInformation: await manager.fetchGatewayInformation(),
+		shardCount: await manager.getShardCount(),
+	};
+}

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

@@ -0,0 +1,14 @@
+import type { FetchingStrategyOptions, IContextFetchingStrategy } from './IContextFetchingStrategy';
+import type { SessionInfo, WebSocketManager } from '../../ws/WebSocketManager';
+
+export class SimpleContextFetchingStrategy implements IContextFetchingStrategy {
+	public constructor(private readonly manager: WebSocketManager, public readonly options: FetchingStrategyOptions) {}
+
+	public async retrieveSessionInfo(shardId: number): Promise<SessionInfo | null> {
+		return this.manager.options.retrieveSessionInfo(shardId);
+	}
+
+	public updateSessionInfo(shardId: number, sessionInfo: SessionInfo | null) {
+		return this.manager.options.updateSessionInfo(shardId, sessionInfo);
+	}
+}

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

@@ -0,0 +1,49 @@
+import { isMainThread, parentPort } from 'node:worker_threads';
+import { Collection } from '@discordjs/collection';
+import type { FetchingStrategyOptions, IContextFetchingStrategy } from './IContextFetchingStrategy';
+import type { SessionInfo } from '../../ws/WebSocketManager';
+import {
+	WorkerRecievePayload,
+	WorkerRecievePayloadOp,
+	WorkerSendPayload,
+	WorkerSendPayloadOp,
+} from '../sharding/WorkerShardingStrategy';
+
+export class WorkerContextFetchingStrategy implements IContextFetchingStrategy {
+	private readonly sessionPromises = new Collection<number, (session: SessionInfo | null) => void>();
+
+	public constructor(public readonly options: FetchingStrategyOptions) {
+		if (isMainThread) {
+			throw new Error('Cannot instantiate WorkerContextFetchingStrategy on the main thread');
+		}
+
+		parentPort!.on('message', (payload: WorkerSendPayload) => {
+			if (payload.op === WorkerSendPayloadOp.SessionInfoResponse) {
+				const resolve = this.sessionPromises.get(payload.nonce);
+				resolve?.(payload.session);
+				this.sessionPromises.delete(payload.nonce);
+			}
+		});
+	}
+
+	public async retrieveSessionInfo(shardId: number): Promise<SessionInfo | null> {
+		const nonce = Math.random();
+		const payload: WorkerRecievePayload = {
+			op: WorkerRecievePayloadOp.RetrieveSessionInfo,
+			shardId,
+			nonce,
+		};
+		const promise = new Promise<SessionInfo | null>((resolve) => this.sessionPromises.set(nonce, resolve));
+		parentPort!.postMessage(payload);
+		return promise;
+	}
+
+	public updateSessionInfo(shardId: number, sessionInfo: SessionInfo | null) {
+		const payload: WorkerRecievePayload = {
+			op: WorkerRecievePayloadOp.UpdateSessionInfo,
+			shardId,
+			session: sessionInfo,
+		};
+		parentPort!.postMessage(payload);
+	}
+}

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

@@ -0,0 +1,25 @@
+import type { GatewaySendPayload } from 'discord-api-types/v10';
+import type { Awaitable } from '../../utils/utils';
+import type { WebSocketShardDestroyOptions } from '../../ws/WebSocketShard';
+
+/**
+ * Strategies responsible for spawning, initializing connections, destroying shards, and relaying events
+ */
+export interface IShardingStrategy {
+	/**
+	 * Spawns all the shards
+	 */
+	spawn: (shardIds: number[]) => Awaitable<void>;
+	/**
+	 * Initializes all the shards
+	 */
+	connect: () => Awaitable<void>;
+	/**
+	 * Destroys all the shards
+	 */
+	destroy: (options?: Omit<WebSocketShardDestroyOptions, 'recover'>) => Awaitable<void>;
+	/**
+	 * Sends a payload to a shard
+	 */
+	send: (shardId: number, payload: GatewaySendPayload) => Awaitable<void>;
+}

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

@@ -0,0 +1,64 @@
+import { Collection } from '@discordjs/collection';
+import type { GatewaySendPayload } from 'discord-api-types/v10';
+import type { IShardingStrategy } from './IShardingStrategy';
+import { IdentifyThrottler } from '../../utils/IdentifyThrottler';
+import type { WebSocketManager } from '../../ws/WebSocketManager';
+import { WebSocketShard, WebSocketShardDestroyOptions, WebSocketShardEvents } from '../../ws/WebSocketShard';
+import { managerToFetchingStrategyOptions } from '../context/IContextFetchingStrategy';
+import { SimpleContextFetchingStrategy } from '../context/SimpleContextFetchingStrategy';
+
+/**
+ * Simple strategy that just spawns shards in the current process
+ */
+export class SimpleShardingStrategy implements IShardingStrategy {
+	private readonly manager: WebSocketManager;
+	private readonly shards = new Collection<number, WebSocketShard>();
+
+	private readonly throttler: IdentifyThrottler;
+
+	public constructor(manager: WebSocketManager) {
+		this.manager = manager;
+		this.throttler = new IdentifyThrottler(manager);
+	}
+
+	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);
+			for (const event of Object.values(WebSocketShardEvents)) {
+				// @ts-expect-error
+				shard.on(event, (payload) => this.manager.emit(event, { ...payload, shardId }));
+			}
+			this.shards.set(shardId, shard);
+		}
+	}
+
+	public async connect() {
+		const promises = [];
+
+		for (const shard of this.shards.values()) {
+			await this.throttler.waitForIdentify();
+			promises.push(shard.connect());
+		}
+
+		await Promise.all(promises);
+	}
+
+	public async destroy(options?: Omit<WebSocketShardDestroyOptions, 'recover'>) {
+		const promises = [];
+
+		for (const shard of this.shards.values()) {
+			promises.push(shard.destroy(options));
+		}
+
+		await Promise.all(promises);
+		this.shards.clear();
+	}
+
+	public send(shardId: number, payload: GatewaySendPayload) {
+		const shard = this.shards.get(shardId);
+		if (!shard) throw new Error(`Shard ${shardId} not found`);
+		return shard.send(payload);
+	}
+}

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

@@ -0,0 +1,203 @@
+import { once } from 'node:events';
+import { join } from 'node:path';
+import { Worker } from 'node:worker_threads';
+import { Collection } from '@discordjs/collection';
+import type { GatewaySendPayload } from 'discord-api-types/v10';
+import type { IShardingStrategy } from './IShardingStrategy';
+import { IdentifyThrottler } from '../../utils/IdentifyThrottler';
+import type { SessionInfo, WebSocketManager } from '../../ws/WebSocketManager';
+import type { WebSocketShardDestroyOptions, WebSocketShardEvents } from '../../ws/WebSocketShard';
+import { FetchingStrategyOptions, managerToFetchingStrategyOptions } from '../context/IContextFetchingStrategy';
+
+export interface WorkerData extends FetchingStrategyOptions {
+	shardIds: number[];
+}
+
+export enum WorkerSendPayloadOp {
+	Connect,
+	Destroy,
+	Send,
+	SessionInfoResponse,
+}
+
+export type WorkerSendPayload =
+	| { op: WorkerSendPayloadOp.Connect; shardId: number }
+	| { op: WorkerSendPayloadOp.Destroy; shardId: number; options?: WebSocketShardDestroyOptions }
+	| { op: WorkerSendPayloadOp.Send; shardId: number; payload: GatewaySendPayload }
+	| { op: WorkerSendPayloadOp.SessionInfoResponse; nonce: number; session: SessionInfo | null };
+
+export enum WorkerRecievePayloadOp {
+	Connected,
+	Destroyed,
+	Event,
+	RetrieveSessionInfo,
+	UpdateSessionInfo,
+}
+
+export type WorkerRecievePayload =
+	| { op: WorkerRecievePayloadOp.Connected; shardId: number }
+	| { op: WorkerRecievePayloadOp.Destroyed; shardId: number }
+	// Can't seem to get a type-safe union based off of the event, so I'm sadly leaving data as any for now
+	| { op: WorkerRecievePayloadOp.Event; shardId: number; event: WebSocketShardEvents; data: any }
+	| { op: WorkerRecievePayloadOp.RetrieveSessionInfo; shardId: number; nonce: number }
+	| { op: WorkerRecievePayloadOp.UpdateSessionInfo; shardId: number; session: SessionInfo | null };
+
+/**
+ * Options for a {@link WorkerShardingStrategy}
+ */
+export interface WorkerShardingStrategyOptions {
+	/**
+	 * Dictates how many shards should be spawned per worker thread.
+	 */
+	shardsPerWorker: number | 'all';
+}
+
+/**
+ * Strategy used to spawn threads in worker_threads
+ */
+export class WorkerShardingStrategy implements IShardingStrategy {
+	private readonly manager: WebSocketManager;
+	private readonly options: WorkerShardingStrategyOptions;
+
+	#workers: Worker[] = [];
+	readonly #workerByShardId = new Collection<number, Worker>();
+
+	private readonly connectPromises = new Collection<number, () => void>();
+	private readonly destroyPromises = new Collection<number, () => void>();
+
+	private readonly throttler: IdentifyThrottler;
+
+	public constructor(manager: WebSocketManager, options: WorkerShardingStrategyOptions) {
+		this.manager = manager;
+		this.throttler = new IdentifyThrottler(manager);
+		this.options = options;
+	}
+
+	public async spawn(shardIds: number[]) {
+		const shardsPerWorker = this.options.shardsPerWorker === 'all' ? shardIds.length : this.options.shardsPerWorker;
+		const strategyOptions = await managerToFetchingStrategyOptions(this.manager);
+
+		let shards = 0;
+		while (shards !== shardIds.length) {
+			const slice = shardIds.slice(shards, shardsPerWorker + shards);
+			const workerData: WorkerData = {
+				...strategyOptions,
+				shardIds: slice,
+			};
+
+			const worker = new Worker(join(__dirname, 'worker.cjs'), { workerData });
+			await once(worker, 'online');
+			worker
+				.on('error', (err) => {
+					throw err;
+				})
+				.on('messageerror', (err) => {
+					throw err;
+				})
+				// eslint-disable-next-line @typescript-eslint/no-misused-promises
+				.on('message', (payload: WorkerRecievePayload) => this.onMessage(worker, payload));
+
+			this.#workers.push(worker);
+			for (const shardId of slice) {
+				this.#workerByShardId.set(shardId, worker);
+			}
+
+			shards += slice.length;
+		}
+	}
+
+	public async connect() {
+		const promises = [];
+
+		for (const [shardId, worker] of this.#workerByShardId.entries()) {
+			await this.throttler.waitForIdentify();
+
+			const payload: WorkerSendPayload = {
+				op: WorkerSendPayloadOp.Connect,
+				shardId,
+			};
+
+			const promise = new Promise<void>((resolve) => this.connectPromises.set(shardId, resolve));
+			worker.postMessage(payload);
+			promises.push(promise);
+		}
+
+		await Promise.all(promises);
+	}
+
+	public async destroy(options: Omit<WebSocketShardDestroyOptions, 'recover'> = {}) {
+		const promises = [];
+
+		for (const [shardId, worker] of this.#workerByShardId.entries()) {
+			const payload: WorkerSendPayload = {
+				op: WorkerSendPayloadOp.Destroy,
+				shardId,
+				options,
+			};
+
+			promises.push(
+				new Promise<void>((resolve) => this.destroyPromises.set(shardId, resolve)).then(() => worker.terminate()),
+			);
+			worker.postMessage(payload);
+		}
+
+		this.#workers = [];
+		this.#workerByShardId.clear();
+
+		await Promise.all(promises);
+	}
+
+	public send(shardId: number, data: GatewaySendPayload) {
+		const worker = this.#workerByShardId.get(shardId);
+		if (!worker) {
+			throw new Error(`No worker found for shard ${shardId}`);
+		}
+
+		const payload: WorkerSendPayload = {
+			op: WorkerSendPayloadOp.Send,
+			shardId,
+			payload: data,
+		};
+		worker.postMessage(payload);
+	}
+
+	private async onMessage(worker: Worker, payload: WorkerRecievePayload) {
+		switch (payload.op) {
+			case WorkerRecievePayloadOp.Connected: {
+				const resolve = this.connectPromises.get(payload.shardId)!;
+				resolve();
+				this.connectPromises.delete(payload.shardId);
+				break;
+			}
+
+			case WorkerRecievePayloadOp.Destroyed: {
+				const resolve = this.destroyPromises.get(payload.shardId)!;
+				resolve();
+				this.destroyPromises.delete(payload.shardId);
+				break;
+			}
+
+			case WorkerRecievePayloadOp.Event: {
+				// eslint-disable-next-line @typescript-eslint/no-unsafe-argument
+				this.manager.emit(payload.event, { ...payload.data, shardId: payload.shardId });
+				break;
+			}
+
+			case WorkerRecievePayloadOp.RetrieveSessionInfo: {
+				const session = await this.manager.options.retrieveSessionInfo(payload.shardId);
+				const response: WorkerSendPayload = {
+					op: WorkerSendPayloadOp.SessionInfoResponse,
+					nonce: payload.nonce,
+					session,
+				};
+				worker.postMessage(response);
+				break;
+			}
+
+			case WorkerRecievePayloadOp.UpdateSessionInfo: {
+				await this.manager.options.updateSessionInfo(payload.shardId, payload.session);
+				break;
+			}
+		}
+	}
+}

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

@@ -0,0 +1,93 @@
+import { isMainThread, workerData, parentPort } from 'node:worker_threads';
+import { Collection } from '@discordjs/collection';
+import {
+	WorkerData,
+	WorkerRecievePayload,
+	WorkerRecievePayloadOp,
+	WorkerSendPayload,
+	WorkerSendPayloadOp,
+} from './WorkerShardingStrategy';
+import { WebSocketShard, WebSocketShardDestroyOptions, WebSocketShardEvents } from '../../ws/WebSocketShard';
+import { WorkerContextFetchingStrategy } from '../context/WorkerContextFetchingStrategy';
+
+if (isMainThread) {
+	throw new Error('Expected worker script to not be ran within the main thread');
+}
+
+const data = workerData as WorkerData;
+const shards = new Collection<number, WebSocketShard>();
+
+async function connect(shardId: number) {
+	const shard = shards.get(shardId);
+	if (!shard) {
+		throw new Error(`Shard ${shardId} does not exist`);
+	}
+	await shard.connect();
+}
+
+async function destroy(shardId: number, options?: WebSocketShardDestroyOptions) {
+	const shard = shards.get(shardId);
+	if (!shard) {
+		throw new Error(`Shard ${shardId} does not exist`);
+	}
+	await shard.destroy(options);
+}
+
+for (const shardId of data.shardIds) {
+	const shard = new WebSocketShard(new WorkerContextFetchingStrategy(data), shardId);
+	for (const event of Object.values(WebSocketShardEvents)) {
+		// @ts-expect-error
+		shard.on(event, (data) => {
+			const payload: WorkerRecievePayload = {
+				op: WorkerRecievePayloadOp.Event,
+				event,
+				data,
+				shardId,
+			};
+			parentPort!.postMessage(payload);
+		});
+	}
+	shards.set(shardId, shard);
+}
+
+parentPort!
+	.on('messageerror', (err) => {
+		throw err;
+	})
+	// eslint-disable-next-line @typescript-eslint/no-misused-promises
+	.on('message', async (payload: WorkerSendPayload) => {
+		switch (payload.op) {
+			case WorkerSendPayloadOp.Connect: {
+				await connect(payload.shardId);
+				const response: WorkerRecievePayload = {
+					op: WorkerRecievePayloadOp.Connected,
+					shardId: payload.shardId,
+				};
+				parentPort!.postMessage(response);
+				break;
+			}
+
+			case WorkerSendPayloadOp.Destroy: {
+				await destroy(payload.shardId, payload.options);
+				const response: WorkerRecievePayload = {
+					op: WorkerRecievePayloadOp.Destroyed,
+					shardId: payload.shardId,
+				};
+				parentPort!.postMessage(response);
+				break;
+			}
+
+			case WorkerSendPayloadOp.Send: {
+				const shard = shards.get(payload.shardId);
+				if (!shard) {
+					throw new Error(`Shard ${payload.shardId} does not exist`);
+				}
+				await shard.send(payload.payload);
+				break;
+			}
+
+			case WorkerSendPayloadOp.SessionInfoResponse: {
+				break;
+			}
+		}
+	});