feat: @discordjs/ws (#8260)

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

packages/ws/__tests__/strategy/WorkerContextFetchingStrategy.test.ts

@@ -0,0 +1,83 @@
+import { REST } from '@discordjs/rest';
+import { MockAgent, Interceptable } from 'undici';
+import { beforeEach, test, vi, expect } from 'vitest';
+import {
+	managerToFetchingStrategyOptions,
+	WorkerContextFetchingStrategy,
+	WorkerRecievePayload,
+	WorkerSendPayload,
+	WebSocketManager,
+	WorkerSendPayloadOp,
+	WorkerRecievePayloadOp,
+} from '../../src';
+
+let mockAgent: MockAgent;
+let mockPool: Interceptable;
+
+beforeEach(() => {
+	mockAgent = new MockAgent();
+	mockAgent.disableNetConnect();
+	mockPool = mockAgent.get('https://discord.com');
+});
+
+const session = {
+	shardId: 0,
+	shardCount: 1,
+	sequence: 123,
+	sessionId: 'abc',
+};
+
+vi.mock('node:worker_threads', async () => {
+	const { EventEmitter }: typeof import('node:events') = await vi.importActual('node:events');
+	class MockParentPort extends EventEmitter {
+		public postMessage(message: WorkerRecievePayload) {
+			if (message.op === WorkerRecievePayloadOp.RetrieveSessionInfo) {
+				const response: WorkerSendPayload = {
+					op: WorkerSendPayloadOp.SessionInfoResponse,
+					nonce: message.nonce,
+					session,
+				};
+				this.emit('message', response);
+			}
+		}
+	}
+
+	return {
+		parentPort: new MockParentPort(),
+		isMainThread: false,
+	};
+});
+
+test('session info', async () => {
+	const rest = new REST().setAgent(mockAgent).setToken('A-Very-Fake-Token');
+	const manager = new WebSocketManager({ token: 'A-Very-Fake-Token', intents: 0, rest });
+
+	mockPool
+		.intercept({
+			path: '/api/v10/gateway/bot',
+			method: 'GET',
+		})
+		.reply(() => ({
+			data: {
+				shards: 1,
+				session_start_limit: {
+					max_concurrency: 3,
+					reset_after: 60,
+					remaining: 3,
+					total: 3,
+				},
+				url: 'wss://gateway.discord.gg',
+			},
+			statusCode: 200,
+			responseOptions: {
+				headers: {
+					'content-type': 'application/json',
+				},
+			},
+		}));
+
+	const strategy = new WorkerContextFetchingStrategy(await managerToFetchingStrategyOptions(manager));
+
+	strategy.updateSessionInfo(0, session);
+	expect(await strategy.retrieveSessionInfo(0)).toEqual(session);
+});

packages/ws/__tests__/strategy/WorkerShardingStrategy.test.ts

@@ -0,0 +1,195 @@
+import { REST } from '@discordjs/rest';
+import {
+	GatewayDispatchEvents,
+	GatewayDispatchPayload,
+	GatewayOpcodes,
+	GatewaySendPayload,
+} from 'discord-api-types/v10';
+import { MockAgent, Interceptable } from 'undici';
+import { beforeEach, test, vi, expect, afterEach } from 'vitest';
+import {
+	WorkerRecievePayload,
+	WorkerSendPayload,
+	WebSocketManager,
+	WorkerSendPayloadOp,
+	WorkerRecievePayloadOp,
+	WorkerShardingStrategy,
+	WebSocketShardEvents,
+	SessionInfo,
+} from '../../src';
+
+let mockAgent: MockAgent;
+let mockPool: Interceptable;
+
+const mockConstructor = vi.fn();
+const mockSend = vi.fn();
+const mockTerminate = vi.fn();
+
+const memberChunkData: GatewayDispatchPayload = {
+	op: GatewayOpcodes.Dispatch,
+	s: 123,
+	t: GatewayDispatchEvents.GuildMembersChunk,
+	d: {
+		guild_id: '123',
+		members: [],
+	},
+};
+
+const sessionInfo: SessionInfo = {
+	shardId: 0,
+	shardCount: 2,
+	sequence: 123,
+	sessionId: 'abc',
+};
+
+vi.mock('node:worker_threads', async () => {
+	const { EventEmitter }: typeof import('node:events') = await vi.importActual('node:events');
+	class MockWorker extends EventEmitter {
+		public constructor(...args: any[]) {
+			super();
+			// eslint-disable-next-line @typescript-eslint/no-unsafe-argument
+			mockConstructor(...args);
+			// need to delay this by an event loop cycle to allow the strategy to attach a listener
+			setImmediate(() => this.emit('online'));
+		}
+
+		public postMessage(message: WorkerSendPayload) {
+			switch (message.op) {
+				case WorkerSendPayloadOp.Connect: {
+					const response: WorkerRecievePayload = {
+						op: WorkerRecievePayloadOp.Connected,
+						shardId: message.shardId,
+					};
+					this.emit('message', response);
+					break;
+				}
+
+				case WorkerSendPayloadOp.Destroy: {
+					const response: WorkerRecievePayload = {
+						op: WorkerRecievePayloadOp.Destroyed,
+						shardId: message.shardId,
+					};
+					this.emit('message', response);
+					break;
+				}
+
+				case WorkerSendPayloadOp.Send: {
+					if (message.payload.op === GatewayOpcodes.RequestGuildMembers) {
+						const response: WorkerRecievePayload = {
+							op: WorkerRecievePayloadOp.Event,
+							shardId: message.shardId,
+							event: WebSocketShardEvents.Dispatch,
+							data: memberChunkData,
+						};
+						this.emit('message', response);
+
+						// Fetch session info
+						const sessionFetch: WorkerRecievePayload = {
+							op: WorkerRecievePayloadOp.RetrieveSessionInfo,
+							shardId: message.shardId,
+							nonce: Math.random(),
+						};
+						this.emit('message', sessionFetch);
+					}
+
+					mockSend(message.shardId, message.payload);
+					break;
+				}
+
+				case WorkerSendPayloadOp.SessionInfoResponse: {
+					message.session ??= sessionInfo;
+
+					const session: WorkerRecievePayload = {
+						op: WorkerRecievePayloadOp.UpdateSessionInfo,
+						shardId: message.session.shardId,
+						session: { ...message.session, sequence: message.session.sequence + 1 },
+					};
+					this.emit('message', session);
+					break;
+				}
+			}
+		}
+
+		public terminate = mockTerminate;
+	}
+
+	return {
+		Worker: MockWorker,
+	};
+});
+
+beforeEach(() => {
+	mockAgent = new MockAgent();
+	mockAgent.disableNetConnect();
+	mockPool = mockAgent.get('https://discord.com');
+});
+
+afterEach(() => {
+	mockConstructor.mockRestore();
+	mockSend.mockRestore();
+	mockTerminate.mockRestore();
+});
+
+test('spawn, connect, send a message, session info, and destroy', async () => {
+	const rest = new REST().setAgent(mockAgent).setToken('A-Very-Fake-Token');
+
+	const mockRetrieveSessionInfo = vi.fn();
+	const mockUpdateSessionInfo = vi.fn();
+	const manager = new WebSocketManager({
+		token: 'A-Very-Fake-Token',
+		intents: 0,
+		rest,
+		shardIds: [0, 1],
+		retrieveSessionInfo: mockRetrieveSessionInfo,
+		updateSessionInfo: mockUpdateSessionInfo,
+	});
+
+	const managerEmitSpy = vi.spyOn(manager, 'emit');
+
+	mockPool
+		.intercept({
+			path: '/api/v10/gateway/bot',
+			method: 'GET',
+		})
+		.reply(() => ({
+			data: {
+				shards: 1,
+				session_start_limit: {
+					max_concurrency: 3,
+					reset_after: 60,
+					remaining: 3,
+					total: 3,
+				},
+				url: 'wss://gateway.discord.gg',
+			},
+			statusCode: 200,
+			responseOptions: {
+				headers: {
+					'content-type': 'application/json',
+				},
+			},
+		}));
+
+	const strategy = new WorkerShardingStrategy(manager, { shardsPerWorker: 'all' });
+	manager.setStrategy(strategy);
+
+	await manager.connect();
+	expect(mockConstructor).toHaveBeenCalledWith(
+		expect.stringContaining('worker.cjs'),
+		// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+		expect.objectContaining({ workerData: expect.objectContaining({ shardIds: [0, 1] }) }),
+	);
+
+	const payload: GatewaySendPayload = { op: GatewayOpcodes.RequestGuildMembers, d: { guild_id: '123', limit: 0 } };
+	await manager.send(0, payload);
+	expect(mockSend).toHaveBeenCalledWith(0, payload);
+	expect(managerEmitSpy).toHaveBeenCalledWith(WebSocketShardEvents.Dispatch, {
+		...memberChunkData,
+		shardId: 0,
+	});
+	expect(mockRetrieveSessionInfo).toHaveBeenCalledWith(0);
+	expect(mockUpdateSessionInfo).toHaveBeenCalledWith(0, { ...sessionInfo, sequence: sessionInfo.sequence + 1 });
+
+	await manager.destroy({ reason: 'souji is a soft boi :3' });
+	expect(mockTerminate).toHaveBeenCalled();
+});

packages/ws/__tests__/util/IdentifyThrottler.test.ts

@@ -0,0 +1,46 @@
+import { setTimeout as sleep } from 'node:timers/promises';
+import { expect, Mock, test, vi } from 'vitest';
+import { IdentifyThrottler, WebSocketManager } from '../../src';
+
+vi.mock('node:timers/promises', () => ({
+	setTimeout: vi.fn(),
+}));
+
+const fetchGatewayInformation = vi.fn();
+
+const manager = {
+	fetchGatewayInformation,
+} as unknown as WebSocketManager;
+
+const throttler = new IdentifyThrottler(manager);
+
+vi.useFakeTimers();
+
+const NOW = vi.fn().mockReturnValue(Date.now());
+global.Date.now = NOW;
+
+test('wait for identify', async () => {
+	fetchGatewayInformation.mockReturnValue({
+		session_start_limit: {
+			max_concurrency: 2,
+		},
+	});
+
+	// First call should never wait
+	await throttler.waitForIdentify();
+	expect(sleep).not.toHaveBeenCalled();
+
+	// Second call still won't wait because max_concurrency is 2
+	await throttler.waitForIdentify();
+	expect(sleep).not.toHaveBeenCalled();
+
+	// Third call should wait
+	await throttler.waitForIdentify();
+	expect(sleep).toHaveBeenCalled();
+
+	(sleep as Mock).mockRestore();
+
+	// Fourth call shouldn't wait, because our max_concurrency is 2 and we waited for a reset
+	await throttler.waitForIdentify();
+	expect(sleep).not.toHaveBeenCalled();
+});

packages/ws/__tests__/ws/WebSocketManager.test.ts

@@ -0,0 +1,197 @@
+import { REST } from '@discordjs/rest';
+import { APIGatewayBotInfo, GatewayOpcodes, GatewaySendPayload } from 'discord-api-types/v10';
+import { MockAgent, Interceptable } from 'undici';
+import { beforeEach, describe, expect, test, vi } from 'vitest';
+import { IShardingStrategy, WebSocketManager } from '../../src';
+
+vi.useFakeTimers();
+
+let mockAgent: MockAgent;
+let mockPool: Interceptable;
+
+beforeEach(() => {
+	mockAgent = new MockAgent();
+	mockAgent.disableNetConnect();
+	mockPool = mockAgent.get('https://discord.com');
+});
+
+const NOW = vi.fn().mockReturnValue(Date.now());
+global.Date.now = NOW;
+
+test('fetch gateway information', async () => {
+	const rest = new REST().setAgent(mockAgent).setToken('A-Very-Fake-Token');
+	const manager = new WebSocketManager({ token: 'A-Very-Fake-Token', intents: 0, rest });
+
+	const data: APIGatewayBotInfo = {
+		shards: 1,
+		session_start_limit: {
+			max_concurrency: 3,
+			reset_after: 60,
+			remaining: 3,
+			total: 3,
+		},
+		url: 'wss://gateway.discord.gg',
+	};
+
+	const fetch = vi.fn(() => ({
+		data,
+		statusCode: 200,
+		responseOptions: {
+			headers: {
+				'content-type': 'application/json',
+			},
+		},
+	}));
+
+	mockPool
+		.intercept({
+			path: '/api/v10/gateway/bot',
+			method: 'GET',
+		})
+		.reply(fetch);
+
+	const initial = await manager.fetchGatewayInformation();
+	expect(initial).toEqual(data);
+	expect(fetch).toHaveBeenCalledOnce();
+
+	fetch.mockRestore();
+
+	const cached = await manager.fetchGatewayInformation();
+	expect(cached).toEqual(data);
+	expect(fetch).not.toHaveBeenCalled();
+
+	fetch.mockRestore();
+	mockPool
+		.intercept({
+			path: '/api/v10/gateway/bot',
+			method: 'GET',
+		})
+		.reply(fetch);
+
+	const forced = await manager.fetchGatewayInformation(true);
+	expect(forced).toEqual(data);
+	expect(fetch).toHaveBeenCalledOnce();
+
+	fetch.mockRestore();
+	mockPool
+		.intercept({
+			path: '/api/v10/gateway/bot',
+			method: 'GET',
+		})
+		.reply(fetch);
+
+	NOW.mockReturnValue(Infinity);
+	const cacheExpired = await manager.fetchGatewayInformation();
+	expect(cacheExpired).toEqual(data);
+	expect(fetch).toHaveBeenCalledOnce();
+});
+
+describe('get shard count', () => {
+	test('with shard count', async () => {
+		const rest = new REST().setAgent(mockAgent).setToken('A-Very-Fake-Token');
+		const manager = new WebSocketManager({ token: 'A-Very-Fake-Token', intents: 0, rest, shardCount: 2 });
+
+		expect(await manager.getShardCount()).toBe(2);
+	});
+
+	test('with shard ids array', async () => {
+		const rest = new REST().setAgent(mockAgent).setToken('A-Very-Fake-Token');
+		const shardIds = [5, 9];
+		const manager = new WebSocketManager({ token: 'A-Very-Fake-Token', intents: 0, rest, shardIds });
+
+		expect(await manager.getShardCount()).toBe(shardIds.at(-1)! + 1);
+	});
+
+	test('with shard id range', async () => {
+		const rest = new REST().setAgent(mockAgent).setToken('A-Very-Fake-Token');
+		const shardIds = { start: 5, end: 9 };
+		const manager = new WebSocketManager({ token: 'A-Very-Fake-Token', intents: 0, rest, shardIds });
+
+		expect(await manager.getShardCount()).toBe(shardIds.end + 1);
+	});
+});
+
+test('update shard count', async () => {
+	const rest = new REST().setAgent(mockAgent).setToken('A-Very-Fake-Token');
+	const manager = new WebSocketManager({ token: 'A-Very-Fake-Token', intents: 0, rest, shardCount: 2 });
+
+	const data: APIGatewayBotInfo = {
+		shards: 1,
+		session_start_limit: {
+			max_concurrency: 3,
+			reset_after: 60,
+			remaining: 3,
+			total: 3,
+		},
+		url: 'wss://gateway.discord.gg',
+	};
+
+	const fetch = vi.fn(() => ({
+		data,
+		statusCode: 200,
+		responseOptions: {
+			headers: {
+				'content-type': 'application/json',
+			},
+		},
+	}));
+
+	mockPool
+		.intercept({
+			path: '/api/v10/gateway/bot',
+			method: 'GET',
+		})
+		.reply(fetch);
+
+	expect(await manager.getShardCount()).toBe(2);
+	expect(fetch).not.toHaveBeenCalled();
+
+	fetch.mockRestore();
+	mockPool
+		.intercept({
+			path: '/api/v10/gateway/bot',
+			method: 'GET',
+		})
+		.reply(fetch);
+
+	await manager.updateShardCount(3);
+	expect(await manager.getShardCount()).toBe(3);
+	expect(fetch).toHaveBeenCalled();
+});
+
+test('it handles passing in both shardIds and shardCount', async () => {
+	const rest = new REST().setAgent(mockAgent).setToken('A-Very-Fake-Token');
+	const shardIds = { start: 2, end: 3 };
+	const manager = new WebSocketManager({ token: 'A-Very-Fake-Token', intents: 0, rest, shardIds, shardCount: 4 });
+
+	expect(await manager.getShardCount()).toBe(4);
+	expect(await manager.getShardIds()).toStrictEqual([2, 3]);
+});
+
+test('strategies', async () => {
+	class MockStrategy implements IShardingStrategy {
+		public spawn = vi.fn();
+		public connect = vi.fn();
+		public destroy = vi.fn();
+		public send = vi.fn();
+	}
+
+	const rest = new REST().setAgent(mockAgent).setToken('A-Very-Fake-Token');
+	const shardIds = [0, 1, 2];
+	const manager = new WebSocketManager({ token: 'A-Very-Fake-Token', intents: 0, rest, shardIds });
+
+	const strategy = new MockStrategy();
+	manager.setStrategy(strategy);
+
+	await manager.connect();
+	expect(strategy.spawn).toHaveBeenCalledWith(shardIds);
+	expect(strategy.connect).toHaveBeenCalled();
+
+	const destroyOptions = { reason: ':3' };
+	await manager.destroy(destroyOptions);
+	expect(strategy.destroy).toHaveBeenCalledWith(destroyOptions);
+
+	const send: GatewaySendPayload = { op: GatewayOpcodes.RequestGuildMembers, d: { guild_id: '1234', limit: 0 } };
+	await manager.send(0, send);
+	expect(strategy.send).toHaveBeenCalledWith(0, send);
+});