diff --git a/module/basicShard.ts b/module/basicShard.ts new file mode 100644 index 000000000..55dfa2f31 --- /dev/null +++ b/module/basicShard.ts @@ -0,0 +1,203 @@ +import { + DiscordBotGatewayData, + GatewayOpcode, + ReadyPayload, +} from "../types/discord.ts"; +import { + eventHandlers, + botGatewayData, + IdentifyPayload, +} from "./client.ts"; +import { delay } from "https://deno.land/std@0.61.0/async/delay.ts"; +import { + connectWebSocket, + isWebSocketCloseEvent, + WebSocket, +} from "https://deno.land/std@0.61.0/ws/mod.ts"; +import { DiscordHeartbeatPayload } from "../types/discord.ts"; +import { logRed } from "../utils/logger.ts"; +import { handleDiscordPayload } from "./shardingManager.ts"; + +const basicShards = new Map(); + +export interface BasicShard { + id: number; + socket: WebSocket; + resumeInterval: number; + sessionID: string; + previousSequenceNumber: number | null; + needToResume: boolean; +} + +export async function createBasicShard( + data: DiscordBotGatewayData, + identifyPayload: IdentifyPayload, + resuming = false, + shardID = 0, +) { + const basicShard: BasicShard = { + id: shardID, + socket: await connectWebSocket(data.url), + resumeInterval: 0, + sessionID: "", + previousSequenceNumber: 0, + needToResume: false, + }; + + basicShards.set(basicShard.id, basicShard); + + if (!resuming) { + // Intial identify with the gateway + await identify(basicShard, identifyPayload); + } else { + await resume(basicShard, identifyPayload); + } + + for await (const message of basicShard.socket) { + if (typeof message === "string") { + const data = JSON.parse(message); + + switch (data.op) { + case GatewayOpcode.Hello: + heartbeat( + basicShard, + identifyPayload, + (data.d as DiscordHeartbeatPayload).heartbeat_interval, + ); + break; + case GatewayOpcode.Reconnect: + case GatewayOpcode.InvalidSession: + // When d is false we need to reidentify + if (!data.d) { + eventHandlers.debug?.( + { type: "invalidSession", data: { shardID: basicShard.id } }, + ); + createBasicShard(botGatewayData, identifyPayload, false, shardID); + break; + } + basicShard.needToResume = true; + resumeConnection(botGatewayData, identifyPayload); + break; + default: + if (data.t === "RESUMED") { + eventHandlers.debug?.( + { type: "resumed", data: { shardID: basicShard.id } }, + ); + + basicShard.needToResume = false; + break; + } + // Important for RESUME + if (data.t === "READY") { + basicShard.sessionID = (data.d as ReadyPayload).session_id; + } + + // Update the sequence number if it is present + if (data.s) basicShard.previousSequenceNumber = data.s; + + handleDiscordPayload(data, basicShard.id); + break; + } + } else if (isWebSocketCloseEvent(message)) { + eventHandlers.debug?.( + { type: "websocketClose", data: { shardID: basicShard.id, message } }, + ); + + // These error codes should just crash the projects + if ([4004, 4005, 4012, 4013, 4014].includes(message.code)) { + logRed(`Close :( ${JSON.stringify(message)}`); + eventHandlers.debug?.( + { + type: "websocketErrored", + data: { shardID: basicShard.id, message }, + }, + ); + + throw new Error( + "Shard.ts: Error occurred that is not resumeable or able to be reconnected.", + ); + } + // These error codes can not be resumed but need to reconnect from start + if ([4003, 4007, 4008, 4009].includes(message.code)) { + eventHandlers.debug?.( + { + type: "websocketReconnecting", + data: { shardID: basicShard.id, message }, + }, + ); + createBasicShard(botGatewayData, identifyPayload, false, shardID); + } else { + basicShard.needToResume = true; + resumeConnection(botGatewayData, identifyPayload); + } + } + } +} + +async function identify(shard: BasicShard, payload: IdentifyPayload) { + await shard.socket.send( + JSON.stringify( + { + op: GatewayOpcode.Identify, + d: { ...payload, shard: [shard.id, payload.shard[1]] }, + }, + ), + ); +} + +async function resume(shard: BasicShard, payload: IdentifyPayload) { + await shard.socket.send(JSON.stringify({ + op: GatewayOpcode.Resume, + d: { + ...payload, + session_id: shard.sessionID, + seq: shard.previousSequenceNumber, + }, + })); +} + +// TODO: If a client does not receive a heartbeat ack between its attempts at sending heartbeats, it should immediately terminate the connection with a non-1000 close code, reconnect, and attempt to resume. +async function heartbeat( + shard: BasicShard, + payload: IdentifyPayload, + interval: number, +) { + await delay(interval); + shard.socket.send( + JSON.stringify( + { op: GatewayOpcode.Heartbeat, d: shard.previousSequenceNumber }, + ), + ); + eventHandlers.debug?.( + { + type: "heartbeat", + data: { + interval, + previousSequenceNumber: shard.previousSequenceNumber, + shardID: shard.id, + }, + }, + ); + + heartbeat(shard, payload, interval); +} + +async function resumeConnection( + botGatewayData: DiscordBotGatewayData, + payload: IdentifyPayload, +) { + const shard = basicShards.get(payload.shard[0]); + if (!shard) { + eventHandlers.debug?.( + { type: "missingShard", data: { shardID: payload.shard[0] } }, + ); + return; + } + + eventHandlers.debug?.({ type: "resuming", data: { shardID: shard.id } }); + // Run it once + createBasicShard(botGatewayData, payload, true, shard.id); + // Then retry every 15 seconds + await delay(1000 * 15); + if (shard.needToResume) resumeConnection(botGatewayData, payload); +} diff --git a/module/client.ts b/module/client.ts index f9de40ae8..46accd5a7 100644 --- a/module/client.ts +++ b/module/client.ts @@ -11,7 +11,7 @@ export let eventHandlers: EventHandlers = {}; export let botGatewayData: DiscordBotGatewayData; -export const identifyPayload = { +export const identifyPayload: IdentifyPayload = { token: "", compress: false, properties: { @@ -23,6 +23,18 @@ export const identifyPayload = { shard: [0, 0], }; +export interface IdentifyPayload { + token: string; + compress: boolean; + properties: { + $os: string; + $browser: string; + $device: string; + }; + intents: number; + shard: [number, number]; +} + export const createClient = async (data: ClientOptions) => { if (data.eventHandlers) eventHandlers = data.eventHandlers; authorization = `Bot ${data.token}`; @@ -37,6 +49,7 @@ export const createClient = async (data: ClientOptions) => { (bits, next) => (bits |= next), 0, ); + identifyPayload.shard = [0, botGatewayData.shards] spawnShards(botGatewayData, identifyPayload); }; diff --git a/module/shardingManager.ts b/module/shardingManager.ts index 6a4afde09..020ce4b76 100644 --- a/module/shardingManager.ts +++ b/module/shardingManager.ts @@ -14,6 +14,7 @@ import { identifyPayload, botID, setBotID, + IdentifyPayload, } from "./client.ts"; import { delay } from "https://deno.land/std@0.61.0/async/delay.ts"; import { @@ -54,6 +55,7 @@ import { } from "../types/message.ts"; import { createMessage } from "../structures/message.ts"; import { GuildUpdateChange } from "../types/options.ts"; +import { createBasicShard } from "./basicShard.ts"; let shardCounter = 0; @@ -104,13 +106,16 @@ export function createShardWorker(shardID?: number) { export const spawnShards = async ( data: DiscordBotGatewayData, - payload: unknown, + payload: IdentifyPayload, id = 1, ) => { if ((data.shards === 1 && id === 1) || id <= data.shards) { if (createNextShard) { createNextShard = false; - createShardWorker(); + if (data.shards >= 25) createShardWorker(); + else { + createBasicShard(data, payload, false, id - 1); + } spawnShards(data, payload, id + 1); } else { await delay(1000); @@ -119,7 +124,10 @@ export const spawnShards = async ( } }; -async function handleDiscordPayload(data: DiscordPayload, shardID: number) { +export async function handleDiscordPayload( + data: DiscordPayload, + shardID: number, +) { eventHandlers.raw?.(data); switch (data.op) { diff --git a/types/options.ts b/types/options.ts index 75c64426b..20c7a75c1 100644 --- a/types/options.ts +++ b/types/options.ts @@ -67,7 +67,8 @@ export interface DebugArg { | "resumed" | "websocketClose" | "websocketErrored" - | "websocketReconnecting"; + | "websocketReconnecting" + | "missingShard"; data: unknown; }