From b469c9852733f6ed6e665d25c33770105adb06f4 Mon Sep 17 00:00:00 2001 From: ITOH Date: Wed, 25 May 2022 22:27:49 +0200 Subject: [PATCH] refactor(gateway,types,util)!: finalize gateway code (#2241) * push some cursed stuff * bucket * fix bugs and hack start ordered * some management improvements * more gw stuff * f * rename manager to gateway * remove basic thing * remove old stuff * f * fix imp * fixes --- bot.ts | 45 +-- gateway/calculateMaxShards.ts | 9 - gateway/closeWs.ts | 6 - gateway/createShard.ts | 58 --- gateway/gatewayManager.ts | 245 ------------ gateway/handleOnMessage.ts | 182 --------- gateway/heartbeat.ts | 61 --- gateway/identify.ts | 71 ---- gateway/manager/calculateTotalShards.ts | 16 + gateway/manager/calculateWorkerId.ts | 13 + gateway/manager/gatewayManager.ts | 296 ++++++++++++++ gateway/manager/mod.ts | 8 + gateway/manager/prepareBuckets.ts | 47 +++ gateway/manager/resharder.ts | 341 +++++++++++++++++ gateway/manager/shardManager.ts | 122 ++++++ gateway/manager/spawnShards.ts | 32 ++ gateway/manager/stop.ts | 8 + gateway/{ => manager}/tellWorkerToIdentify.ts | 6 +- gateway/mod.ts | 18 +- gateway/processGatewayQueue.ts | 51 --- gateway/resharder.ts | 155 -------- gateway/resume.ts | 67 ---- gateway/safeRequestsPerShard.ts | 9 - gateway/sendShardMessage.ts | 17 - gateway/shard.ts | 48 --- gateway/shard/calculateSafeRequests.ts | 9 + gateway/shard/close.ts | 7 + gateway/shard/connect.ts | 34 ++ gateway/shard/createShard.ts | 362 ++++++++++++++++++ gateway/{ => shard}/deps.ts | 0 gateway/shard/handleClose.ts | 68 ++++ gateway/shard/handleMessage.ts | 156 ++++++++ gateway/shard/identify.ts | 50 +++ gateway/shard/isOpen.ts | 5 + gateway/shard/mod.ts | 14 + gateway/shard/resume.ts | 48 +++ gateway/shard/send.ts | 27 ++ gateway/shard/shutdown.ts | 6 + gateway/shard/startHeartbeating.ts | 64 ++++ gateway/shard/stopHeartbeating.ts | 9 + gateway/shard/types.ts | 152 ++++++++ gateway/spawnShards.ts | 64 ---- gateway/stopGateway.ts | 21 - helpers/members/fetchMembers.ts | 7 +- helpers/misc/editBotStatus.ts | 79 +--- helpers/misc/editShardStatus.ts | 78 ++++ helpers/misc/mod.ts | 1 + helpers/voice/connectToVoiceChannel.ts | 8 +- types/shared.ts | 2 + util/bucket.ts | 175 +++++++++ util/calculateShardId.ts | 6 +- 51 files changed, 2194 insertions(+), 1189 deletions(-) delete mode 100644 gateway/calculateMaxShards.ts delete mode 100644 gateway/closeWs.ts delete mode 100644 gateway/createShard.ts delete mode 100644 gateway/gatewayManager.ts delete mode 100644 gateway/handleOnMessage.ts delete mode 100644 gateway/heartbeat.ts delete mode 100644 gateway/identify.ts create mode 100644 gateway/manager/calculateTotalShards.ts create mode 100644 gateway/manager/calculateWorkerId.ts create mode 100644 gateway/manager/gatewayManager.ts create mode 100644 gateway/manager/mod.ts create mode 100644 gateway/manager/prepareBuckets.ts create mode 100644 gateway/manager/resharder.ts create mode 100644 gateway/manager/shardManager.ts create mode 100644 gateway/manager/spawnShards.ts create mode 100644 gateway/manager/stop.ts rename gateway/{ => manager}/tellWorkerToIdentify.ts (67%) delete mode 100644 gateway/processGatewayQueue.ts delete mode 100644 gateway/resharder.ts delete mode 100644 gateway/resume.ts delete mode 100644 gateway/safeRequestsPerShard.ts delete mode 100644 gateway/sendShardMessage.ts delete mode 100644 gateway/shard.ts create mode 100644 gateway/shard/calculateSafeRequests.ts create mode 100644 gateway/shard/close.ts create mode 100644 gateway/shard/connect.ts create mode 100644 gateway/shard/createShard.ts rename gateway/{ => shard}/deps.ts (100%) create mode 100644 gateway/shard/handleClose.ts create mode 100644 gateway/shard/handleMessage.ts create mode 100644 gateway/shard/identify.ts create mode 100644 gateway/shard/isOpen.ts create mode 100644 gateway/shard/mod.ts create mode 100644 gateway/shard/resume.ts create mode 100644 gateway/shard/send.ts create mode 100644 gateway/shard/shutdown.ts create mode 100644 gateway/shard/startHeartbeating.ts create mode 100644 gateway/shard/stopHeartbeating.ts create mode 100644 gateway/shard/types.ts delete mode 100644 gateway/spawnShards.ts delete mode 100644 gateway/stopGateway.ts create mode 100644 helpers/misc/editShardStatus.ts create mode 100644 util/bucket.ts diff --git a/bot.ts b/bot.ts index 18b9f90c3..43fd20b85 100644 --- a/bot.ts +++ b/bot.ts @@ -30,7 +30,7 @@ import { SLASH_COMMANDS_NAME_REGEX, USER_AGENT, } from "./util/constants.ts"; -import { createGatewayManager, GatewayManager } from "./gateway/mod.ts"; +import { createGatewayManager, GatewayManager } from "./gateway/manager/gatewayManager.ts"; import { validateLength } from "./util/validateLength.ts"; import { delay, formatImageURL, hasProperty } from "./util/utils.ts"; import { iconBigintToHash, iconHashToBigInt } from "./util/hash.ts"; @@ -138,8 +138,9 @@ import { import { transformEmbedToDiscordEmbed } from "./transformers/reverse/embed.ts"; import { transformComponentToDiscordComponent } from "./transformers/reverse/component.ts"; import { getBotIdFromToken, removeTokenPrefix } from "./util/token.ts"; +import { CreateShardManager } from "./gateway/manager/shardManager.ts"; -export function createBot(options: CreateBotOptions): Bot { +export async function createBot(options: CreateBotOptions): Promise { const bot = { id: options.botId ?? getBotIdFromToken(options.token), applicationId: options.applicationId || options.botId, @@ -167,22 +168,26 @@ export function createBot(options: CreateBotOptions): Bot { bot.helpers = createHelpers(bot, options.helpers ?? {}); bot.gateway = createGatewayManager({ - token: bot.token, - intents: bot.intents, + gatewayBot: bot.botGatewayData ?? await bot.helpers.getGatewayBot(), + gatewayConfig: { + token: options.token, + }, + debug: bot.events.debug, + handleDiscordPayload: bot.handleDiscordPayload ?? - async function (_, data: DiscordGatewayPayload, shardId: number) { + async function (shard, data: DiscordGatewayPayload) { // TRIGGER RAW EVENT - bot.events.raw(bot as Bot, data, shardId); + bot.events.raw(bot as Bot, data, shard.id); if (!data.t) return; // RUN DISPATCH CHECK - await bot.events.dispatchRequirements(bot as Bot, data, shardId); + await bot.events.dispatchRequirements(bot as Bot, data, shard.id); bot.handlers[data.t as GatewayDispatchEventNames]?.( bot as Bot, data, - shardId, + shard.id, ); }, }); @@ -253,22 +258,8 @@ export function createEventHandlers( }; } -export async function startBot(bot: Bot) { - if (!bot.botGatewayData) { - bot.botGatewayData = await bot.helpers.getGatewayBot(); - } - - // SETUP GATEWAY LOGIN INFO - bot.gateway.urlWSS = bot.botGatewayData.url; - bot.gateway.shardsRecommended = bot.botGatewayData.shards; - bot.gateway.sessionStartLimitTotal = bot.botGatewayData.sessionStartLimit.total; - bot.gateway.sessionStartLimitRemaining = bot.botGatewayData.sessionStartLimit.remaining; - bot.gateway.sessionStartLimitResetAfter = bot.botGatewayData.sessionStartLimit.resetAfter; - bot.gateway.maxConcurrency = bot.botGatewayData.sessionStartLimit.maxConcurrency; - bot.gateway.lastShardId = bot.botGatewayData.shards === 1 ? 0 : bot.botGatewayData.shards - 1; - bot.gateway.maxShards = bot.botGatewayData.shards; - - bot.gateway.spawnShards(bot.gateway); +export function startBot(bot: Bot) { + bot.gateway.spawnShards(); } export function createUtils(options: Partial) { @@ -304,7 +295,7 @@ export interface HelperUtils { } export async function stopBot(bot: Bot) { - await bot.gateway.stopGateway(bot.gateway); + await bot.gateway.stop(1000, "User requested bot stop"); return bot; } @@ -318,7 +309,7 @@ export interface CreateBotOptions { intents?: GatewayIntents; botGatewayData?: GetGatewayBot; rest?: Omit; - handleDiscordPayload?: GatewayManager["handleDiscordPayload"]; + handleDiscordPayload?: CreateShardManager["handleMessage"]; utils?: Partial>; transformers?: Partial>; helpers?: Partial; @@ -348,7 +339,7 @@ export interface Bot { fetchAllMembersProcessingRequests: Map; }; enabledPlugins: Set; - handleDiscordPayload?: GatewayManager["handleDiscordPayload"]; + handleDiscordPayload?: CreateShardManager["handleMessage"]; } export const defaultHelpers = { ...helpers }; diff --git a/gateway/calculateMaxShards.ts b/gateway/calculateMaxShards.ts deleted file mode 100644 index 34b2664ec..000000000 --- a/gateway/calculateMaxShards.ts +++ /dev/null @@ -1,9 +0,0 @@ -/** Handler used to determine max number of shards to use based upon the max concurrency. */ -export function calculateMaxShards(maxShards: number, maxConcurrency: number): number { - if (maxShards < 100) return maxShards; - - return Math.ceil( - maxShards / - (maxConcurrency === 1 ? 16 : maxConcurrency), - ) * maxConcurrency; -} diff --git a/gateway/closeWs.ts b/gateway/closeWs.ts deleted file mode 100644 index 399174439..000000000 --- a/gateway/closeWs.ts +++ /dev/null @@ -1,6 +0,0 @@ -/** Use this function to close a ws connection properly */ -export function closeWS(ws: WebSocket, code?: number, reason?: string) { - if (ws.readyState !== WebSocket.OPEN) return; - - ws.close(code, reason); -} diff --git a/gateway/createShard.ts b/gateway/createShard.ts deleted file mode 100644 index b120ef805..000000000 --- a/gateway/createShard.ts +++ /dev/null @@ -1,58 +0,0 @@ -import { GatewayCloseEventCodes } from "../types/shared.ts"; -import { GatewayManager } from "./gatewayManager.ts"; - -export function createShard(gateway: GatewayManager, shardId: number) { - const socket = new WebSocket(`${gateway.urlWSS}/?v=9&encoding=json`); - - socket.onerror = (errorEvent) => { - gateway.debug("GW ERROR", { shardId, error: errorEvent }); - }; - - socket.onmessage = ({ data: message }) => gateway.handleOnMessage(gateway, message, shardId); - - socket.onclose = async (event) => { - gateway.debug("GW CLOSED", { shardId, payload: event }); - - if (event.code === 3064 || event.reason === "Discordeno Testing Finished! Do Not RESUME!") { - return; - } - - if (event.code === 3065 || ["Resharded!", "Resuming the shard, closing old shard."].includes(event.reason)) { - return gateway.debug("GW CLOSED_RECONNECT", { shardId, payload: event }); - } - - switch (event.code) { - // Discordeno tests finished - case 3061: - return; - case 3063: // Resharded - case 3064: // Resuming - case 3065: // Re-identifying - case 3066: // Missing ACK - // Will restart shard manually - return gateway.debug("GW CLOSED_RECONNECT", { shardId, payload: event }); - case GatewayCloseEventCodes.UnknownOpcode: - case GatewayCloseEventCodes.DecodeError: - case GatewayCloseEventCodes.AuthenticationFailed: - case GatewayCloseEventCodes.AlreadyAuthenticated: - case GatewayCloseEventCodes.InvalidShard: - case GatewayCloseEventCodes.ShardingRequired: - case GatewayCloseEventCodes.InvalidApiVersion: - case GatewayCloseEventCodes.InvalidIntents: - case GatewayCloseEventCodes.DisallowedIntents: - throw new Error(event.reason || "Discord gave no reason! GG! You broke Discord!"); - // THESE ERRORS CAN NO BE RESUMED! THEY MUST RE-IDENTIFY! - case GatewayCloseEventCodes.NotAuthenticated: - case GatewayCloseEventCodes.InvalidSeq: - case GatewayCloseEventCodes.RateLimited: - case GatewayCloseEventCodes.SessionTimedOut: - await gateway.identify(gateway, shardId, gateway.maxShards); - break; - default: - gateway.resume(gateway, shardId); - break; - } - }; - - return socket; -} diff --git a/gateway/gatewayManager.ts b/gateway/gatewayManager.ts deleted file mode 100644 index 0d3748425..000000000 --- a/gateway/gatewayManager.ts +++ /dev/null @@ -1,245 +0,0 @@ -import { Collection } from "../util/collection.ts"; -import { safeRequestsPerShard } from "./safeRequestsPerShard.ts"; -import { closeWS } from "./closeWs.ts"; -import { createShard } from "./createShard.ts"; -import { handleOnMessage } from "./handleOnMessage.ts"; -import { heartbeat } from "./heartbeat.ts"; -import { identify } from "./identify.ts"; -import { processGatewayQueue } from "./processGatewayQueue.ts"; -import { - markNewGuildShardId, - resharder, - resharderCloseOldShards, - resharderIsPending, - reshardingEditGuildShardIds, - startReshardingChecks, -} from "./resharder.ts"; -import { resume } from "./resume.ts"; -import { sendShardMessage } from "./sendShardMessage.ts"; -import { prepareBuckets, spawnShards } from "./spawnShards.ts"; -import { stopGateway } from "./stopGateway.ts"; -import { tellWorkerToIdentify } from "./tellWorkerToIdentify.ts"; -import { DiscordenoShard } from "./shard.ts"; -import { GatewayIntents } from "../types/shared.ts"; -import { StatusUpdate } from "../helpers/misc/editBotStatus.ts"; -import { DiscordGatewayPayload } from "../types/discord.ts"; -import { calculateMaxShards } from "./calculateMaxShards.ts"; -import { removeTokenPrefix } from "../util/token.ts"; - -/** Create a new Gateway Manager. - * - * @param options: Customize every bit of the manager. If something is not - * provided, it will fallback to a default which should be suitable for most - * bots. - */ -export function createGatewayManager( - options: Partial & Pick, -): GatewayManager { - return { - queueResetInterval: 60000, - maxRequestsPerInterval: 120, - cache: { - guildIds: new Set(), - loadingGuildIds: new Set(), - editedMessages: new Collection(), - }, - secretKey: options.secretKey ?? "", - url: options.url ?? "", - reshard: options.reshard ?? true, - reshardPercentage: options.reshardPercentage ?? 80, - spawnShardDelay: options.spawnShardDelay ?? 5000, - maxShards: options.maxShards ?? options.shardsRecommended ?? 0, - useOptimalLargeBotSharding: options.useOptimalLargeBotSharding ?? true, - shardsPerWorker: options.shardsPerWorker ?? 25, - maxWorkers: options.maxWorkers ?? 4, - firstShardId: options.firstShardId ?? 0, - lastShardId: options.lastShardId ?? options.maxShards ?? options.shardsRecommended ?? 1, - token: removeTokenPrefix(options.token, "GATEWAY"), - compress: options.compress ?? false, - $os: options.$os ?? "linux", - $browser: options.$browser ?? "Discordeno", - $device: options.$device ?? "Discordeno", - intents: options.intents ?? 0, - shard: options.shard ?? [0, options.shardsRecommended ?? 1], - presence: options.presence, - urlWSS: options.urlWSS ?? "wss://gateway.discord.gg/?v=9&encoding=json", - shardsRecommended: options.shardsRecommended ?? 1, - sessionStartLimitTotal: options.sessionStartLimitTotal ?? 1000, - sessionStartLimitRemaining: options.sessionStartLimitRemaining ?? 1000, - sessionStartLimitResetAfter: options.sessionStartLimitResetAfter ?? 0, - maxConcurrency: options.maxConcurrency ?? 1, - shards: options.shards ?? new Collection(), - loadingShards: options.loadingShards ?? new Collection(), - buckets: new Collection(), - utf8decoder: new TextDecoder(), - - prepareBuckets: options.prepareBuckets ?? prepareBuckets, - spawnShards: options.spawnShards ?? spawnShards, - createShard: options.createShard ?? createShard, - identify: options.identify ?? identify, - heartbeat: options.heartbeat ?? heartbeat, - tellWorkerToIdentify, - debug: options.debug || function () {}, - resharding: { - resharder: options.resharding?.resharder ?? resharder, - isPending: options.resharding?.isPending ?? resharderIsPending, - closeOldShards: options.resharding?.closeOldShards ?? resharderCloseOldShards, - check: options.resharding?.check ?? startReshardingChecks, - markNewGuildShardId: options.resharding?.markNewGuildShardId ?? markNewGuildShardId, - editGuildShardIds: options.resharding?.editGuildShardIds ?? reshardingEditGuildShardIds, - }, - handleOnMessage: options.handleOnMessage ?? handleOnMessage, - processGatewayQueue: options.processGatewayQueue ?? processGatewayQueue, - closeWS: options.closeWS ?? closeWS, - stopGateway: options.stopGateway ?? stopGateway, - sendShardMessage: options.sendShardMessage ?? sendShardMessage, - resume: options.resume ?? resume, - safeRequestsPerShard: options.safeRequestsPerShard ?? safeRequestsPerShard, - handleDiscordPayload: options.handleDiscordPayload, - calculateMaxShards: options.calculateMaxShards ?? calculateMaxShards, - }; -} - -export interface GatewayManager { - /** The secret key authorization header the bot will expect when sending payloads. */ - secretKey: string; - /** The url that all discord payloads for the dispatch type should be sent to. */ - url: string; - /** Whether or not to automatically reshard. */ - reshard: boolean; - /** The percentage at which resharding should occur. */ - reshardPercentage: number; - /** The delay in milliseconds to wait before spawning next shard. OPTIMAL IS ABOVE 2500. YOU DON"T WANT TO HIT THE RATE LIMIT!!! */ - spawnShardDelay: number; - /** The maximum shard Id number. Useful for zero-downtime updates or resharding. */ - maxShards: number; - /** Whether or not the resharder should automatically switch to LARGE BOT SHARDING when you are above 100K servers. */ - useOptimalLargeBotSharding: boolean; - /** The amount of shards to load per worker. */ - shardsPerWorker: number; - /** The maximum amount of workers to use for your bot. */ - maxWorkers: number; - /** The first shard Id to start spawning. */ - firstShardId: number; - /** The last shard Id for this worker. */ - lastShardId: number; - token: string; - compress: boolean; - $os: string; - $browser: string; - $device: string; - intents: GatewayIntents; - shard: [number, number]; - presence?: Omit; - - /** The WSS URL that can be used for connecting to the gateway. */ - urlWSS: string; - /** The recommended number of shards to use when connecting. */ - shardsRecommended: number; - /** The total number of session starts the current user is allowed. */ - sessionStartLimitTotal: number; - /** The remaining number of session starts the current user is allowed. */ - sessionStartLimitRemaining: number; - /** Milliseconds left until limit is reset. */ - sessionStartLimitResetAfter: number; - /** The number of identify requests allowed per 5 seconds. - * So, if you had a max concurrency of 16, and 16 shards for example, you could start them all up at the same time. - * Whereas if you had 32 shards, if you tried to start up shard 0 and 16 at the same time for example, it would not work. You can start shards 0-15 concurrently, then 16-31... - */ - maxConcurrency: number; - shards: Collection; - loadingShards: Collection< - number, - { - shardId: number; - resolve: (value: unknown) => void; - } - >; - /** Stored as bucketId: { workers: [workerId, [ShardIds]], createNextShard: boolean } */ - buckets: Collection< - number, - { - workers: number[][]; - createNextShard: (() => Promise)[]; - } - >; - utf8decoder: TextDecoder; - /** The amount of milliseconds the gateway rate limit will reset in. By default 60000 or 1 minute. */ - queueResetInterval: number; - /** The maximum amount of requests that the gateway can make before being rate limited. By default 120. */ - maxRequestsPerInterval: number; - - cache: { - guildIds: Set; - loadingGuildIds: Set; - editedMessages: Collection; - }; - - // METHODS - - /** Prepares the buckets for identifying */ - prepareBuckets: typeof prepareBuckets; - /** The handler for spawning ALL the shards. */ - spawnShards: typeof spawnShards; - /** Create the websocket and adds the proper handlers to the websocket. */ - createShard: typeof createShard; - /** Begins identification of the shard to discord. */ - identify: typeof identify; - /** Begins heartbeating of the shard to keep it alive. */ - heartbeat: typeof heartbeat; - /** Sends the discord payload to another server. */ - handleDiscordPayload: (gateway: GatewayManager, data: DiscordGatewayPayload, shardId: number) => any; - /** Tell the worker to begin identifying this shard */ - tellWorkerToIdentify: typeof tellWorkerToIdentify; - /** Handle the different logs. Used for debugging. */ - debug: (text: GatewayDebugEvents, ...args: any[]) => unknown; - /** The methods related to resharding. */ - resharding: { - /** Handles resharding the bot when necessary. */ - resharder: typeof resharder; - /** Handles checking if all new shards are online in the new gateway. */ - isPending: typeof resharderIsPending; - /** Handles closing all shards in the old gateway. */ - closeOldShards: typeof resharderCloseOldShards; - /** Handles checking if it is time to reshard and triggers the resharder. */ - check: typeof startReshardingChecks; - /** Handler to mark a guild id with its new shard id in cache. */ - markNewGuildShardId: typeof markNewGuildShardId; - /** Handler to update all guilds in cache with the new shard id. */ - editGuildShardIds: typeof reshardingEditGuildShardIds; - }; - /** Handles the message events from websocket. */ - handleOnMessage: typeof handleOnMessage; - /** Handles processing queue of requests send to this shard. */ - processGatewayQueue: typeof processGatewayQueue; - /** Closes shard WebSocket connection properly. */ - closeWS: typeof closeWS; - /** Use this function to stop the gateway properly. */ - stopGateway: typeof stopGateway; - /** Properly adds a message to the shards queue. */ - sendShardMessage: typeof sendShardMessage; - /** Properly resume an old shards session. */ - resume: typeof resume; - /** Calculates the number of requests in a shard that are safe to be used. */ - safeRequestsPerShard: typeof safeRequestsPerShard; - /** Calculates the number of shards to use based on the max concurrency */ - calculateMaxShards: typeof calculateMaxShards; -} - -export type GatewayDebugEvents = - | "GW ERROR" - | "GW CLOSED" - | "GW CLOSED_RECONNECT" - | "GW RAW" - | "GW RECONNECT" - | "GW INVALID_SESSION" - | "GW RESUMED" - | "GW RESUMING" - | "GW IDENTIFYING" - | "GW RAW_SEND" - | "GW MAX REQUESTS" - | "GW DEBUG" - | "GW HEARTBEATING" - | "GW HEARTBEATING_STARTED" - | "GW HEARTBEATING_DETAILS" - | "GW HEARTBEATING_CLOSED"; diff --git a/gateway/handleOnMessage.ts b/gateway/handleOnMessage.ts deleted file mode 100644 index fdd5e33eb..000000000 --- a/gateway/handleOnMessage.ts +++ /dev/null @@ -1,182 +0,0 @@ -import { GatewayManager } from "./gatewayManager.ts"; -import { snowflakeToBigint } from "../util/bigint.ts"; -import { delay } from "../util/utils.ts"; -import { decompressWith } from "./deps.ts"; -import { - DiscordGatewayPayload, - DiscordGuild, - DiscordHello, - DiscordMessage, - DiscordReady, - DiscordUnavailableGuild, -} from "../types/discord.ts"; -import { GatewayEventNames, GatewayOpcodes } from "../types/shared.ts"; - -/** Handler for handling every message event from websocket. */ -// deno-lint-ignore no-explicit-any -export async function handleOnMessage(gateway: GatewayManager, message: any, shardId: number) { - if (gateway.compress && message instanceof Blob) { - message = decompressWith( - new Uint8Array(await message.arrayBuffer()), - 0, - (slice: Uint8Array) => gateway.utf8decoder.decode(slice), - ); - } - - if (typeof message !== "string") return; - - const shard = gateway.shards.get(shardId); - - if (shard) { - // Edge case for big bots when too many events that 45 seconds are not enough for receving the heartbeat ack. As long as we are receving events no point in closing a connection. - shard.heartbeat.acknowledged = true; - } - - const messageData = JSON.parse(message) as DiscordGatewayPayload; - gateway.debug("GW RAW", { shardId, payload: messageData }); - - switch (messageData.op) { - case GatewayOpcodes.Heartbeat: - if (shard?.ws.readyState !== WebSocket.OPEN) return; - - shard.heartbeat.lastSentAt = Date.now(); - // Discord randomly sends this requiring an immediate heartbeat back - gateway.sendShardMessage( - gateway, - shard, - { - op: GatewayOpcodes.Heartbeat, - d: shard?.previousSequenceNumber, - }, - true, - ); - break; - case GatewayOpcodes.Hello: - gateway.heartbeat(gateway, shardId, (messageData.d as DiscordHello).heartbeat_interval); - // UPDATES THE SAFE AMOUNT OF SHARDS BASED ON THE INTERVAL - if (shard) shard.safeRequestsPerShard = gateway.safeRequestsPerShard(gateway, shard); - break; - case GatewayOpcodes.HeartbeatACK: - if (shard) { - shard.heartbeat.acknowledged = true; - shard.heartbeat.lastReceivedAt = Date.now(); - } - break; - case GatewayOpcodes.Reconnect: - gateway.debug("GW RECONNECT", { shardId }); - - if (gateway.shards.has(shardId)) { - gateway.shards.get(shardId)!.resuming = true; - } - - gateway.resume(gateway, shardId); - break; - case GatewayOpcodes.InvalidSession: - gateway.debug("GW INVALID_SESSION", { shardId, payload: messageData }); - - // We need to wait for a random amount of time between 1 and 5: https://discord.com/developers/docs/topics/gateway#resuming - await delay(Math.floor((Math.random() * 4 + 1) * 1000)); - - // When d is false we need to re-identify - if (!messageData.d) { - await gateway.identify(gateway, shardId, gateway.maxShards); - break; - } - - if (gateway.shards.has(shardId)) { - gateway.shards.get(shardId)!.resuming = true; - } - - gateway.resume(gateway, shardId); - break; - default: - if (messageData.t === "RESUMED") { - gateway.debug("GW RESUMED", { shardId }); - - if (gateway.shards.has(shardId)) { - gateway.shards.get(shardId)!.resuming = false; - } - break; - } - - // Important for RESUME - if (messageData.t === "READY") { - // Wait few seconds to spawn next shard - const bucket = gateway.buckets.get(shardId % gateway.maxConcurrency); - if (bucket?.createNextShard.length) { - // await delay(gateway.spawnShardDelay); - // setTimeout(() => { - bucket.createNextShard.shift()?.(); - // }, gateway.spawnShardDelay); - } - - const shard = gateway.shards.get(shardId); - const payload = messageData.d as DiscordReady; - - if (shard) { - shard.sessionId = payload.session_id; - shard.ready = true; - } - - payload.guilds.forEach((g) => gateway.cache.loadingGuildIds.add(snowflakeToBigint(g.id))); - - gateway.loadingShards.get(shardId)?.resolve(true); - gateway.loadingShards.delete(shardId); - } - - // Update the sequence number if it is present - if (messageData.s) { - const shard = gateway.shards.get(shardId); - if (shard) { - shard.previousSequenceNumber = messageData.s; - } - } - - // MUST HANDLE GUILD_CREATE EVENTS AS THEY ARE EXPENSIVE WITHOUT GATEWAY CACHE - if (messageData.t === "GUILD_CREATE") { - const id = snowflakeToBigint((messageData.d as DiscordGuild).id); - - // SHARD RESUMED MOST LIKELY, THEY EMIT GUILD CREATES. OR GUILD BECAME AVAILABLE AGAIN - if (gateway.cache.guildIds.has(id)) return; - - // GUILD WAS MARKED LOADING IN READY EVENT, THIS WAS THE FIRST GUILD_CREATE TO ARRIVE - if (gateway.cache.loadingGuildIds.has(id)) { - messageData.t = "GUILD_LOADED_DD" as GatewayEventNames; - gateway.cache.loadingGuildIds.delete(id); - } - - gateway.cache.guildIds.add(id); - } - - // MESSAGE_UPDATE CAN SPAM FOR NO REASON USE THIS TO IGNORE - if (messageData.t === "MESSAGE_UPDATE") { - const payload = messageData.d as DiscordMessage; - - const id = snowflakeToBigint(payload.id); - const content = payload.content || ""; - const cached = gateway.cache.editedMessages.get(id); - - if (cached === content) return; - else { - // ADD TO LOCAL CACHE FOR FUTURE EVENTS. - gateway.cache.editedMessages.set(id, content); - // REMOVE AFTER 10 SECONDS FROM CACHE - setTimeout(() => { - gateway.cache.editedMessages.delete(id); - }, 10000); - } - } - - // MUST HANDLE GUILD_DELETE EVENTS FOR UNAVAILABLE - if (messageData.t === "GUILD_DELETE") { - if ((messageData.d as DiscordUnavailableGuild).unavailable) return; - } - - // IF NO TYPE THEN THIS SHOULD NOT BE SENT FORWARD - if (!messageData.t) return; - - await gateway.handleDiscordPayload(gateway, messageData, shardId); - - break; - } -} diff --git a/gateway/heartbeat.ts b/gateway/heartbeat.ts deleted file mode 100644 index ef0dfaaad..000000000 --- a/gateway/heartbeat.ts +++ /dev/null @@ -1,61 +0,0 @@ -import { GatewayOpcodes } from "../types/shared.ts"; -import { delay } from "../util/utils.ts"; -import { GatewayManager } from "./gatewayManager.ts"; - -export async function heartbeat(gateway: GatewayManager, shardId: number, interval: number) { - gateway.debug("GW HEARTBEATING_STARTED", { shardId, interval }); - - const shard = gateway.shards.get(shardId); - if (!shard) return; - - gateway.debug("GW HEARTBEATING_DETAILS", { shardId, interval, shard }); - - // The first heartbeat is special so we send it without set Interval: https://discord.com/developers/docs/topics/gateway#heartbeating - await delay(Math.floor(shard.heartbeat.interval * Math.random())); - - if (shard.ws.readyState !== WebSocket.OPEN) return; - - shard.ws.send( - JSON.stringify({ - op: GatewayOpcodes.Heartbeat, - d: shard.previousSequenceNumber, - }), - ); - - shard.heartbeat.keepAlive = true; - shard.heartbeat.acknowledged = false; - shard.heartbeat.lastSentAt = Date.now(); - shard.heartbeat.interval = interval; - - shard.heartbeat.intervalId = setInterval(async () => { - gateway.debug("GW DEBUG", `Running setInterval in heartbeat file. Shard: ${shardId}`); - const currentShard = gateway.shards.get(shardId); - if (!currentShard) return; - - gateway.debug("GW HEARTBEATING", { shardId, shard: currentShard }); - - if (currentShard.ws.readyState === WebSocket.CLOSED || !currentShard.heartbeat.keepAlive) { - gateway.debug("GW HEARTBEATING_CLOSED", { shardId, shard: currentShard }); - - // STOP THE HEARTBEAT - return clearInterval(shard.heartbeat.intervalId); - } - - if (!currentShard.heartbeat.acknowledged) { - gateway.closeWS(currentShard.ws, 3066, "Did not receive an ACK in time."); - return await gateway.identify(gateway, shardId, gateway.maxShards); - } - - if (currentShard.ws.readyState !== WebSocket.OPEN) return; - - currentShard.heartbeat.acknowledged = false; - currentShard.heartbeat.lastSentAt = Date.now(); - - currentShard.ws.send( - JSON.stringify({ - op: GatewayOpcodes.Heartbeat, - d: currentShard.previousSequenceNumber, - }), - ); - }, shard.heartbeat.interval); -} diff --git a/gateway/identify.ts b/gateway/identify.ts deleted file mode 100644 index d72de5544..000000000 --- a/gateway/identify.ts +++ /dev/null @@ -1,71 +0,0 @@ -import { GatewayOpcodes } from "../types/shared.ts"; -import { GatewayManager } from "./gatewayManager.ts"; - -export function identify(gateway: GatewayManager, shardId: number, maxShards: number) { - gateway.debug("GW IDENTIFYING", { shardId, maxShards }); - - // Need to clear the old heartbeat interval - const oldShard = gateway.shards.get(shardId); - if (oldShard) { - gateway.closeWS(oldShard.ws, 3065, "Re-identifying closure of old shard"); - clearInterval(oldShard.heartbeat.intervalId); - } - - // CREATE A SHARD - const socket = gateway.createShard(gateway, shardId); - - // Identify can just set/reset the settings for the shard - gateway.shards.set(shardId, { - id: shardId, - ws: socket, - sessionId: "", - previousSequenceNumber: 0, - resuming: false, - ready: false, - unavailableGuildIds: new Set(), - heartbeat: { - lastSentAt: 0, - lastReceivedAt: 0, - acknowledged: false, - keepAlive: false, - interval: 0, - intervalId: 0, - }, - queue: [], - processingQueue: false, - queueStartedAt: Date.now(), - queueCounter: 0, - // BY DEFAULT SET TO 120. EDIT IN HELLO - safeRequestsPerShard: 120, - }); - - socket.onopen = () => { - gateway.sendShardMessage( - gateway, - shardId, - { - op: GatewayOpcodes.Identify, - d: { - token: `Bot ${gateway.token}`, - compress: gateway.compress, - properties: { - $os: gateway.$os, - $browser: gateway.$browser, - $device: gateway.$device, - }, - intents: gateway.intents, - shard: [shardId, maxShards], - presence: gateway.presence, - }, - }, - true, - ); - }; - - return new Promise((resolve) => { - gateway.loadingShards.set(shardId, { - shardId, - resolve, - }); - }); -} diff --git a/gateway/manager/calculateTotalShards.ts b/gateway/manager/calculateTotalShards.ts new file mode 100644 index 000000000..44817f528 --- /dev/null +++ b/gateway/manager/calculateTotalShards.ts @@ -0,0 +1,16 @@ +import { GatewayManager } from "./gatewayManager.ts"; + +/** Handler used to determine max number of shards to use based upon the max concurrency. */ +export function calculateTotalShards(gateway: GatewayManager): number { + // Bots under 100k servers do not have access to total shards. + if (gateway.manager.totalShards < 100) return gateway.manager.totalShards; + + // Calculate a multiple of `maxConcurrency` which can be used to connect to the gateway. + return Math.ceil( + gateway.manager.totalShards / + // If `maxConcurrency` is 1 we can safely use 16. + (gateway.gatewayBot.sessionStartLimit.maxConcurrency === 1 + ? 16 + : gateway.gatewayBot.sessionStartLimit.maxConcurrency), + ) * gateway.gatewayBot.sessionStartLimit.maxConcurrency; +} diff --git a/gateway/manager/calculateWorkerId.ts b/gateway/manager/calculateWorkerId.ts new file mode 100644 index 000000000..cd2ac2e62 --- /dev/null +++ b/gateway/manager/calculateWorkerId.ts @@ -0,0 +1,13 @@ +import { GatewayManager } from "./gatewayManager.ts"; + +export function calculateWorkerId(manager: GatewayManager, shardId: number) { + // Ignore decimal numbers. + let workerId = Math.floor((shardId) / manager.shardsPerWorker); + // If the workerId overflows the maximal allowed workers we by default just use to last worker. + if (workerId >= manager.totalWorkers) { + // The Id of the last available worker is total -1 + workerId = manager.totalWorkers - 1; + } + + return workerId; +} diff --git a/gateway/manager/gatewayManager.ts b/gateway/manager/gatewayManager.ts new file mode 100644 index 000000000..8507b9aa2 --- /dev/null +++ b/gateway/manager/gatewayManager.ts @@ -0,0 +1,296 @@ +import { GetGatewayBot } from "../../transformers/gatewayBot.ts"; +import { DiscordGatewayPayload } from "../../types/discord.ts"; +import { GatewayIntents, MakeRequired, OmitFirstFnArg, PickPartial } from "../../types/shared.ts"; +import { LeakyBucket } from "../../util/bucket.ts"; +import { Collection } from "../../util/collection.ts"; +import { CreateShard, createShard } from "../shard/createShard.ts"; +import { Shard, ShardGatewayConfig } from "../shard/types.ts"; +import { calculateTotalShards } from "./calculateTotalShards.ts"; +import { calculateWorkerId } from "./calculateWorkerId.ts"; +// import { +// markNewGuildShardId, +// resharder, +// resharderCloseOldShards, +// resharderIsPending, +// reshardingEditGuildShardIds, +// } from "./resharder.ts"; +import { spawnShards } from "./spawnShards.ts"; +import { prepareBuckets } from "./prepareBuckets.ts"; +import { tellWorkerToIdentify } from "./tellWorkerToIdentify.ts"; +import { createShardManager, ShardManager } from "./shardManager.ts"; +import { stop } from "./stop.ts"; + +export type GatewayManager = ReturnType; + +/** Create a new Gateway Manager. + * + * @param options: Customize every bit of the manager. If something is not + * provided, it will fallback to a default which should be suitable for most + * bots. + */ +export function createGatewayManager( + options: PickPartial, +) { + const prepareBucketsOverwritten = options.prepareBuckets ?? prepareBuckets; + const spawnShardsOverwritten = options.spawnShards ?? spawnShards; + const stopOverwritten = options.stop ?? stop; + const tellWorkerToIdentifyOverwritten = options.tellWorkerToIdentify ?? tellWorkerToIdentify; + const calculateTotalShardsOverwritten = options.calculateTotalShards ?? calculateTotalShards; + const calculateWorkerIdOverwritten = options.calculateWorkerId ?? calculateWorkerId; + + const totalShards = options.totalShards ?? options.gatewayBot.shards ?? 1; + + const gatewayManager = { + // ---------- + // PROPERTIES + // ---------- + + /** The max concurrency buckets. + * Those will be created when the `spawnShards` (which calls `prepareBuckets` under the hood) function gets called. + */ + buckets: new Collection< + number, + { + workers: { id: number; queue: number[] }[]; + leak: LeakyBucket; + } + >(), + /** Id of the first Shard which should get controlled by this manager. + * + * NOTE: This is intended for testing purposes + * if big bots want to test the gateway on smaller scale. + * This is not recommended to be used in production. + */ + firstShardId: options.firstShardId ?? 0, + /** Important data which is used by the manager to connect shards to the gateway. */ + gatewayBot: options.gatewayBot, + /** Id of the last Shard which should get controlled by this manager. + * + * NOTE: This is intended for testing purposes + * if big bots want to test the gateway on smaller scale. + * This is not recommended to be used in production. + */ + lastShardId: options.lastShardId ?? totalShards - 1 ?? 1, + /** This is where the Shards get stored. + * This will not be used when having a custom workers solution. + */ + manager: {} as ShardManager, + /** Delay in milliseconds to wait before spawning next shard. + * OPTIMAL IS ABOVE 5100. YOU DON'T WANT TO HIT THE RATE LIMIT!!! + */ + spawnShardDelay: options.spawnShardDelay ?? 5300, + /** How many Shards should get assigned to a Worker. + * + * IMPORTANT: Discordeno will NOT spawn Workers for you. + * Instead you have to overwrite the `tellWorkerToIdentify` function to make that for you. + * Look at the [BigBot template gateway solution](https://github.com/discordeno/discordeno/tree/main/template/bigbot/src/gateway) for reference. + * + * NOTE: The last Worker will IGNORE this value, + * which means that the last worker can get assigned an unlimited amount of shards. + * This is not a bug but intended behavior and means you have to assign more workers to this manager. + */ + shardsPerWorker: options.shardsPerWorker ?? 25, + /** The total amount of Workers which get controlled by this manager. + * + * IMPORTANT: Discordeno will NOT spawn Workers for you. + * Instead you have to overwrite the `tellWorkerToIdentify` function to make that for you. + * Look at the [BigBot template gateway solution](https://github.com/discordeno/discordeno/tree/main/template/bigbot/src/gateway) for reference. + */ + totalWorkers: options.totalWorkers ?? 4, + + // ---------- + // PROPERTIES + // ---------- + /** Prepares the buckets for identifying. + * + * NOTE: Most of the time this function does not need to be called, + * since it gets called by the `spawnShards` function indirectly. + */ + prepareBuckets: function () { + return prepareBucketsOverwritten(this); + }, + /** This function starts to spawn the Shards assigned to this manager. + * + * The managers `buckets` will be created and + * + * if `resharding.useOptimalLargeBotSharding` is set to true, + * `totalShards` gets double checked and adjusted accordingly if wrong. + */ + spawnShards: function () { + return spawnShardsOverwritten(this); + }, + /** Stop the gateway. This closes all shards. */ + stop: function (code: number, reason: string) { + return stopOverwritten(this, code, reason); + }, + /** Tell the Worker with this Id to identify this Shard. + * + * Useful if a custom Worker solution should be used. + * + * IMPORTANT: Discordeno will NOT spawn Workers for you. + * Instead you have to overwrite the `tellWorkerToIdentify` function to make that for you. + * Look at the [BigBot template gateway solution](https://github.com/discordeno/discordeno/tree/main/template/bigbot/src/gateway) for reference. + */ + tellWorkerToIdentify: function (workerId: number, shardId: number, bucketId: number) { + return tellWorkerToIdentifyOverwritten(this, workerId, shardId, bucketId); + }, + // TODO: fix debug + /** Handle the different logs. Used for debugging. */ + debug: options.debug || function () {}, + + // /** The methods related to resharding. */ + // resharding: { + // /** Whether the resharder should automatically switch to LARGE BOT SHARDING when the bot is above 100K servers. */ + // useOptimalLargeBotSharding: options.resharding?.useOptimalLargeBotSharding ?? true, + // /** Whether or not to automatically reshard. + // * + // * @default true + // */ + // reshard: options.resharding?.reshard ?? true, + // /** The percentage at which resharding should occur. + // * + // * @default 80 + // */ + // reshardPercentage: options.resharding?.reshardPercentage ?? 80, + // /** Handles resharding the bot when necessary. */ + // resharder: options.resharding?.resharder ?? resharder, + // /** Handles checking if all new shards are online in the new gateway. */ + // isPending: options.resharding?.isPending ?? resharderIsPending, + // /** Handles closing all shards in the old gateway. */ + // closeOldShards: options.resharding?.closeOldShards ?? resharderCloseOldShards, + // /** Handles checking if it is time to reshard and triggers the resharder. */ + // check: options.resharding?.check ?? startReshardingChecks, + // /** Handler to mark a guild id with its new shard id in cache. */ + // markNewGuildShardId: options.resharding?.markNewGuildShardId ?? markNewGuildShardId, + // /** Handler to update all guilds in cache with the new shard id. */ + // editGuildShardIds: options.resharding?.editGuildShardIds ?? reshardingEditGuildShardIds, + // }, + + /** Calculate the amount of Shards which should be used based on the bot's max concurrency. */ + calculateTotalShards: function () { + return calculateTotalShardsOverwritten(this); + }, + + /** Calculate the Id of the Worker related to this Shard. */ + calculateWorkerId: function (shardId: number) { + return calculateWorkerIdOverwritten(this, shardId); + }, + }; + + gatewayManager.manager = createShardManager({ + createShardOptions: options.createShardOptions, + gatewayConfig: options.gatewayConfig, + shardIds: [], + totalShards, + + handleMessage: function (shard, message) { + return options.handleDiscordPayload(shard, message); + }, + + requestIdentify: async (shardId) => { + // TODO: improve + await gatewayManager.buckets.get(shardId % gatewayManager.gatewayBot.sessionStartLimit.maxConcurrency)!.leak + .acquire(1); + }, + }); + + return gatewayManager; +} + +export interface CreateGatewayManager { + /** Delay in milliseconds to wait before spawning next shard. OPTIMAL IS ABOVE 5100. YOU DON'T WANT TO HIT THE RATE LIMIT!!! */ + spawnShardDelay: number; + /** Total amount of shards your bot uses. Useful for zero-downtime updates or resharding. */ + totalShards: number; + /** The amount of shards to load per worker. */ + shardsPerWorker: number; + /** The total amount of workers to use for your bot. */ + totalWorkers: number; + /** Id of the first Shard which should get controlled by this manager. + * + * NOTE: This is intended for testing purposes + * if big bots want to test the gateway on smaller scale. + * This is not recommended to be used in production. + */ + firstShardId: number; + /** Id of the last Shard which should get controlled by this manager. + * + * NOTE: This is intended for testing purposes + * if big bots want to test the gateway on smaller scale. + * This is not recommended to be used in production. + */ + lastShardId: number; + + /** Important data which is used by the manager to connect shards to the gateway. */ + gatewayBot: GetGatewayBot; + + gatewayConfig: PickPartial; + + /** Options which are used to create a new shard. */ + createShardOptions?: Omit; + + /** Stored as bucketId: { workers: [workerId, [ShardIds]], createNextShard: boolean } */ + buckets: Collection< + number, + { + workers: { id: number; queue: number[] }[]; + leak: LeakyBucket; + } + >; + // METHODS + + /** Prepares the buckets for identifying */ + prepareBuckets: typeof prepareBuckets; + /** The handler for spawning ALL the shards. */ + spawnShards: typeof spawnShards; + /** The handler to close all shards. */ + stop: typeof stop; + /** Sends the discord payload to another server. */ + handleDiscordPayload: (shard: Shard, data: DiscordGatewayPayload) => any; + /** Tell the worker to begin identifying this shard */ + tellWorkerToIdentify: typeof tellWorkerToIdentify; + /** Handle the different logs. Used for debugging. */ + debug: (text: GatewayDebugEvents, ...args: any[]) => unknown; + /** The methods related to resharding. */ + // resharding: { + // /** Whether the resharder should automatically switch to LARGE BOT SHARDING when you are above 100K servers. */ + // useOptimalLargeBotSharding: boolean; + // /** Whether or not to automatically reshard. */ + // reshard: boolean; + // /** The percentage at which resharding should occur. */ + // reshardPercentage: number; + // /** Handles resharding the bot when necessary. */ + // resharder: typeof resharder; + // /** Handles checking if all new shards are online in the new gateway. */ + // isPending: typeof resharderIsPending; + // /** Handles closing all shards in the old gateway. */ + // closeOldShards: typeof resharderCloseOldShards; + // /** Handler to mark a guild id with its new shard id in cache. */ + // markNewGuildShardId: typeof markNewGuildShardId; + // /** Handler to update all guilds in cache with the new shard id. */ + // editGuildShardIds: typeof reshardingEditGuildShardIds; + // }; + /** Calculates the number of shards to use based on the max concurrency */ + calculateTotalShards: typeof calculateTotalShards; + + /** Calculate the id of the worker related ot this Shard. */ + calculateWorkerId: typeof calculateWorkerId; +} + +export type GatewayDebugEvents = + | "GW ERROR" + | "GW CLOSED" + | "GW CLOSED_RECONNECT" + | "GW RAW" + | "GW RECONNECT" + | "GW INVALID_SESSION" + | "GW RESUMED" + | "GW RESUMING" + | "GW IDENTIFYING" + | "GW RAW_SEND" + | "GW MAX REQUESTS" + | "GW DEBUG" + | "GW HEARTBEATING" + | "GW HEARTBEATING_STARTED" + | "GW HEARTBEATING_DETAILS" + | "GW HEARTBEATING_CLOSED"; diff --git a/gateway/manager/mod.ts b/gateway/manager/mod.ts new file mode 100644 index 000000000..f3c056780 --- /dev/null +++ b/gateway/manager/mod.ts @@ -0,0 +1,8 @@ +export * from "./calculateTotalShards.ts"; +export * from "./calculateWorkerId.ts"; +export * from "./gatewayManager.ts"; +export * from "./prepareBuckets.ts"; +export * from "./shardManager.ts"; +export * from "./spawnShards.ts"; +export * from "./stop.ts"; +export * from "./tellWorkerToIdentify.ts"; diff --git a/gateway/manager/prepareBuckets.ts b/gateway/manager/prepareBuckets.ts new file mode 100644 index 000000000..308ea3ae9 --- /dev/null +++ b/gateway/manager/prepareBuckets.ts @@ -0,0 +1,47 @@ +import { createLeakyBucket } from "../../util/bucket.ts"; +import { GatewayManager } from "./gatewayManager.ts"; + +export function prepareBuckets(gateway: GatewayManager) { + for (let i = 0; i < gateway.gatewayBot.sessionStartLimit.maxConcurrency; ++i) { + gateway.buckets.set(i, { + workers: [], + leak: createLeakyBucket({ + max: 1, + refillAmount: 1, + // special number which is proven to be working dont change + refillInterval: gateway.spawnShardDelay, + }), + }); + } + + // ORGANIZE ALL SHARDS INTO THEIR OWN BUCKETS + for (let shardId = gateway.firstShardId; shardId <= gateway.lastShardId; ++shardId) { + // gateway.debug("GW DEBUG", `1. Running for loop in spawnShards function for shardId ${i}.`); + if (shardId >= gateway.manager.totalShards) { + throw new Error( + `Shard (id: ${shardId}) is bigger or equal to the used amount of used shards which is ${gateway.manager.totalShards}`, + ); + } + + const bucketId = shardId % gateway.gatewayBot.sessionStartLimit.maxConcurrency; + const bucket = gateway.buckets.get(bucketId); + if (!bucket) { + throw new Error( + `Shard (id: ${shardId}) got assigned to an illegal bucket id: ${bucketId}, expected a bucket id between 0 and ${ + gateway.gatewayBot.sessionStartLimit.maxConcurrency - 1 + }`, + ); + } + + // FIND A QUEUE IN THIS BUCKET THAT HAS SPACE + // const worker = bucket.workers.find((w) => w.queue.length < gateway.shardsPerWorker); + const workerId = gateway.calculateWorkerId(shardId); + const worker = bucket.workers.find((w) => w.id === workerId); + if (worker) { + // IF THE QUEUE HAS SPACE JUST ADD IT TO THIS QUEUE + worker.queue.push(shardId); + } else { + bucket.workers.push({ id: workerId, queue: [shardId] }); + } + } +} diff --git a/gateway/manager/resharder.ts b/gateway/manager/resharder.ts new file mode 100644 index 000000000..34949ffb6 --- /dev/null +++ b/gateway/manager/resharder.ts @@ -0,0 +1,341 @@ +import { GetGatewayBot, transformGatewayBot } from "../../transformers/gatewayBot.ts"; +import { DiscordReady } from "../../types/discord.ts"; +import { Collection } from "../../util/collection.ts"; +import { createGatewayManager, GatewayManager } from "./gatewayManager.ts"; + +export type Resharder = ReturnType; + +export function activateResharder(options: ActivateResharderOptions) { + const resharder = { + // ---------- + // PROPERTIES + // ---------- + + /** Interval in milliseconds of when to check whether it's time to reshard. + * + * @default 28800000 (8 hours) + */ + checkInterval: options.checkInterval || 28800000, + + /** Gateway manager which is currently processing all shards and events. */ + gateway: options.gatewayManager, + + /** Timeout of the reshard checker. */ + intervalId: undefined as number | undefined, + + /** Percentage at which resharding should occur. + * @default 80 + */ + percentage: options.percentage ?? 80, + + /** Whether the resharder should automatically switch to LARGE BOT SHARDING when the bot is above 100K servers. */ + useOptimalLargeBotSharding: options.useOptimalLargeBotSharding ?? true, + + // ---------- + // METHODS + // ---------- + + /** Activate the resharder and delay the next reshard check. */ + activate: function () { + return activate(this); + }, + + /** Function which is used to fetch the current gateway information of the bot. + * This function is mainly used by the reshard checker. + */ + getGatewayBot: options.getGatewayBot, + + /** Reshard the bots gateway. */ + reshard: function (gatewayBot: GetGatewayBot) { + return reshard(this, gatewayBot); + }, + + tellWorkerToPrepare: options.tellWorkerToPrepare, + }; + + resharder.activate(); + + return resharder; +} + +// /** The methods related to resharding. */ +// resharding: { +// /** Whether the resharder should automatically switch to LARGE BOT SHARDING when the bot is above 100K servers. */ +// useOptimalLargeBotSharding: options.resharding?.useOptimalLargeBotSharding ?? true, +// /** Whether or not to automatically reshard. +// * +// * @default true +// */ +// reshard: options.resharding?.reshard ?? true, +// /** The percentage at which resharding should occur. +// * +// * @default 80 +// */ +// reshardPercentage: options.resharding?.reshardPercentage ?? 80, +// /** Handles resharding the bot when necessary. */ +// resharder: options.resharding?.resharder ?? resharder, +// /** Handles checking if all new shards are online in the new gateway. */ +// isPending: options.resharding?.isPending ?? resharderIsPending, +// /** Handles closing all shards in the old gateway. */ +// closeOldShards: options.resharding?.closeOldShards ?? resharderCloseOldShards, +// /** Handles checking if it is time to reshard and triggers the resharder. */ +// check: options.resharding?.check ?? startReshardingChecks, +// /** Handler to mark a guild id with its new shard id in cache. */ +// markNewGuildShardId: options.resharding?.markNewGuildShardId ?? markNewGuildShardId, +// /** Handler to update all guilds in cache with the new shard id. */ +// editGuildShardIds: options.resharding?.editGuildShardIds ?? reshardingEditGuildShardIds, +// }, + +export interface ActivateResharderOptions { + /** Interval in milliseconds of when to check whether it's time to reshard. + * + * @default 28800000 (8 hours) + */ + checkInterval?: number; + /** Gateway manager which the resharder should be bound to. */ + gatewayManager: GatewayManager; + /** Percentage at which resharding should occur. + * @default 80 + */ + percentage?: number; + /** Whether the resharder should automatically switch to LARGE BOT SHARDING when the bot is above 100K servers. */ + useOptimalLargeBotSharding?: boolean; + + /** Function which can be used to fetch the current gateway information of the bot. + * This function is mainly used by the reshard checker. + */ + getGatewayBot(): Promise; + + /** Function which is used to tell a Worker that it should identify a resharder Shard to the gateway and wait for further instructions. + * The worker should **NOT** process any events coming from this Shard. + */ + tellWorkerToPrepare( + gatewayManager: GatewayManager, + workerId: number, + shardId: number, + bucketId: number, + ): Promise; +} + +/** Handler that by default will check to see if resharding should occur. Can be overridden if you have multiple servers and you want to communicate through redis pubsub or whatever you prefer. */ +export function activate(resharder: Resharder): void { + if (resharder.intervalId !== undefined) { + throw new Error("[RESHARDER] Cannot activate the resharder more than one time."); + } + + resharder.intervalId = setInterval(async () => { + // gateway.debug("GW DEBUG", "[Resharding] Checking if resharding is needed."); + + // TODO: is it possible to route this to REST? + const result = await resharder.getGatewayBot(); + + const percentage = + ((result.shards - resharder.gateway.manager.totalShards) / resharder.gateway.manager.totalShards) * 100; + // Less than necessary% being used so do nothing + if (percentage < resharder.percentage) return; + + // Don't have enough identify rate limits to reshard + if (result.sessionStartLimit.remaining < result.shards) return; + + // MULTI-SERVER BOTS OVERRIDE THIS IF YOU NEED TO RESHARD SERVER BY SERVER + return resharder.reshard(result); + }, resharder.checkInterval); +} + +export async function reshard(resharder: Resharder, gatewayBot: GetGatewayBot) { + // oldGateway.debug("GW DEBUG", "[Resharding] Starting the reshard process."); + + // Create a temporary gateway manager for easier handling. + const tmpManager = createGatewayManager({ + gatewayBot: gatewayBot, + gatewayConfig: resharder.gateway.manager.gatewayConfig, + handleDiscordPayload: () => {}, + tellWorkerToIdentify: resharder.tellWorkerToPrepare, + }); + + // Begin resharding + + // If more than 100K servers, begin switching to 16x sharding + if (resharder.useOptimalLargeBotSharding) { + // gateway.debug("GW DEBUG", "[Resharding] Using optimal large bot sharding solution."); + tmpManager.manager.totalShards = resharder.gateway.calculateTotalShards(resharder.gateway); + } + + tmpManager.spawnShards(tmpManager); + + return new Promise((resolve) => { + // TIMER TO KEEP CHECKING WHEN ALL SHARDS HAVE RESHARDED + const timer = setInterval(async () => { + const pending = await gateway.resharding.isPending(gateway, oldGateway); + // STILL PENDING ON SOME SHARDS TO BE CREATED + if (pending) return; + + // ENABLE EVENTS ON NEW SHARDS AND IGNORE EVENTS ON OLD + const oldHandler = oldGateway.handleDiscordPayload; + gateway.handleDiscordPayload = oldHandler; + oldGateway.handleDiscordPayload = function (og, data, shardId) { + // ALLOW EXCEPTION FOR CHUNKING TO PREVENT REQUESTS FREEZING + if (data.t !== "GUILD_MEMBERS_CHUNK") return; + oldHandler(og, data, shardId); + }; + + // STOP TIMER + clearInterval(timer); + await gateway.resharding.editGuildShardIds(); + await gateway.resharding.closeOldShards(oldGateway); + gateway.debug("GW DEBUG", "[Resharding] Complete."); + resolve(gateway); + }, 30000); + }) as Promise; +} + +// /** The handler to automatically reshard when necessary. */ +// export async function resharder( +// oldGateway: GatewayManager, +// results: GetGatewayBot, +// ) { +// oldGateway.debug("GW DEBUG", "[Resharding] Starting the reshard process."); + +// const gateway = createGatewayManager({ +// ...oldGateway, +// // RESET THE SETS AND COLLECTIONS +// // cache: { +// // guildIds: new Set(), +// // loadingGuildIds: new Set(), +// // editedMessages: new Collection(), +// // }, +// shards: new Collection(), +// // loadingShards: new Collection(), +// buckets: new Collection(), +// // utf8decoder: new TextDecoder(), +// }); + +// for (const [key, value] of Object.entries(oldGateway)) { +// if (key === "handleDiscordPayload") { +// gateway.handleDiscordPayload = async function (_, data, shardId) { +// if (data.t === "READY") { +// const payload = data.d as DiscordReady; +// await gateway.resharding.markNewGuildShardId(payload.guilds.map((g) => BigInt(g.id)), shardId); +// } +// }; +// continue; +// } + +// // USE ANY CUSTOMIZED OPTIONS FROM OLD GATEWAY +// // @ts-ignore TODO: fix this dynamical assignment +// gateway[key] = oldGateway[key as keyof typeof oldGateway]; +// } + +// // Begin resharding +// gateway.maxShards = results.shards; +// // FOR MANUAL SHARD CONTROL, OVERRIDE THIS SHARD ID! +// gateway.lastShardId = oldGateway.lastShardId === oldGateway.maxShards ? gateway.maxShards : oldGateway.lastShardId; +// gateway.shardsRecommended = results.shards; +// gateway.sessionStartLimitTotal = results.sessionStartLimit.total; +// gateway.sessionStartLimitRemaining = results.sessionStartLimit.remaining; +// gateway.sessionStartLimitResetAfter = results.sessionStartLimit.resetAfter; +// gateway.maxConcurrency = results.sessionStartLimit.maxConcurrency; +// // If more than 100K servers, begin switching to 16x sharding +// if (gateway.useOptimalLargeBotSharding) { +// gateway.debug("GW DEBUG", "[Resharding] Using optimal large bot sharding solution."); +// gateway.maxShards = gateway.calculateTotalShards(gateway.maxShards, results.sessionStartLimit.maxConcurrency); +// } + +// gateway.spawnShards(gateway, gateway.firstShardId); + +// return new Promise((resolve) => { +// // TIMER TO KEEP CHECKING WHEN ALL SHARDS HAVE RESHARDED +// const timer = setInterval(async () => { +// const pending = await gateway.resharding.isPending(gateway, oldGateway); +// // STILL PENDING ON SOME SHARDS TO BE CREATED +// if (pending) return; + +// // ENABLE EVENTS ON NEW SHARDS AND IGNORE EVENTS ON OLD +// const oldHandler = oldGateway.handleDiscordPayload; +// gateway.handleDiscordPayload = oldHandler; +// oldGateway.handleDiscordPayload = function (og, data, shardId) { +// // ALLOW EXCEPTION FOR CHUNKING TO PREVENT REQUESTS FREEZING +// if (data.t !== "GUILD_MEMBERS_CHUNK") return; +// oldHandler(og, data, shardId); +// }; + +// // STOP TIMER +// clearInterval(timer); +// await gateway.resharding.editGuildShardIds(); +// await gateway.resharding.closeOldShards(oldGateway); +// gateway.debug("GW DEBUG", "[Resharding] Complete."); +// resolve(gateway); +// }, 30000); +// }) as Promise; +// } + +/** Handler that by default will check all new shards are online in the new gateway. The handler can be overridden if you have multiple servers to communicate through redis pubsub or whatever you prefer. */ +export async function resharderIsPending( + gateway: GatewayManager, + oldGateway: GatewayManager, +) { + for (let i = gateway.firstShardId; i < gateway.lastShardId; i++) { + const shard = gateway.shards.get(i); + if (!shard?.ready) { + return true; + } + } + + return false; +} + +/** Handler that by default closes all shards in the old gateway. Can be overridden if you have multiple servers and you want to communicate through redis pubsub or whatever you prefer. */ +export async function resharderCloseOldShards(oldGateway: GatewayManager) { + // SHUT DOWN ALL SHARDS IF NOTHING IN QUEUE + oldGateway.shards.forEach((shard) => { + // CLOSE THIS SHARD IT HAS NO QUEUE + if (!shard.processingQueue && !shard.queue.length) { + return oldGateway.closeWS( + shard.ws, + 3066, + "Shard has been resharded. Closing shard since it has no queue.", + ); + } + + // IF QUEUE EXISTS GIVE IT 5 MINUTES TO COMPLETE + setTimeout(() => { + oldGateway.closeWS( + shard.ws, + 3066, + "Shard has been resharded. Delayed closing shard since it had a queue.", + ); + }, 300000); + }); +} + +// /** Handler that by default will check to see if resharding should occur. Can be overridden if you have multiple servers and you want to communicate through redis pubsub or whatever you prefer. */ +// export async function startReshardingChecks(gateway: GatewayManager) { +// gateway.debug("GW DEBUG", "[Resharding] Checking if resharding is needed."); +// +// // TODO: is it possible to route this to REST? +// const results = (await fetch(`https://discord.com/api/gateway/bot`, { +// headers: { +// Authorization: `Bot ${gateway.token}`, +// }, +// }).then((res) => res.json()).then((res) => transformGatewayBot(res))) as GetGatewayBot; +// +// const percentage = ((results.shards - gateway.maxShards) / gateway.maxShards) * 100; +// // Less than necessary% being used so do nothing +// if (percentage < gateway.reshardPercentage) return; +// +// // Don't have enough identify rate limits to reshard +// if (results.sessionStartLimit.remaining < results.shards) return; +// +// // MULTI-SERVER BOTS OVERRIDE THIS IF YOU NEED TO RESHARD SERVER BY SERVER +// return gateway.resharding.resharder(gateway, results); +// } + +/** Handler that by default will save the new shard id for each guild this becomes ready in new gateway. This can be overridden to save the shard ids in a redis cache layer or whatever you prefer. These ids will be used later to update all guilds. */ +export async function markNewGuildShardId(guildIds: bigint[], shardId: number) { + // PLACEHOLDER TO LET YOU MARK A GUILD ID AND SHARD ID FOR LATER USE ONCE RESHARDED +} + +/** Handler that by default does not do anything since by default the library will not cache. */ +export async function reshardingEditGuildShardIds() { + // PLACEHOLDER TO LET YOU UPDATE CACHED GUILDS +} diff --git a/gateway/manager/shardManager.ts b/gateway/manager/shardManager.ts new file mode 100644 index 000000000..38af3f2e1 --- /dev/null +++ b/gateway/manager/shardManager.ts @@ -0,0 +1,122 @@ +import { DiscordGatewayPayload } from "../../types/discord.ts"; +import { PickPartial } from "../../types/shared.ts"; +import { Collection } from "../../util/collection.ts"; +import { CreateShard, createShard } from "../shard/createShard.ts"; +import { Shard, ShardGatewayConfig } from "../shard/types.ts"; + +// TODO: debug + +/** This is a Shard manager. + * This does not manage a specific range of Shard but the provided Shards on create or when an identify is requested. + * The aim of this is to provide an easy to use manager which can be used by workers or any other kind of separate process. + */ +export type ShardManager = ReturnType; + +/** Create a new Shard manager. + * This does not manage a specific range of Shard but the provided Shards on create or when an identify is requested. + * The aim of this is to provide an easy to use manager which can be used by workers or any other kind of separate process. + */ +export function createShardManager(options: CreateShardManager) { + return { + // ---------- + // PROPERTIES + // ---------- + + /** Options which are used to create a new Shard. */ + createShardOptions: { + ...options.createShardOptions, + events: { + ...options.createShardOptions?.events, + message: options.createShardOptions?.events?.message ?? options.handleMessage, + }, + }, + /** Gateway configuration which is used when creating a Shard. */ + gatewayConfig: options.gatewayConfig, + /** Managed Shards. */ + shards: new Collection( + options.shardIds.map((shardId) => { + const shard = createShard({ + ...options.createShardOptions, + id: shardId, + totalShards: options.totalShards, + gatewayConfig: options.gatewayConfig, + requestIdentify: async function () { + return await options.requestIdentify(shardId); + }, + }); + + return [shardId, shard] as const; + }), + ), + /** Total amount of Shards used by the bot. */ + totalShards: options.totalShards, + + // ---------- + // METHODS + // ---------- + + /** Tell the manager to identify a Shard. + * If this Shard is not already managed this will also add the Shard to the manager. + */ + identify: async function (shardId: number) { + let shard = this.shards.get(shardId); + if (!shard) { + shard = createShard({ + ...this.createShardOptions, + id: shardId, + totalShards: this.totalShards, + gatewayConfig: this.gatewayConfig, + requestIdentify: async function () { + return await options.requestIdentify(shardId); + }, + }); + + this.shards.set(shardId, shard); + } + + return await shard.identify(); + }, + + /** Kill a shard. + * Close a shards connection to Discord's gateway (if any) and remove it from the manager. + */ + kill: async function (shardId: number) { + const shard = this.shards.get(shardId); + if (!shard) return; + + this.shards.delete(shardId); + return await shard.shutdown(); + }, + + /** This function communicates with the parent manager, + * in order to know whether this manager is allowed to identify a new shard. + */ + requestIdentify: options.requestIdentify, + }; +} + +export interface CreateShardManager { + // ---------- + // PROPERTIES + // ---------- + /** Options which are used to create a new Shard. */ + createShardOptions?: Omit; + /** Gateway configuration which is used when creating a Shard. */ + gatewayConfig: PickPartial; + /** Ids of the Shards which should be managed. */ + shardIds: number[]; + /** Total amount of Shard used by the bot. */ + totalShards: number; + + // ---------- + // METHODS + // ---------- + + /** This function is used when a shard receives any message from Discord. */ + handleMessage(shard: Shard, message: DiscordGatewayPayload): unknown; + + /** This function communicates with the parent manager, + * in order to know whether this manager is allowed to identify a new shard. # + */ + requestIdentify(shardId: number): Promise; +} diff --git a/gateway/manager/spawnShards.ts b/gateway/manager/spawnShards.ts new file mode 100644 index 000000000..62df1edbe --- /dev/null +++ b/gateway/manager/spawnShards.ts @@ -0,0 +1,32 @@ +import { GatewayIntents } from "../../types/shared.ts"; +import { createLeakyBucket } from "../../util/bucket.ts"; +import { createShard } from "../shard/createShard.ts"; +import { Shard } from "../shard/types.ts"; +import { createGatewayManager, GatewayManager } from "./gatewayManager.ts"; + +/** Begin spawning shards. */ +export function spawnShards(gateway: GatewayManager) { + // PREPARES THE MAX SHARD COUNT BY CONCURRENCY + // if (manager.resharding.useOptimalLargeBotSharding) { + // // gateway.debug("GW DEBUG", "[Spawning] Using optimal large bot sharding solution."); + // manager.manager.totalShards = manager.calculateTotalShards( + // manager, + // ); + // } + + // PREPARES ALL SHARDS IN SPECIFIC BUCKETS + gateway.prepareBuckets(); + + // SPREAD THIS OUT TO DIFFERENT WORKERS TO BEGIN STARTING UP + gateway.buckets.forEach(async (bucket, bucketId) => { + // gateway.debug("GW DEBUG", `2. Running forEach loop in spawnShards function.`); + + for (const worker of bucket.workers) { + // gateway.debug("GW DEBUG", `3. Running for of loop in spawnShards function.`); + + for (const shardId of worker.queue) { + await gateway.tellWorkerToIdentify(worker.id, shardId, bucketId); + } + } + }); +} diff --git a/gateway/manager/stop.ts b/gateway/manager/stop.ts new file mode 100644 index 000000000..99508b0da --- /dev/null +++ b/gateway/manager/stop.ts @@ -0,0 +1,8 @@ +import { delay } from "../../util/utils.ts"; +import { GatewayManager } from "./gatewayManager.ts"; + +export async function stop(gateway: GatewayManager, code: number, reason: string) { + gateway.manager.shards.forEach((shard) => shard.close(code, reason)); + + await delay(5000); +} diff --git a/gateway/tellWorkerToIdentify.ts b/gateway/manager/tellWorkerToIdentify.ts similarity index 67% rename from gateway/tellWorkerToIdentify.ts rename to gateway/manager/tellWorkerToIdentify.ts index 46a43a21f..69054b648 100644 --- a/gateway/tellWorkerToIdentify.ts +++ b/gateway/manager/tellWorkerToIdentify.ts @@ -1,3 +1,5 @@ +import { GatewayIntents } from "../../types/shared.ts"; +import { createShard } from "../shard/createShard.ts"; import { GatewayManager } from "./gatewayManager.ts"; /** Allows users to hook in and change to communicate to different workers across different servers or anything they like. For example using redis pubsub to talk to other servers. */ @@ -6,6 +8,6 @@ export async function tellWorkerToIdentify( _workerId: number, shardId: number, _bucketId: number, -) { - await gateway.identify(gateway, shardId, gateway.maxShards); +): Promise { + return await gateway.manager.identify(shardId); } diff --git a/gateway/mod.ts b/gateway/mod.ts index 90713e518..9315ff337 100644 --- a/gateway/mod.ts +++ b/gateway/mod.ts @@ -1,16 +1,2 @@ -export * from "./calculateMaxShards.ts"; -export * from "./closeWs.ts"; -export * from "./createShard.ts"; -export * from "./gatewayManager.ts"; -export * from "./handleOnMessage.ts"; -export * from "./heartbeat.ts"; -export * from "./identify.ts"; -export * from "./processGatewayQueue.ts"; -export * from "./resharder.ts"; -export * from "./resume.ts"; -export * from "./safeRequestsPerShard.ts"; -export * from "./sendShardMessage.ts"; -export * from "./shard.ts"; -export * from "./spawnShards.ts"; -export * from "./stopGateway.ts"; -export * from "./tellWorkerToIdentify.ts"; +export * from "./manager/mod.ts"; +export * from "./shard/mod.ts"; diff --git a/gateway/processGatewayQueue.ts b/gateway/processGatewayQueue.ts deleted file mode 100644 index 9c2c3d348..000000000 --- a/gateway/processGatewayQueue.ts +++ /dev/null @@ -1,51 +0,0 @@ -import { delay } from "../util/utils.ts"; -import { GatewayManager } from "./gatewayManager.ts"; - -export async function processGatewayQueue(gateway: GatewayManager, id: number) { - const shard = gateway.shards.get(id); - // If no items or its already processing then exit - if (!shard?.queue.length || shard.processingQueue) return; - - shard.processingQueue = true; - - while (shard.queue.length) { - if (shard.ws.readyState !== WebSocket.OPEN) { - shard.processingQueue = false; - return; - } - - const now = Date.now(); - if (now - shard.queueStartedAt >= gateway.queueResetInterval) { - shard.queueStartedAt = now; - shard.queueCounter = 0; - } - - // Send a request that is next in line - const request = shard.queue.shift(); - if (!request) return; - - gateway.debug("GW RAW_SEND", shard.id, request); - - shard.ws.send(JSON.stringify(request)); - - // Counter is useful for preventing max requests. - shard.queueCounter++; - - // Handle if the requests have been maxed - if (shard.queueCounter >= shard.safeRequestsPerShard) { - const remaining = shard.queueStartedAt + gateway.queueResetInterval - Date.now(); - if (remaining > 0) { - gateway.debug("GW MAX REQUESTS", { - message: `Max gateway requests per minute reached setting timeout for ${remaining}ms`, - shardId: shard.id, - }); - await delay(remaining); - } - - shard.queueCounter = 0; - continue; - } - } - - shard.processingQueue = false; -} diff --git a/gateway/resharder.ts b/gateway/resharder.ts deleted file mode 100644 index 0aa097bcf..000000000 --- a/gateway/resharder.ts +++ /dev/null @@ -1,155 +0,0 @@ -import { GetGatewayBot, transformGatewayBot } from "../transformers/gatewayBot.ts"; -import { DiscordReady } from "../types/discord.ts"; -import { Collection } from "../util/collection.ts"; -import { createGatewayManager, GatewayManager } from "./gatewayManager.ts"; - -/** The handler to automatically reshard when necessary. */ -export async function resharder( - oldGateway: GatewayManager, - results: GetGatewayBot, -) { - oldGateway.debug("GW DEBUG", "[Resharding] Starting the reshard process."); - - const gateway = createGatewayManager({ - ...oldGateway, - // RESET THE SETS AND COLLECTIONS - cache: { - guildIds: new Set(), - loadingGuildIds: new Set(), - editedMessages: new Collection(), - }, - shards: new Collection(), - loadingShards: new Collection(), - buckets: new Collection(), - utf8decoder: new TextDecoder(), - }); - - for (const [key, value] of Object.entries(oldGateway)) { - if (key === "handleDiscordPayload") { - gateway.handleDiscordPayload = async function (_, data, shardId) { - if (data.t === "READY") { - const payload = data.d as DiscordReady; - await gateway.resharding.markNewGuildShardId(payload.guilds.map((g) => BigInt(g.id)), shardId); - } - }; - continue; - } - - // USE ANY CUSTOMIZED OPTIONS FROM OLD GATEWAY - // @ts-ignore TODO: fix this dynamical assignment - gateway[key] = oldGateway[key as keyof typeof oldGateway]; - } - - // Begin resharding - gateway.maxShards = results.shards; - // FOR MANUAL SHARD CONTROL, OVERRIDE THIS SHARD ID! - gateway.lastShardId = oldGateway.lastShardId === oldGateway.maxShards ? gateway.maxShards : oldGateway.lastShardId; - gateway.shardsRecommended = results.shards; - gateway.sessionStartLimitTotal = results.sessionStartLimit.total; - gateway.sessionStartLimitRemaining = results.sessionStartLimit.remaining; - gateway.sessionStartLimitResetAfter = results.sessionStartLimit.resetAfter; - gateway.maxConcurrency = results.sessionStartLimit.maxConcurrency; - // If more than 100K servers, begin switching to 16x sharding - if (gateway.useOptimalLargeBotSharding) { - gateway.debug("GW DEBUG", "[Resharding] Using optimal large bot sharding solution."); - gateway.maxShards = gateway.calculateMaxShards(gateway.maxShards, results.sessionStartLimit.maxConcurrency); - } - - gateway.spawnShards(gateway, gateway.firstShardId); - - return new Promise((resolve) => { - // TIMER TO KEEP CHECKING WHEN ALL SHARDS HAVE RESHARDED - const timer = setInterval(async () => { - const pending = await gateway.resharding.isPending(gateway, oldGateway); - // STILL PENDING ON SOME SHARDS TO BE CREATED - if (pending) return; - - // ENABLE EVENTS ON NEW SHARDS AND IGNORE EVENTS ON OLD - const oldHandler = oldGateway.handleDiscordPayload; - gateway.handleDiscordPayload = oldHandler; - oldGateway.handleDiscordPayload = function (og, data, shardId) { - // ALLOW EXCEPTION FOR CHUNKING TO PREVENT REQUESTS FREEZING - if (data.t !== "GUILD_MEMBERS_CHUNK") return; - oldHandler(og, data, shardId); - }; - - // STOP TIMER - clearInterval(timer); - await gateway.resharding.editGuildShardIds(); - await gateway.resharding.closeOldShards(oldGateway); - gateway.debug("GW DEBUG", "[Resharding] Complete."); - resolve(gateway); - }, 30000); - }) as Promise; -} - -/** Handler that by default will check all new shards are online in the new gateway. The handler can be overridden if you have multiple servers to communicate through redis pubsub or whatever you prefer. */ -export async function resharderIsPending( - gateway: GatewayManager, - oldGateway: GatewayManager, -) { - for (let i = gateway.firstShardId; i < gateway.lastShardId; i++) { - const shard = gateway.shards.get(i); - if (!shard?.ready) { - return true; - } - } - - return false; -} - -/** Handler that by default closes all shards in the old gateway. Can be overridden if you have multiple servers and you want to communicate through redis pubsub or whatever you prefer. */ -export async function resharderCloseOldShards(oldGateway: GatewayManager) { - // SHUT DOWN ALL SHARDS IF NOTHING IN QUEUE - oldGateway.shards.forEach((shard) => { - // CLOSE THIS SHARD IT HAS NO QUEUE - if (!shard.processingQueue && !shard.queue.length) { - return oldGateway.closeWS( - shard.ws, - 3066, - "Shard has been resharded. Closing shard since it has no queue.", - ); - } - - // IF QUEUE EXISTS GIVE IT 5 MINUTES TO COMPLETE - setTimeout(() => { - oldGateway.closeWS( - shard.ws, - 3066, - "Shard has been resharded. Delayed closing shard since it had a queue.", - ); - }, 300000); - }); -} - -/** Handler that by default will check to see if resharding should occur. Can be overridden if you have multiple servers and you want to communicate through redis pubsub or whatever you prefer. */ -export async function startReshardingChecks(gateway: GatewayManager) { - gateway.debug("GW DEBUG", "[Resharding] Checking if resharding is needed."); - - // TODO: is it possible to route this to REST? - const results = (await fetch(`https://discord.com/api/gateway/bot`, { - headers: { - Authorization: `Bot ${gateway.token}`, - }, - }).then((res) => res.json()).then((res) => transformGatewayBot(res))) as GetGatewayBot; - - const percentage = ((results.shards - gateway.maxShards) / gateway.maxShards) * 100; - // Less than necessary% being used so do nothing - if (percentage < gateway.reshardPercentage) return; - - // Don't have enough identify rate limits to reshard - if (results.sessionStartLimit.remaining < results.shards) return; - - // MULTI-SERVER BOTS OVERRIDE THIS IF YOU NEED TO RESHARD SERVER BY SERVER - return gateway.resharding.resharder(gateway, results); -} - -/** Handler that by default will save the new shard id for each guild this becomes ready in new gateway. This can be overridden to save the shard ids in a redis cache layer or whatever you prefer. These ids will be used later to update all guilds. */ -export async function markNewGuildShardId(guildIds: bigint[], shardId: number) { - // PLACEHOLDER TO LET YOU MARK A GUILD ID AND SHARD ID FOR LATER USE ONCE RESHARDED -} - -/** Handler that by default does not do anything since by default the library will not cache. */ -export async function reshardingEditGuildShardIds() { - // PLACEHOLDER TO LET YOU UPDATE CACHED GUILDS -} diff --git a/gateway/resume.ts b/gateway/resume.ts deleted file mode 100644 index 03ea355a6..000000000 --- a/gateway/resume.ts +++ /dev/null @@ -1,67 +0,0 @@ -import { GatewayOpcodes } from "../types/shared.ts"; -import { GatewayManager } from "./gatewayManager.ts"; - -export function resume(gateway: GatewayManager, shardId: number) { - gateway.debug("GW RESUMING", { shardId }); - - // NOW WE HANDLE RESUMING THIS SHARD - // Get the old data for this shard necessary for resuming - const oldShard = gateway.shards.get(shardId); - if (!oldShard) { - return gateway.debug( - "GW DEBUG", - `[Error] Trying to resume a shard (id: ${shardId}) that was not first identified.`, - ); - } - - // HOW TO CLOSE OLD SHARD SOCKET!!! - gateway.closeWS(oldShard.ws, 3064, "Resuming the shard, closing old shard."); - // STOP OLD HEARTBEAT - clearInterval(oldShard.heartbeat.intervalId); - - // CREATE A SHARD - const socket = gateway.createShard(gateway, shardId); - - const sessionId = oldShard.sessionId || ""; - const previousSequenceNumber = oldShard.previousSequenceNumber || 0; - - gateway.shards.set(shardId, { - id: shardId, - ws: socket, - sessionId: sessionId, - previousSequenceNumber: previousSequenceNumber, - resuming: false, - ready: false, - unavailableGuildIds: new Set(), - heartbeat: { - lastSentAt: 0, - lastReceivedAt: 0, - acknowledged: false, - keepAlive: false, - interval: 0, - intervalId: 0, - }, - queue: oldShard.queue || [], - processingQueue: false, - queueStartedAt: Date.now(), - queueCounter: 0, - safeRequestsPerShard: oldShard.safeRequestsPerShard || 120, - }); - - // Resume on open - socket.onopen = () => { - gateway.sendShardMessage( - gateway, - shardId, - { - op: GatewayOpcodes.Resume, - d: { - token: `Bot ${gateway.token}`, - session_id: sessionId, - seq: previousSequenceNumber, - }, - }, - true, - ); - }; -} diff --git a/gateway/safeRequestsPerShard.ts b/gateway/safeRequestsPerShard.ts deleted file mode 100644 index 07eece102..000000000 --- a/gateway/safeRequestsPerShard.ts +++ /dev/null @@ -1,9 +0,0 @@ -import { GatewayManager } from "./gatewayManager.ts"; -import { DiscordenoShard } from "./shard.ts"; - -export function safeRequestsPerShard(gateway: GatewayManager, shard: DiscordenoShard) { - // * 2 adds extra safety layer for discords OP 1 requests that we need to respond to - const safeRequests = gateway.maxRequestsPerInterval - - Math.ceil(gateway.queueResetInterval / shard.heartbeat.interval) * 2; - return safeRequests > 0 ? safeRequests : 0; -} diff --git a/gateway/sendShardMessage.ts b/gateway/sendShardMessage.ts deleted file mode 100644 index f9af2b29c..000000000 --- a/gateway/sendShardMessage.ts +++ /dev/null @@ -1,17 +0,0 @@ -import { GatewayManager } from "./gatewayManager.ts"; -import { DiscordenoShard, WebSocketRequest } from "./shard.ts"; - -export function sendShardMessage( - gateway: GatewayManager, - shard: number | DiscordenoShard, - message: WebSocketRequest, - highPriority = false, -) { - if (typeof shard === "number") shard = gateway.shards.get(shard)!; - if (!shard) return; - - if (highPriority) shard.queue.unshift(message); - else shard.queue.push(message); - - gateway.processGatewayQueue(gateway, shard.id); -} diff --git a/gateway/shard.ts b/gateway/shard.ts deleted file mode 100644 index 1b9d3ae26..000000000 --- a/gateway/shard.ts +++ /dev/null @@ -1,48 +0,0 @@ -import { GatewayOpcodes } from "../types/shared.ts"; - -export interface DiscordenoShard { - /** The shard id number. */ - id: number; - /** The websocket for this shard. */ - ws: WebSocket; - /** The session id important for resuming connections. */ - sessionId: string; - /** The previous sequence number, important for resuming connections. */ - previousSequenceNumber: number | null; - /** Whether the shard is currently resuming. */ - resuming: boolean; - /** Whether the shard has received the ready event. */ - ready: boolean; - /** The list of guild ids that are currently unavailable due to an outage. */ - unavailableGuildIds: Set; - failedToLoadTimeoutId?: number; - heartbeat: { - /** The exact timestamp the last heartbeat was sent. */ - lastSentAt: number; - /** The timestamp the last heartbeat ACK was received from discord. */ - lastReceivedAt: number; - /** Whether or not the heartbeat was acknowledged by discord in time. */ - acknowledged: boolean; - /** Whether or not to keep heartbeating. Useful for when needing to stop heartbeating. */ - keepAlive: boolean; - /** The interval between heartbeats requested by discord. */ - interval: number; - /** The id of the interval, useful for stopping the interval if ws closed. */ - intervalId: number; - }; - /** The items/requests that are in queue to be sent to this shard websocket. */ - queue: WebSocketRequest[]; - /** Whether or not the queue for this shard is being processed. */ - processingQueue: boolean; - /** When the first request for this minute has been sent. */ - queueStartedAt: number; - /** The request counter of the queue. */ - queueCounter: number; - /** The safe number of requests that can be made while preserving some for required things like heartbeating. */ - safeRequestsPerShard: number; -} - -export interface WebSocketRequest { - op: GatewayOpcodes; - d: unknown; -} diff --git a/gateway/shard/calculateSafeRequests.ts b/gateway/shard/calculateSafeRequests.ts new file mode 100644 index 000000000..ea0ca6a14 --- /dev/null +++ b/gateway/shard/calculateSafeRequests.ts @@ -0,0 +1,9 @@ +import { Shard } from "./types.ts"; + +export function calculateSafeRequests(shard: Shard) { + // * 2 adds extra safety layer for discords OP 1 requests that we need to respond to + const safeRequests = shard.maxRequestsPerRateLimitTick - + Math.ceil(shard.rateLimitResetInterval / shard.heart.interval) * 2; + + return safeRequests < 0 ? 0 : safeRequests; +} diff --git a/gateway/shard/close.ts b/gateway/shard/close.ts new file mode 100644 index 000000000..fc8d0b2a7 --- /dev/null +++ b/gateway/shard/close.ts @@ -0,0 +1,7 @@ +import { Shard } from "./types.ts"; + +export function close(shard: Shard, code: number, reason: string): void { + if (shard.socket?.readyState !== WebSocket.OPEN) return; + + return shard.socket?.close(code, reason); +} diff --git a/gateway/shard/connect.ts b/gateway/shard/connect.ts new file mode 100644 index 000000000..980466b23 --- /dev/null +++ b/gateway/shard/connect.ts @@ -0,0 +1,34 @@ +import { Shard, ShardState } from "./types.ts"; + +export async function connect(shard: Shard): Promise { + // Only set the shard to `Connecting` state, + // if the connection request does not come from an identify or resume action. + if (![ShardState.Identifying, ShardState.Resuming].includes(shard.state)) { + shard.state = ShardState.Connecting; + } + shard.events.connecting?.(shard); + + // Explicitly setting the encoding to json, since we do not support ETF. + const socket = new WebSocket(`${shard.gatewayConfig.url}/?v=${shard.gatewayConfig.version}&encoding=json`); + shard.socket = socket; + + // TODO: proper event handling + socket.onerror = (event) => console.log({ error: event }); + + socket.onclose = (event) => shard.handleClose(event); + + socket.onmessage = (message) => shard.handleMessage(message); + + return new Promise((resolve) => { + socket.onopen = () => { + // Only set the shard to `Unidentified` state, + // if the connection request does not come from an identify or resume action. + if (![ShardState.Identifying, ShardState.Resuming].includes(shard.state)) { + shard.state = ShardState.Unidentified; + } + shard.events.connected?.(shard); + + resolve(); + }; + }); +} diff --git a/gateway/shard/createShard.ts b/gateway/shard/createShard.ts new file mode 100644 index 000000000..de230d5c4 --- /dev/null +++ b/gateway/shard/createShard.ts @@ -0,0 +1,362 @@ +import { identify } from "./identify.ts"; +import { handleMessage } from "./handleMessage.ts"; +import { + DEFAULT_HEARTBEAT_INTERVAL, + GATEWAY_RATE_LIMIT_RESET_INTERVAL, + MAX_GATEWAY_REQUESTS_PER_INTERVAL, + Shard, + ShardEvents, + ShardGatewayConfig, + ShardHeart, + ShardSocketCloseCodes, + ShardSocketRequest, + ShardState, +} from "./types.ts"; +import { StatusUpdate } from "../../helpers/misc/editShardStatus.ts"; +import { startHeartbeating } from "./startHeartbeating.ts"; +import { stopHeartbeating } from "./stopHeartbeating.ts"; +import { resume } from "./resume.ts"; +import { createLeakyBucket, LeakyBucket } from "../../util/bucket.ts"; +import { calculateSafeRequests } from "./calculateSafeRequests.ts"; +import { send } from "./send.ts"; +import { handleClose } from "./handleClose.ts"; +import { connect } from "./connect.ts"; +import { close } from "./close.ts"; +import { shutdown } from "./shutdown.ts"; +import { isOpen } from "./isOpen.ts"; +import { DiscordGatewayPayload } from "../../types/discord.ts"; +import { GatewayIntents, PickPartial } from "../../types/shared.ts"; +import { API_VERSION } from "../../util/constants.ts"; + +// TODO: debug +// TODO: function overwrite +// TODO: improve shard event resolving + +/** */ +export function createShard( + options: CreateShard, +) { + // This is done for performance reasons + const calculateSafeRequestsOverwritten = options.calculateSafeRequests ?? calculateSafeRequests; + const closeOverwritten = options.close ?? close; + const connectOverwritten = options.connect ?? connect; + const identifyOverwritten = options.identify ?? identify; + const sendOverwritten = options.send ?? send; + const shutdownOverwritten = options.shutdown ?? shutdown; + const resumeOverwritten = options.resume ?? resume; + const handleCloseOverwritten = options.handleClose ?? handleClose; + const handleMessageOverwritten = options.handleMessage ?? handleMessage; + const isOpenOverwritten = options.isOpen ?? isOpen; + const startHeartbeatingOverwritten = options.startHeartbeating ?? startHeartbeating; + const stopHeartbeatingOverwritten = options.stopHeartbeating ?? stopHeartbeating; + + return { + // ---------- + // PROPERTIES + // ---------- + + /** The gateway configuration which is used to connect to Discord. */ + gatewayConfig: { + compress: options.gatewayConfig.compress ?? false, + intents: options.gatewayConfig.intents ?? 0, + properties: { + $os: options.gatewayConfig?.properties?.$os ?? Deno.build.os, + $browser: options.gatewayConfig?.properties?.$browser ?? "Discordeno", + $device: options.gatewayConfig?.properties?.$device ?? "Discordeno", + }, + token: options.gatewayConfig.token, + url: options.gatewayConfig.url ?? "wss://gateway.discord.gg", + version: options.gatewayConfig.version ?? API_VERSION, + } as ShardGatewayConfig, + /** This contains all the heartbeat information */ + heart: { + acknowledged: false, + interval: DEFAULT_HEARTBEAT_INTERVAL, + } as ShardHeart, + /** Id of the shard. */ + id: options.id, + /** The maximum of requests which can be send to discord per rate limit tick. + * Typically this value should not be changed. + */ + maxRequestsPerRateLimitTick: MAX_GATEWAY_REQUESTS_PER_INTERVAL, + /** The previous payload sequence number. */ + previousSequenceNumber: options.previousSequenceNumber || null, + /** In which interval (in milliseconds) the gateway resets it's rate limit. */ + rateLimitResetInterval: GATEWAY_RATE_LIMIT_RESET_INTERVAL, + /** Current session id of the shard if present. */ + sessionId: undefined as string | undefined, + /** This contains the WebSocket connection to Discord, if currently connected. */ + socket: undefined as WebSocket | undefined, + /** Current internal state of the shard. */ + state: ShardState.Offline, + /** The total amount of shards which are used to communicate with Discord. */ + totalShards: options.totalShards, + + // ---------- + // METHODS + // ---------- + + /** The shard related event handlers. */ + events: options.events ?? {} as ShardEvents, + + /** Calculate the amount of requests which can safely be made per rate limit interval, + * before the gateway gets disconnected due to an exceeded rate limit. + */ + calculateSafeRequests: function () { + return calculateSafeRequestsOverwritten(this); + }, + + /** Close the socket connection to discord if present. */ + close: function (code: number, reason: string) { + return closeOverwritten(this, code, reason); + }, + + /** Connect the shard with the gateway and start heartbeating. + * This will not identify the shard to the gateway. + */ + connect: async function () { + return await connectOverwritten(this); + }, + + /** Identify the shard to the gateway. + * If not connected, this will also connect the shard to the gateway. + */ + identify: async function () { + return await identifyOverwritten(this); + }, + + /** Check whether the connection to Discord is currently open. */ + isOpen: function () { + return isOpenOverwritten(this); + }, + + /** Function which can be overwritten in order to get the shards presence. */ + // This function allows to be async, in case the devs create the presence based on eg. database values. + // Passing the shard's id there to make it easier for the dev to use this function. + makePresence: options.makePresence, + + /** Attempt to resume the previous shards session with the gateway. */ + resume: async function () { + return await resumeOverwritten(this); + }, + + /** Send a message to Discord. + * @param {boolean} [highPriority=false] - Whether this message should be send asap. + */ + send: async function (message: ShardSocketRequest, highPriority: boolean = false) { + return await sendOverwritten(this, message, highPriority); + }, + + /** Shutdown the shard. + * Forcefully disconnect the shard from Discord. + * The shard may not attempt to reconnect with Discord. + */ + shutdown: async function () { + return await shutdownOverwritten(this); + }, + + /** @private Internal shard bucket. + * Only access this if you know what you are doing. + * + * Bucket for handling shard request rate limits. + */ + bucket: createLeakyBucket({ + max: MAX_GATEWAY_REQUESTS_PER_INTERVAL, + refillInterval: GATEWAY_RATE_LIMIT_RESET_INTERVAL, + refillAmount: MAX_GATEWAY_REQUESTS_PER_INTERVAL, + }), + + /** @private Internal shard function. + * Only use this function if you know what you are doing. + * + * Handle a gateway connection close. + */ + handleClose: async function (close: CloseEvent) { + return await handleCloseOverwritten(this, close); + }, + + /** @private Internal shard function. + * Only use this function if you know what you are doing. + * + * Handle an incoming gateway message. + */ + handleMessage: async function (message: MessageEvent) { + return await handleMessageOverwritten(this, message); + }, + + /** This function communicates with the management process, in order to know whether its free to identify. */ + requestIdentify: async function () { + return await options.requestIdentify(this.id); + }, + + /** @private Internal state. + * Only use this if you know what you are doing. + * + * Cache for pending gateway requests which should have been send while the gateway went offline. + */ + offlineSendQueue: [] as ((_?: unknown) => void)[], + + /** @private Internal shard map. + * Only use this map if you know what you are doing. + * + * This is used to resolve internal waiting states. + * Mapped by SelectedEvents => ResolveFunction + */ + resolves: new Map<"READY" | "RESUMED" | "INVALID_SESSION", (payload: DiscordGatewayPayload) => void>(), + + /** @private Internal shard function. + * Only use this function if you know what you are doing. + * + * Start sending heartbeat payloads to Discord in the provided interval. + */ + startHeartbeating: function (interval: number) { + return startHeartbeatingOverwritten(this, interval); + }, + + /** @private Internal shard function. + * Only use this function if you know what you are doing. + * + * Stop the heartbeating process with discord. + */ + stopHeartbeating: function () { + return stopHeartbeatingOverwritten(this); + }, + }; +} + +// const shard = createShard({ +// id: 0, +// gatewayConfig: { +// compress: true, +// url: "wss://gateway.discord.gg", +// version: 10, +// intents: 1 << 1, +// properties: { +// $os: "Discordeno", +// $browser: "Discordeno", +// $device: "Discordeno", +// }, +// token: TOKEN, +// }, +// totalShards: 1, +// makePresence: (shardId) => ({ +// activities: [ +// { +// name: `Cards Against Humanity #${shardId}`, +// type: 0, +// createdAt: Date.now(), +// }, +// ], +// status: "dnd", +// }), +// requestIdentify: async () => {}, +// }); + +export interface CreateShard { + /** Id of the shard which should be created. */ + id: number; + + /** Gateway configuration for the shard. */ + gatewayConfig: PickPartial; + + /** The total amount of shards which are used to communicate with Discord. */ + totalShards: number; + + /** This function communicates with the management process, in order to know whether its free to identify. + * When this function resolves, this means that the shard is allowed to send an identify payload to discord. + */ + requestIdentify: (shardId: number) => Promise; + + /** Calculate the amount of requests which can safely be made per rate limit interval, + * before the gateway gets disconnected due to an exceeded rate limit. + */ + calculateSafeRequests?: typeof calculateSafeRequests; + + /** Close the socket connection to discord if present. */ + close?: typeof close; + + /** Connect the shard with the gateway and start heartbeating. + * This will not identify the shard to the gateway. + */ + connect?: typeof connect; + + /** @private Internal shard function. + * Only use this function if you know what you are doing. + * + * Handle a gateway connection close. + */ + handleClose?: typeof handleClose; + + /** @private Internal shard function. + * Only use this function if you know what you are doing. + * + * Handle an incoming gateway message. + */ + handleMessage?: typeof handleMessage; + + /** Identify the shard to the gateway. + * If not connected, this will also connect the shard to the gateway. + */ + identify?: typeof identify; + + /** Check whether the connection to Discord is currently open. */ + isOpen?: typeof isOpen; + + /** Function which can be overwritten in order to get the shards presence. */ + makePresence?(shardId: number): Promise | StatusUpdate; + + /** The maximum of requests which can be send to discord per rate limit tick. + * Typically this value should not be changed. + */ + maxRequestsPerRateLimitTick?: number; + + /** The previous payload sequence number. */ + previousSequenceNumber?: number; + + /** In which interval (in milliseconds) the gateway resets it's rate limit. */ + rateLimitResetInterval?: number; + + /** Attempt to resume the previous shards session with the gateway. */ + resume?: typeof resume; + + /** Send a message to Discord. + * @param {boolean} [highPriority=false] - Whether this message should be send asap. + */ + send?: typeof send; + + /** Shutdown the shard. + * Forcefully disconnect the shard from Discord. + * The shard may not attempt to reconnect with Discord. + */ + shutdown?: typeof shutdown; + + /** @private Internal shard function. + * Only use this function if you know what you are doing. + * + * Start sending heartbeat payloads to Discord in the provided interval. + */ + startHeartbeating?: typeof startHeartbeating; + + /** Current internal state of the shard. */ + state?: ShardState; + + /** @private Internal shard function. + * Only use this function if you know what you are doing. + * + * Stop the heartbeating process with discord. + */ + stopHeartbeating?: typeof stopHeartbeating; + + /** The shard related event handlers. */ + events?: ShardEvents; + /** This contains all the heartbeat information */ + heart?: ShardHeart; + /** Bucket for handling shard request rate limits. */ + bucket?: LeakyBucket; + /** Cache for pending gateway requests which should have been send while the gateway went offline. */ + offlineSendQueue?: ShardSocketRequest[]; + /** This is used to resolve internal waiting states. + * Mapped by SelectedEvents => ResolveFunction + */ + resolves?: Shard["resolves"]; +} diff --git a/gateway/deps.ts b/gateway/shard/deps.ts similarity index 100% rename from gateway/deps.ts rename to gateway/shard/deps.ts diff --git a/gateway/shard/handleClose.ts b/gateway/shard/handleClose.ts new file mode 100644 index 000000000..a5c9ee446 --- /dev/null +++ b/gateway/shard/handleClose.ts @@ -0,0 +1,68 @@ +import { GatewayCloseEventCodes } from "../../types/shared.ts"; +import { Shard, ShardSocketCloseCodes, ShardState } from "./types.ts"; + +export async function handleClose(shard: Shard, close: CloseEvent): Promise { + // gateway.debug("GW CLOSED", { shardId, payload: event }); + + shard.stopHeartbeating(); + + console.log({ id: shard.id, close }); + + switch (close.code) { + case ShardSocketCloseCodes.Shutdown: { + throw new Error("SHUTDOWN"); + } + case ShardSocketCloseCodes.TestingFinished: { + shard.state = ShardState.Offline; + shard.events.disconnected?.(shard); + + return; + } + // On these codes a manual start will be done. + case ShardSocketCloseCodes.ReIdentifying: + case ShardSocketCloseCodes.Resharded: + case ShardSocketCloseCodes.ResumeClosingOldConnection: + case ShardSocketCloseCodes.ZombiedConnection: + case ShardSocketCloseCodes.Shutdown: { + shard.state = ShardState.Disconnected; + shard.events.disconnected?.(shard); + + // gateway.debug("GW CLOSED_RECONNECT", { shardId, payload: event }); + return; + } + // Gateway connection closes which require a new identify. + case GatewayCloseEventCodes.UnknownOpcode: + case GatewayCloseEventCodes.NotAuthenticated: + case GatewayCloseEventCodes.InvalidSeq: + case GatewayCloseEventCodes.RateLimited: + case GatewayCloseEventCodes.SessionTimedOut: { + shard.state = ShardState.Identifying; + shard.events.disconnected?.(shard); + + return await shard.identify(); + } + // When these codes are received something went really wrong. + // On those we cannot start a reconnect attempt. + case GatewayCloseEventCodes.AuthenticationFailed: + case GatewayCloseEventCodes.InvalidShard: + case GatewayCloseEventCodes.ShardingRequired: + case GatewayCloseEventCodes.InvalidApiVersion: + case GatewayCloseEventCodes.InvalidIntents: + case GatewayCloseEventCodes.DisallowedIntents: { + shard.state = ShardState.Offline; + shard.events.disconnected?.(shard); + + throw new Error(close.reason || "Discord gave no reason! GG! You broke Discord!"); + } + // Gateway connection closes on which a resume is allowed. + case GatewayCloseEventCodes.UnknownError: + case GatewayCloseEventCodes.DecodeError: + case GatewayCloseEventCodes.AlreadyAuthenticated: + default: { + shard.state = ShardState.Resuming; + shard.events.disconnected?.(shard); + + return await shard.resume(); + } + } +} diff --git a/gateway/shard/handleMessage.ts b/gateway/shard/handleMessage.ts new file mode 100644 index 000000000..a50c9e8db --- /dev/null +++ b/gateway/shard/handleMessage.ts @@ -0,0 +1,156 @@ +import { DiscordGatewayPayload, DiscordHello, DiscordReady } from "../../types/discord.ts"; +import { GatewayOpcodes } from "../../types/shared.ts"; +import { createLeakyBucket } from "../../util/bucket.ts"; +import { delay } from "../../util/utils.ts"; +import { decompressWith } from "./deps.ts"; +import { GATEWAY_RATE_LIMIT_RESET_INTERVAL, Shard, ShardState } from "./types.ts"; + +const decoder = new TextDecoder(); + +export async function handleMessage(shard: Shard, message: MessageEvent): Promise { + message = message.data; + + // If message compression is enabled, + // Discord might send zlib compressed payloads. + if (shard.gatewayConfig.compress && message instanceof Blob) { + message = decompressWith( + new Uint8Array(await message.arrayBuffer()), + 0, + (slice: Uint8Array) => decoder.decode(slice), + ); + } + + // Safeguard incase decompression failed to make a string. + if (typeof message !== "string") return; + + const messageData = JSON.parse(message) as DiscordGatewayPayload; + // gateway.debug("GW RAW", { shardId, payload: messageData }); + + // TODO: remove + // console.log({ messageData: censor(messageData) }); + + switch (messageData.op) { + case GatewayOpcodes.Heartbeat: { + // TODO: can this actually happen + if (!shard.isOpen()) return; + + shard.heart.lastBeat = Date.now(); + // Discord randomly sends this requiring an immediate heartbeat back. + // Using a direct socket.send call here because heartbeat requests are reserved by us. + shard.socket?.send( + JSON.stringify({ + op: GatewayOpcodes.Heartbeat, + d: shard.previousSequenceNumber, + }), + ); + shard.events.heartbeat?.(shard); + + break; + } + case GatewayOpcodes.Hello: { + const interval = (messageData.d as DiscordHello).heartbeat_interval; + + shard.startHeartbeating(interval); + + if (shard.state !== ShardState.Resuming) { + // HELLO has been send on a non resume action. + // This means that the shard starts a new session, + // therefore the rate limit interval has been reset too. + shard.bucket = createLeakyBucket({ + max: shard.calculateSafeRequests(), + refillInterval: GATEWAY_RATE_LIMIT_RESET_INTERVAL, + refillAmount: shard.calculateSafeRequests(), + // Waiting acquires should not be lost on a re-identify. + waiting: shard.bucket.waiting, + }); + } + + shard.events.hello?.(shard); + + break; + } + case GatewayOpcodes.HeartbeatACK: { + shard.heart.acknowledged = true; + shard.heart.lastAck = Date.now(); + // Manually calculating the round trip time for users who need it. + if (shard.heart.lastBeat) { + shard.heart.rtt = shard.heart.lastAck - shard.heart.lastBeat; + } + + shard.events.heartbeatAck?.(shard); + + break; + } + case GatewayOpcodes.Reconnect: { + // gateway.debug("GW RECONNECT", { shardId }); + + shard.events.requestedReconnect?.(shard); + + await shard.resume(); + + break; + } + case GatewayOpcodes.InvalidSession: { + // gateway.debug("GW INVALID_SESSION", { shardId, payload: messageData }); + const resumable = messageData.d as boolean; + + shard.events.invalidSession?.(shard, resumable); + + // We need to wait for a random amount of time between 1 and 5 + // Reference: https://discord.com/developers/docs/topics/gateway#resuming + await delay(Math.floor((Math.random() * 4 + 1) * 1000)); + + shard.resolves.get("INVALID_SESSION")?.(messageData); + shard.resolves.delete("INVALID_SESSION"); + + // When resumable is false we need to re-identify + if (!resumable) { + await shard.identify(); + + break; + } + + // The session is invalid but apparently it is resumable + await shard.resume(); + + break; + } + } + + if (messageData.t === "RESUMED") { + // gateway.debug("GW RESUMED", { shardId }); + + shard.state = ShardState.Connected; + shard.events.resumed?.(shard); + + // Continue the requests which have been queued since the shard went offline. + shard.offlineSendQueue.map((resolve) => resolve()); + + shard.resolves.get("RESUMED")?.(messageData); + shard.resolves.delete("RESUMED"); + } // Important for future resumes. + else if (messageData.t === "READY") { + const payload = messageData.d as DiscordReady; + + shard.sessionId = payload.session_id; + shard.state = ShardState.Connected; + + // Continue the requests which have been queued since the shard went offline. + // Important when this is a re-identify + shard.offlineSendQueue.map((resolve) => resolve()); + + shard.resolves.get("READY")?.(messageData); + shard.resolves.delete("READY"); + } + + // Update the sequence number if it is present + // `s` can be either `null` or a `number`. + // In order to prevent update misses when `s` is `0` we check against null. + if (messageData.s !== null) { + shard.previousSequenceNumber = messageData.s; + } + + // The necessary handling required for the Shards connection has been finished. + // Now the event can be safely forwarded. + shard.events.message?.(shard, messageData); +} diff --git a/gateway/shard/identify.ts b/gateway/shard/identify.ts new file mode 100644 index 000000000..6cd75897a --- /dev/null +++ b/gateway/shard/identify.ts @@ -0,0 +1,50 @@ +import { GatewayOpcodes } from "../../types/shared.ts"; +import { Shard, ShardSocketCloseCodes, ShardState } from "./types.ts"; + +export async function identify(shard: Shard): Promise { + // A new identify has been requested even though there is already a connection open. + // Therefore we need to close the old connection and heartbeating before creating a new one. + if (shard.state === ShardState.Connected) { + console.log("CLOSING EXISTING SHARD: #" + shard.id); + shard.close(ShardSocketCloseCodes.ReIdentifying, "Re-identifying closure of old connection."); + } + + shard.state = ShardState.Identifying; + shard.events.identifying?.(shard); + + // It is possible that the shard is in Heartbeating state but not identified, + // so check whether there is already a gateway connection existing. + // If not we need to create one before we identify. + if (!shard.isOpen()) { + await shard.connect(); + } + + // Wait until an identify is free for this shard. + await shard.requestIdentify(); + + shard.send({ + op: GatewayOpcodes.Identify, + d: { + token: `Bot ${shard.gatewayConfig.token}`, + compress: shard.gatewayConfig.compress, + properties: shard.gatewayConfig.properties, + intents: shard.gatewayConfig.intents, + shard: [shard.id, shard.totalShards], + presence: await shard.makePresence?.(shard.id), + }, + }, true); + + return new Promise((resolve) => { + shard.resolves.set("READY", () => { + shard.events.identified?.(shard); + resolve(); + }); + // When identifying too fast, + // Discord sends an invalid session payload. + // This can safely be ignored though and the shard starts a new identify action. + shard.resolves.set("INVALID_SESSION", () => { + shard.resolves.delete("READY"); + resolve(); + }); + }); +} diff --git a/gateway/shard/isOpen.ts b/gateway/shard/isOpen.ts new file mode 100644 index 000000000..c5bb149fc --- /dev/null +++ b/gateway/shard/isOpen.ts @@ -0,0 +1,5 @@ +import { Shard } from "./types.ts"; + +export function isOpen(shard: Shard): boolean { + return shard.socket?.readyState === WebSocket.OPEN; +} diff --git a/gateway/shard/mod.ts b/gateway/shard/mod.ts new file mode 100644 index 000000000..753770e50 --- /dev/null +++ b/gateway/shard/mod.ts @@ -0,0 +1,14 @@ +export * from "./calculateSafeRequests.ts"; +export * from "./close.ts"; +export * from "./connect.ts"; +export * from "./createShard.ts"; +export * from "./handleClose.ts"; +export * from "./handleMessage.ts"; +export * from "./identify.ts"; +export * from "./isOpen.ts"; +export * from "./resume.ts"; +export * from "./send.ts"; +export * from "./shutdown.ts"; +export * from "./startHeartbeating.ts"; +export * from "./stopHeartbeating.ts"; +export * from "./types.ts"; diff --git a/gateway/shard/resume.ts b/gateway/shard/resume.ts new file mode 100644 index 000000000..0739c90db --- /dev/null +++ b/gateway/shard/resume.ts @@ -0,0 +1,48 @@ +import { GatewayOpcodes } from "../../types/shared.ts"; +import { Shard, ShardSocketCloseCodes, ShardState } from "./types.ts"; + +export async function resume(shard: Shard): Promise { + // gateway.debug("GW RESUMING", { shardId }); + // It has been requested to resume the Shards session. + // It's possible that the shard is still connected with Discord's gateway therefore we need to forcefully close it. + if (shard.isOpen()) { + shard.close(ShardSocketCloseCodes.ResumeClosingOldConnection, "Reconnecting the shard, closing old connection."); + } + + // Shard has never identified, so we cannot resume. + if (!shard.sessionId) { + // gateway.debug( + // "GW DEBUG", + // `[Error] Trying to resume a shard (id: ${shardId}) that was not first identified.`, + // ); + + return await shard.identify(); + + // throw new Error(`[SHARD] Trying to resume a shard (id: ${shard.id}) which was never identified`); + } + + shard.state = ShardState.Resuming; + + // Before we can resume, we need to create a new connection with Discord's gateway. + await shard.connect(); + + shard.send({ + op: GatewayOpcodes.Resume, + d: { + token: `Bot ${shard.gatewayConfig.token}`, + session_id: shard.sessionId, + seq: shard.previousSequenceNumber ?? 0, + }, + }, true); + + return new Promise((resolve) => { + shard.resolves.set("RESUMED", () => resolve()); + // If it is attempted to resume with an invalid session id, + // Discord sends an invalid session payload + // Not erroring here since it is easy that this happens, also it would be not catchable + shard.resolves.set("INVALID_SESSION", () => { + shard.resolves.delete("RESUMED"); + resolve(); + }); + }); +} diff --git a/gateway/shard/send.ts b/gateway/shard/send.ts new file mode 100644 index 000000000..4adc3a7e3 --- /dev/null +++ b/gateway/shard/send.ts @@ -0,0 +1,27 @@ +import { Shard, ShardSocketRequest } from "./types.ts"; + +async function checkOffline(shard: Shard, highPriority: boolean): Promise { + if (!shard.isOpen()) { + await new Promise((resolve) => { + if (highPriority) { + // Higher priority requests get added at the beginning of the array. + shard.offlineSendQueue.unshift(resolve); + } else { + shard.offlineSendQueue.push(resolve); + } + }); + } +} + +export async function send(shard: Shard, message: ShardSocketRequest, highPriority: boolean): Promise { + // Before acquiring a token from the bucket, check whether the shard is currently offline or not. + // Else bucket and token wait time just get wasted. + await checkOffline(shard, highPriority); + + await shard.bucket.acquire(1, highPriority); + + // It's possible, that the shard went offline after a token has been acquired from the bucket. + await checkOffline(shard, highPriority); + + shard.socket?.send(JSON.stringify(message)); +} diff --git a/gateway/shard/shutdown.ts b/gateway/shard/shutdown.ts new file mode 100644 index 000000000..13b9d8a84 --- /dev/null +++ b/gateway/shard/shutdown.ts @@ -0,0 +1,6 @@ +import { Shard, ShardSocketCloseCodes, ShardState } from "./types.ts"; + +export async function shutdown(shard: Shard): Promise { + shard.close(ShardSocketCloseCodes.Shutdown, "Shard shutting down."); + shard.state = ShardState.Offline; +} diff --git a/gateway/shard/startHeartbeating.ts b/gateway/shard/startHeartbeating.ts new file mode 100644 index 000000000..8ee804f53 --- /dev/null +++ b/gateway/shard/startHeartbeating.ts @@ -0,0 +1,64 @@ +import { GatewayOpcodes } from "../../types/shared.ts"; +import { Shard, ShardSocketCloseCodes, ShardState } from "./types.ts"; + +export function startHeartbeating(shard: Shard, interval: number) { + // gateway.debug("GW HEARTBEATING_STARTED", { shardId, interval }); + + shard.heart.interval = interval; + + // Only set the shard's state to `Unidentified` + // if heartbeating has not been started due to an identify or resume action. + if ([ShardState.Disconnected, ShardState.Offline].includes(shard.state)) { + shard.state = ShardState.Unidentified; + } + + // The first heartbeat needs to be send with a random delay between `0` and `interval` + // Using a `setTimeout(_, jitter)` here to accomplish that. + // `Math.random()` can be `0` so we use `0.5` if this happens + // Reference: https://discord.com/developers/docs/topics/gateway#heartbeating + const jitter = Math.ceil(shard.heart.interval * (Math.random() || 0.5)); + shard.heart.timeoutId = setTimeout(() => { + // Using a direct socket.send call here because heartbeat requests are reserved by us. + shard.socket?.send(JSON.stringify({ + op: GatewayOpcodes.Heartbeat, + d: shard.previousSequenceNumber, + })); + + shard.heart.lastBeat = Date.now(); + shard.heart.acknowledged = false; + + // After the random heartbeat jitter we can start a normal interval. + shard.heart.intervalId = setInterval(async () => { + // gateway.debug("GW DEBUG", `Running setInterval in heartbeat file. Shard: ${shardId}`); + + // gateway.debug("GW HEARTBEATING", { shardId, shard: currentShard }); + + // The Shard did not receive a heartbeat ACK from Discord in time, + // therefore we have to assume that the connection has failed or got "zombied". + // The Shard needs to start a re-identify action accordingly. + // Reference: https://discord.com/developers/docs/topics/gateway#heartbeating-example-gateway-heartbeat-ack + if (!shard.heart.acknowledged) { + shard.close( + ShardSocketCloseCodes.ZombiedConnection, + "Zombied connection, did not receive an heartbeat ACK in time.", + ); + + return await shard.identify(); + } + + shard.heart.acknowledged = false; + + // Using a direct socket.send call here because heartbeat requests are reserved by us. + shard.socket?.send( + JSON.stringify({ + op: GatewayOpcodes.Heartbeat, + d: shard.previousSequenceNumber, + }), + ); + + shard.heart.lastBeat = Date.now(); + + shard.events.heartbeat?.(shard); + }, shard.heart.interval); + }, jitter); +} diff --git a/gateway/shard/stopHeartbeating.ts b/gateway/shard/stopHeartbeating.ts new file mode 100644 index 000000000..f22739ac1 --- /dev/null +++ b/gateway/shard/stopHeartbeating.ts @@ -0,0 +1,9 @@ +import { Shard } from "./types.ts"; + +export function stopHeartbeating(shard: Shard): void { + // Clear the regular heartbeat interval. + clearInterval(shard.heart.intervalId); + // It's possible that the Shard got closed before the first jittered heartbeat. + // To go safe we should clear the related timeout too. + clearTimeout(shard.heart.timeoutId); +} diff --git a/gateway/shard/types.ts b/gateway/shard/types.ts new file mode 100644 index 000000000..c0f8dd576 --- /dev/null +++ b/gateway/shard/types.ts @@ -0,0 +1,152 @@ +import { StatusUpdate } from "../../helpers/misc/editShardStatus.ts"; +import { DiscordGatewayPayload } from "../../types/discord.ts"; +import { GatewayOpcodes } from "../../types/shared.ts"; +import { LeakyBucket } from "../../util/bucket.ts"; +import { createShard } from "./createShard.ts"; + +// TODO: think whether we also need an identifiedShard function + +export const MAX_GATEWAY_REQUESTS_PER_INTERVAL = 120; +export const GATEWAY_RATE_LIMIT_RESET_INTERVAL = 60_000; // 60 seconds +export const DEFAULT_HEARTBEAT_INTERVAL = 45000; + +export type Shard = ReturnType; + +export enum ShardState { + /** Shard is fully connected to the gateway and receiving events from Discord. */ + Connected = 0, + /** Shard started to connect to the gateway. + * This is only used if the shard is not currently trying to identify or resume. + */ + Connecting = 1, + /** Shard got disconnected and reconnection actions have been started. */ + Disconnected = 2, + /** The shard is connected to the gateway but only heartbeating. + * At this state the shard has not been identified with discord. + */ + Unidentified = 3, + /** Shard is trying to identify with the gateway to create a new session. */ + Identifying = 4, + /** Shard is trying to resume a session with the gateway. */ + Resuming = 5, + /** Shard got shut down studied or due to a not (self) fixable error and may not attempt to reconnect on its own. */ + Offline = 6, +} + +export interface ShardGatewayConfig { + /** Whether incoming payloads are compressed using zlib. + * + * @default false + */ + compress: boolean; + /** The calculated intent value of the events which the shard should receive. + * + * @default 0 + */ + intents: number; + /** Identify properties to use */ + properties: { + /** Operating system the shard runs on. + * + * @default "darwin" | "linux" | "windows" + */ + $os: string; + /** The "browser" where this shard is running on. + * + * @default "Discordeno" + */ + $browser: string; + /** The device on which the shard is running. + * + * @default "Discordeno" + */ + $device: string; + }; + /** Bot token which is used to connect to Discord */ + token: string; + /** The URL of the gateway which should be connected to. + * + * @default "wss://gateway.discord.gg" + */ + url: string; + /** The gateway version which should be used. + * + * @default 10 + */ + version: number; +} + +export interface ShardHeart { + /** Whether or not the heartbeat was acknowledged by Discord in time. */ + acknowledged: boolean; + /** Interval between heartbeats requested by Discord. */ + interval: number; + /** Id of the interval, which is used for sending the heartbeats. */ + intervalId?: number; + /** Unix (in milliseconds) timestamp when the last heartbeat ACK was received from Discord. */ + lastAck?: number; + /** Unix timestamp (in milliseconds) when the last heartbeat was sent. */ + lastBeat?: number; + /** Round trip time (in milliseconds) from Shard to Discord and back. + * Calculated using the heartbeat system. + * Note: this value is undefined until the first heartbeat to Discord has happened. + */ + rtt?: number; + /** Id of the timeout which is used for sending the first heartbeat to Discord since it's "special". */ + timeoutId?: number; +} + +export interface ShardEvents { + /** A heartbeat has been send. */ + heartbeat?(shard: Shard): unknown; + /** A heartbeat ACK was received. */ + heartbeatAck?(shard: Shard): unknown; + /** Shard has received a Hello payload. */ + hello?(shard: Shard): unknown; + /** The Shards session has been invalidated. */ + invalidSession?(shard: Shard, resumable: boolean): unknown; + /** The shard has started a resume action. */ + resuming?(shard: Shard): unknown; + /** The shard has successfully resumed an old session. */ + resumed?(shard: Shard): unknown; + /** Discord has requested the Shard to reconnect. */ + requestedReconnect?(shard: Shard): unknown; + /** The shard started to connect to Discord's gateway. */ + connecting?(shard: Shard): unknown; + /** The shard is connected with Discord's gateway. */ + connected?(shard: Shard): unknown; + /** The shard has been disconnected from Discord's gateway. */ + disconnected?(shard: Shard): unknown; + /** The shard has started to identify itself to Discord. */ + identifying?(shard: Shard): unknown; + /** The shard has successfully been identified itself with Discord. */ + identified?(shard: Shard): unknown; + /** The shard has received a message from Discord. */ + message?(shard: Shard, payload: DiscordGatewayPayload): unknown; +} + +export enum ShardSocketCloseCodes { + /** A regular Shard shutdown. + * Discord will display this Shard as offline for other users. + */ + Shutdown = 1000, + /** A resume has been requested and therefore the old connection needs to be closed. */ + ResumeClosingOldConnection = 3024, + /** Did not receive a heartbeat ACK in time. + * Closing the shard and creating a new session. + */ + ZombiedConnection = 3010, + /** Discordeno's gateway tests hae been finished, therefore the Shard can be turned off. */ + TestingFinished = 3064, + /** Special close code reserved for Discordeno's zero-downtime resharding system. */ + Resharded = 3065, + /** Shard is re-identifying therefore the old connection needs to be closed. */ + ReIdentifying = 3066, +} + +export interface ShardSocketRequest { + /** The OP-Code for the payload to send. */ + op: GatewayOpcodes; + /** Payload data. */ + d: unknown; +} diff --git a/gateway/spawnShards.ts b/gateway/spawnShards.ts deleted file mode 100644 index ae3606019..000000000 --- a/gateway/spawnShards.ts +++ /dev/null @@ -1,64 +0,0 @@ -import { GatewayManager } from "./gatewayManager.ts"; - -export function prepareBuckets(gateway: GatewayManager, firstShardId: number, lastShardId: number) { - /** Stored as bucketId: [workerId, [ShardIds]] */ - let worker = 0; - - for (let i = 0; i < gateway.maxConcurrency; i++) { - gateway.buckets.set(i, { - workers: [], - createNextShard: [], - }); - } - - // ORGANIZE ALL SHARDS INTO THEIR OWN BUCKETS - for (let i = firstShardId; i <= lastShardId; i++) { - gateway.debug("GW DEBUG", `1. Running for loop in spawnShards function for shardId ${i}.`); - if (i >= gateway.maxShards) { - continue; - } - - const bucketId = i % gateway.maxConcurrency; - const bucket = gateway.buckets.get(bucketId); - if (!bucket) throw new Error("Bucket not found when spawning shards."); - - // FIND A QUEUE IN THIS BUCKET THAT HAS SPACE - // + 1 cause .workers first item is worker id [workerId, shardId, shardId2...] - const queue = bucket.workers.find((q) => q.length < gateway.shardsPerWorker + 1); - if (queue) { - // IF THE QUEUE HAS SPACE JUST ADD IT TO THIS QUEUE - queue.push(i); - } else { - if (worker + 1 <= gateway.maxWorkers) worker++; - // ADD A NEW QUEUE FOR THIS SHARD - bucket.workers.push([worker, i]); - } - } -} - -/** Begin spawning shards. */ -export function spawnShards(gateway: GatewayManager, firstShardId = 0) { - // PREPARES THE MAX SHARD COUNT BY CONCURRENCY - if (gateway.useOptimalLargeBotSharding) { - gateway.debug("GW DEBUG", "[Spawning] Using optimal large bot sharding solution."); - gateway.maxShards = gateway.calculateMaxShards(gateway.maxShards, gateway.maxConcurrency); - } - - // PREPARES ALL SHARDS IN SPECIFIC BUCKETS - prepareBuckets(gateway, firstShardId, gateway.lastShardId ? gateway.lastShardId : gateway.maxShards - 1); - - // SPREAD THIS OUT TO DIFFERENT WORKERS TO BEGIN STARTING UP - gateway.buckets.forEach(async (bucket, bucketId) => { - gateway.debug("GW DEBUG", `2. Running forEach loop in spawnShards function.`); - for (const [workerId, ...queue] of bucket.workers) { - gateway.debug("GW DEBUG", `3. Running for of loop in spawnShards function.`); - - for (const shardId of queue) { - bucket.createNextShard.push(async () => { - await gateway.tellWorkerToIdentify(gateway, workerId, shardId, bucketId); - }); - } - } - await bucket.createNextShard.shift()?.(); - }); -} diff --git a/gateway/stopGateway.ts b/gateway/stopGateway.ts deleted file mode 100644 index 058ae4591..000000000 --- a/gateway/stopGateway.ts +++ /dev/null @@ -1,21 +0,0 @@ -import { delay } from "../util/utils.ts"; -import { GatewayManager } from "./gatewayManager.ts"; - -/** Use this function to stop the gateway properly */ -export async function stopGateway( - gateway: GatewayManager, - code = 3061, - reason = "Discordeno Testing Finished! Do Not RESUME!", -) { - // STOP WS - gateway.shards.forEach((shard) => { - clearInterval(shard.heartbeat.intervalId); - gateway.closeWS( - shard.ws, - code, - reason, - ); - }); - - await delay(5000); -} diff --git a/helpers/members/fetchMembers.ts b/helpers/members/fetchMembers.ts index 0db10ff25..ee7dc3a5b 100644 --- a/helpers/members/fetchMembers.ts +++ b/helpers/members/fetchMembers.ts @@ -26,7 +26,12 @@ export function fetchMembers( const nonce = `${guildId}-${Date.now()}`; bot.cache.fetchAllMembersProcessingRequests.set(nonce, resolve); - bot.gateway.sendShardMessage(bot.gateway, shardId, { + const shard = bot.gateway.manager.shards.get(shardId); + if (!shard) { + throw new Error(`Shard (id: ${shardId}) not found.`); + } + + shard.send({ op: GatewayOpcodes.RequestGuildMembers, d: { guild_id: guildId.toString(), diff --git a/helpers/misc/editBotStatus.ts b/helpers/misc/editBotStatus.ts index 82f940a77..99448c3b8 100644 --- a/helpers/misc/editBotStatus.ts +++ b/helpers/misc/editBotStatus.ts @@ -1,77 +1,6 @@ -import type { Bot } from "../../bot.ts"; -import { Activity } from "../../transformers/activity.ts"; -import { StatusTypes } from "../../transformers/presence.ts"; -import { GatewayOpcodes } from "../../types/shared.ts"; +import { Bot } from "../../bot.ts"; +import { StatusUpdate } from "./editShardStatus.ts"; -export function editBotStatus(bot: Bot, data: StatusUpdate) { - bot.gateway.shards.forEach((shard) => { - bot.events.debug(`Running forEach loop in editBotStatus function.`); - - bot.gateway.sendShardMessage(bot.gateway, shard, { - op: GatewayOpcodes.PresenceUpdate, - d: { - since: null, - afk: false, - activities: data.activities.map((activity) => ({ - name: activity.name, - type: activity.type, - url: activity.url, - created_at: activity.createdAt, - timestamps: activity.startedAt || activity.endedAt - ? { - start: activity.startedAt, - end: activity.endedAt, - } - : undefined, - application_id: activity.applicationId?.toString(), - details: activity.details, - state: activity.state, - emoji: activity.emoji - ? { - name: activity.emoji.name, - id: activity.emoji.id?.toString(), - animated: activity.emoji.animated, - } - : undefined, - party: activity.partyId - ? { - id: activity.partyId.toString(), - size: activity.partyMaxSize, - } - : undefined, - assets: activity.largeImage || activity.largeText || activity.smallImage || activity.smallText - ? { - large_image: activity.largeImage, - large_text: activity.largeText, - small_image: activity.smallImage, - small_text: activity.smallText, - } - : undefined, - secrets: activity.join || activity.spectate || activity.match - ? { - join: activity.join, - spectate: activity.spectate, - match: activity.match, - } - : undefined, - instance: activity.instance, - flags: activity.flags, - buttons: activity.buttons, - })), - status: data.status, - }, - }); - }); -} - -/** https://discord.com/developers/docs/topics/gateway#update-status */ -export interface StatusUpdate { - // /** Unix time (in milliseconds) of when the client went idle, or null if the client is not idle */ - // since: number | null; - /** The user's activities */ - activities: Activity[]; - /** The user's new status */ - status: StatusTypes; - // /** Whether or not the client is afk */ - // afk: boolean; +export async function editBotStatus(bot: Bot, data: StatusUpdate) { + await Promise.all(bot.gateway.manager.shards.map((shard) => bot.helpers.editShardStatus(shard.id, data))); } diff --git a/helpers/misc/editShardStatus.ts b/helpers/misc/editShardStatus.ts new file mode 100644 index 000000000..3daa23f8e --- /dev/null +++ b/helpers/misc/editShardStatus.ts @@ -0,0 +1,78 @@ +import type { Bot } from "../../bot.ts"; +import { Activity } from "../../transformers/activity.ts"; +import { StatusTypes } from "../../transformers/presence.ts"; +import { GatewayOpcodes } from "../../types/shared.ts"; + +export function editShardStatus(bot: Bot, shardId: number, data: StatusUpdate) { + const shard = bot.gateway.manager.shards.get(shardId); + if (!shard) { + throw new Error(`Shard (id: ${shardId}) not found.`); + } + + shard.send({ + op: GatewayOpcodes.PresenceUpdate, + d: { + since: null, + afk: false, + activities: data.activities.map((activity) => ({ + name: activity.name, + type: activity.type, + url: activity.url, + created_at: activity.createdAt, + timestamps: activity.startedAt || activity.endedAt + ? { + start: activity.startedAt, + end: activity.endedAt, + } + : undefined, + application_id: activity.applicationId?.toString(), + details: activity.details, + state: activity.state, + emoji: activity.emoji + ? { + name: activity.emoji.name, + id: activity.emoji.id?.toString(), + animated: activity.emoji.animated, + } + : undefined, + party: activity.partyId + ? { + id: activity.partyId.toString(), + size: activity.partyMaxSize, + } + : undefined, + assets: activity.largeImage || activity.largeText || activity.smallImage || activity.smallText + ? { + large_image: activity.largeImage, + large_text: activity.largeText, + small_image: activity.smallImage, + small_text: activity.smallText, + } + : undefined, + secrets: activity.join || activity.spectate || activity.match + ? { + join: activity.join, + spectate: activity.spectate, + match: activity.match, + } + : undefined, + instance: activity.instance, + flags: activity.flags, + buttons: activity.buttons, + })), + status: data.status, + }, + }); +} + +/** https://discord.com/developers/docs/topics/gateway#update-status */ +export interface StatusUpdate { + // /** Unix time (in milliseconds) of when the client went idle, or null if the client is not idle */ + // since: number | null; + /** The user's activities */ + activities: Activity[]; + /** The user's new status */ + status: StatusTypes; + // /** Whether or not the client is afk */ + // afk: boolean; +} diff --git a/helpers/misc/mod.ts b/helpers/misc/mod.ts index 7321589ea..f77958501 100644 --- a/helpers/misc/mod.ts +++ b/helpers/misc/mod.ts @@ -1,5 +1,6 @@ export * from "./editBotProfile.ts"; export * from "./editBotStatus.ts"; +export * from "./editShardStatus.ts"; export * from "./getGatewayBot.ts"; export * from "./getUser.ts"; export * from "./nitroStickerPacks.ts"; diff --git a/helpers/voice/connectToVoiceChannel.ts b/helpers/voice/connectToVoiceChannel.ts index 9ab881fc3..f96d0e1fe 100644 --- a/helpers/voice/connectToVoiceChannel.ts +++ b/helpers/voice/connectToVoiceChannel.ts @@ -8,7 +8,13 @@ export async function connectToVoiceChannel( channelId: bigint, options?: AtLeastOne>, ) { - bot.gateway.sendShardMessage(bot.gateway, bot.utils.calculateShardId(bot.gateway, guildId), { + const shardId = bot.utils.calculateShardId(bot.gateway, guildId); + const shard = bot.gateway.manager.shards.get(shardId); + if (!shard) { + throw new Error(`Shard (id: ${shardId} not found`); + } + + shard.send({ op: GatewayOpcodes.VoiceStateUpdate, d: { guild_id: guildId.toString(), diff --git a/types/shared.ts b/types/shared.ts index 51f3b3a73..6a742d040 100644 --- a/types/shared.ts +++ b/types/shared.ts @@ -1388,3 +1388,5 @@ export type PickPartial = [P in keyof T]?: T[P] | undefined; } & { [P in K]: T[P] }; + +export type OmitFirstFnArg = F extends (x: any, ...args: infer P) => infer R ? (...args: P) => R : never; diff --git a/util/bucket.ts b/util/bucket.ts new file mode 100644 index 000000000..5429c4416 --- /dev/null +++ b/util/bucket.ts @@ -0,0 +1,175 @@ +import { PickPartial } from "../types/shared.ts"; +import { delay } from "./utils.ts"; + +/** A Leaky Bucket. + * Useful for rate limiting purposes. + * This uses `performance.now()` instead of `Date.now()` for higher accuracy. + * + * NOTE: This bucket is lazy, means it only updates when a related method is called. + */ +export interface LeakyBucket { + // ---------- + // PROPERTIES + // ---------- + + /** How many tokens this bucket can hold. */ + max: number; + /** Amount of tokens gained per interval. + * If bigger than `max` it will be pressed to `max`. + */ + refillAmount: number; + /** Interval at which the bucket gains tokens. */ + refillInterval: number; + + // ---------- + // METHODS + // ---------- + + /** Acquire tokens from the bucket. + * Resolves when the tokens are acquired and available. + * @param {boolean} [highPriority=false] Whether this acquire is should be done asap. + */ + acquire(amount: number, highPriority?: boolean): Promise; + + /** Returns the number of milliseconds until the next refill. */ + nextRefill(): number; + + /** Current tokens in the bucket. */ + tokens(): number; + + // ---------- + // INTERNAL STATES + // ---------- + + /** @private Internal track of when the last refill of tokens was. + * DO NOT TOUCH THIS! Unless you know what you are doing ofc :P + */ + lastRefill: number; + + /** @private Internal state of whether currently it is allowed to acquire tokens. + * DO NOT TOUCH THIS! Unless you know what you are doing ofc :P + */ + allowAcquire: boolean; + + /** @private Internal number of currently available tokens. + * DO NOT TOUCH THIS! Unless you know what you are doing ofc :P + */ + tokensState: number; + + /** @private Internal array of promises necessary to guarantee no race conditions. + * DO NOT TOUCH THIS! Unless you know what you are doing ofc :P + */ + waiting: ((_?: unknown) => void)[]; +} + +export function createLeakyBucket( + { max, refillInterval, refillAmount, tokens, waiting, ...rest }: + & Omit< + PickPartial< + LeakyBucket, + "max" | "refillInterval" | "refillAmount" + >, + "tokens" + > + & { + /** Current tokens in the bucket. + * @default max + */ + tokens?: number; + }, +): LeakyBucket { + return { + max, + refillInterval, + refillAmount: refillAmount > max ? max : refillAmount, + lastRefill: performance.now(), + allowAcquire: true, + + nextRefill: function () { + return nextRefill(this); + }, + + tokens: function () { + return updateTokens(this); + }, + + acquire: async function (amount, highPriority) { + return await acquire(this, amount, highPriority); + }, + + tokensState: tokens ?? max, + waiting: waiting ?? [], + + ...rest, + }; +} + +/** Update the tokens of that bucket. + * @returns {number} The amount of current available tokens. + */ +function updateTokens(bucket: LeakyBucket): number { + const timePassed = performance.now() - bucket.lastRefill; + const missedRefills = Math.floor(timePassed / bucket.refillInterval); + + // The refill shall not exceed the max amount of tokens. + bucket.tokensState = Math.min(bucket.tokensState + (bucket.refillAmount * missedRefills), bucket.max); + bucket.lastRefill += bucket.refillInterval * missedRefills; + + return bucket.tokensState; +} + +function nextRefill(bucket: LeakyBucket): number { + // Since this bucket is lazy update the tokens before calculating the next refill. + updateTokens(bucket); + + return (performance.now() - bucket.lastRefill) + bucket.refillInterval; +} + +async function acquire(bucket: LeakyBucket, amount: number, highPriority = false): Promise { + // To prevent the race condition of 2 acquires happening at once, + // check whether its currently allowed to acquire. + if (!bucket.allowAcquire) { + // create, push, and wait until the current running acquiring is finished. + await new Promise((resolve) => { + if (highPriority) { + bucket.waiting.unshift(resolve); + } else { + bucket.waiting.push(resolve); + } + }); + + // Somehow another acquire has started, + // so need to wait again. + if (!bucket.allowAcquire) { + return await acquire(bucket, amount); + } + } + + bucket.allowAcquire = false; + // Since the bucket is lazy update the tokens now, + // and also get the current amount of available tokens + let currentTokens = updateTokens(bucket); + + // It's possible that more than available tokens have been acquired, + // so calculate the amount of milliseconds to wait until this acquire is good to go. + if (currentTokens < amount) { + const tokensNeeded = amount - currentTokens; + let refillsNeeded = Math.ceil(tokensNeeded / bucket.refillAmount); + + const waitTime = bucket.refillInterval * refillsNeeded; + await delay(waitTime); + + // Update the tokens again to ensure nothing has been missed. + updateTokens(bucket); + } + + // In order to not subtract too much from the tokens, + // calculate what is actually needed to subtract. + const toSubtract = (amount % bucket.refillAmount) || amount; + bucket.tokensState -= toSubtract; + + // Allow the next acquire to happen. + bucket.allowAcquire = true; + // If there is an acquire waiting, let it continue. + bucket.waiting.shift()?.(); +} diff --git a/util/calculateShardId.ts b/util/calculateShardId.ts index a469c567f..820d2124e 100644 --- a/util/calculateShardId.ts +++ b/util/calculateShardId.ts @@ -1,7 +1,7 @@ -import { GatewayManager } from "../gateway/gatewayManager.ts"; +import { GatewayManager } from "../gateway/manager/gatewayManager.ts"; export function calculateShardId(gateway: GatewayManager, guildId: bigint) { - if (gateway.maxShards === 1) return 0; + if (gateway.manager.totalShards === 1) return 0; - return Number((guildId >> 22n) % BigInt(gateway.maxShards - 1)); + return Number((guildId >> 22n) % BigInt(gateway.manager.totalShards - 1)); }