From 471ef5cb6cc8e17f306e03d2ffdf88a0e01b498d Mon Sep 17 00:00:00 2001 From: Skillz4Killz <23035000+Skillz4Killz@users.noreply.github.com> Date: Thu, 10 Feb 2022 17:43:29 -0500 Subject: [PATCH] fix(gatway): resharder bug (#1979) * fix: resharder bug. TY giveawayboat 4 catching it * fix: make handler set properly once resharded * Update src/ws/resharder.ts * Update src/ws/resharder.ts * Update src/ws/resharder.ts * Update src/ws/resharder.ts * fix: allow flexible resharding for multi vps bots * fix: allow overriding on close * fix: separate resharding checks * fix: dupes from conflict * fix: requested changes * Update src/ws/resharder.ts Co-authored-by: ITOH * fix: gateway.resharding.x * fix: allow editing guilds cached with new shard id * fix: fmt * Update src/ws/gateway_manager.ts Co-authored-by: ITOH * Update src/ws/gateway_manager.ts Co-authored-by: ITOH * Update src/ws/resharder.ts Co-authored-by: ITOH * fix: use guildIds[] instead of per guildi d * fix: use og name * fix: dumb deno fmt Co-authored-by: ITOH --- src/bot.ts | 192 ++++++++++++++++++++++++++++---------- src/ws/gateway_manager.ts | 35 ++++++- src/ws/resharder.ts | 136 ++++++++++++++++++++++++--- 3 files changed, 295 insertions(+), 68 deletions(-) diff --git a/src/bot.ts b/src/bot.ts index 1ba57608c..88a2a85bd 100644 --- a/src/bot.ts +++ b/src/bot.ts @@ -73,7 +73,10 @@ export function createBot(options: CreateBotOptions): Bot { applicationId: options.applicationId || options.botId, token: options.token, events: createEventHandlers(options.events), - intents: options.intents.reduce((bits, next) => (bits |= GatewayIntents[next]), 0), + intents: options.intents.reduce( + (bits, next) => (bits |= GatewayIntents[next]), + 0, + ), botGatewayData: options.botGatewayData, activeGuildIds: new Set(), constants: createBotConstants(), @@ -107,14 +110,20 @@ export function createBot(options: CreateBotOptions): Bot { // RUN DISPATCH CHECK await bot.events.dispatchRequirements(bot as Bot, data, shardId); - bot.handlers[data.t as GatewayDispatchEventNames]?.(bot as Bot, data, shardId); + bot.handlers[data.t as GatewayDispatchEventNames]?.( + bot as Bot, + data, + shardId, + ); }, }); return bot as Bot; } -export function createEventHandlers(events: Partial): EventHandlers { +export function createEventHandlers( + events: Partial, +): EventHandlers { function ignore() {} return { @@ -175,7 +184,9 @@ export function createEventHandlers(events: Partial): EventHandle } export async function startBot(bot: Bot) { - if (!bot.botGatewayData) bot.botGatewayData = await bot.helpers.getGatewayBot(); + if (!bot.botGatewayData) { + bot.botGatewayData = await bot.helpers.getGatewayBot(); + } // SETUP GATEWAY LOGIN INFO bot.gateway.urlWSS = bot.botGatewayData.url; @@ -226,7 +237,11 @@ export async function stopBot(bot: Bot) { // STOP WS bot.gateway.shards.forEach((shard) => { clearInterval(shard.heartbeat.intervalId); - bot.gateway.closeWS(shard.ws, 3061, "Discordeno Testing Finished! Do Not RESUME!"); + bot.gateway.closeWS( + shard.ws, + 3061, + "Discordeno Testing Finished! Do Not RESUME!", + ); }); await delay(5000); @@ -249,7 +264,8 @@ export interface CreateBotOptions { helpers?: Partial; } -export type UnPromise> = T extends Promise ? K : never; +export type UnPromise> = T extends Promise ? K + : never; export interface Bot { id: bigint; @@ -280,11 +296,20 @@ export type DefaultHelpers = typeof defaultHelpers; // deno-lint-ignore no-empty-interface export interface Helpers extends DefaultHelpers {} // Use interface for declaration merging -export function createHelpers(bot: Bot, customHelpers?: Partial): FinalHelpers { +export function createHelpers( + bot: Bot, + customHelpers?: Partial, +): FinalHelpers { const converted = {} as FinalHelpers; - for (const [name, fun] of Object.entries({ ...createBaseHelpers(customHelpers || {}) })) { + for ( + const [name, fun] of Object.entries({ + ...createBaseHelpers(customHelpers || {}), + }) + ) { // @ts-ignore - TODO: make the types better - converted[name as keyof FinalHelpers] = (...args: RemoveFirstFromTuple>) => + converted[name as keyof FinalHelpers] = ( + ...args: RemoveFirstFromTuple> + ) => // @ts-ignore - TODO: make the types better fun(bot, ...args); } @@ -356,9 +381,12 @@ export function createTransformers(options: Partial) { snowflake: options.snowflake || snowflakeToBigint, webhook: options.webhook || transformWebhook, auditlogEntry: options.auditlogEntry || transformAuditlogEntry, - applicationCommand: options.applicationCommand || transformApplicationCommand, - applicationCommandOption: options.applicationCommandOption || transformApplicationCommandOption, - applicationCommandPermission: options.applicationCommandPermission || transformApplicationCommandPermission, + applicationCommand: options.applicationCommand || + transformApplicationCommand, + applicationCommandOption: options.applicationCommandOption || + transformApplicationCommandOption, + applicationCommandPermission: options.applicationCommandPermission || + transformApplicationCommandPermission, scheduledEvent: options.scheduledEvent || transformScheduledEvent, threadMember: options.threadMember || transformThreadMember, welcomeScreen: options.welcomeScreen || transformWelcomeScreen, @@ -421,7 +449,10 @@ export interface EventHandlers { ) => any; interactionCreate: (bot: Bot, interaction: DiscordenoInteraction) => any; integrationCreate: (bot: Bot, integration: DiscordenoIntegration) => any; - integrationDelete: (bot: Bot, payload: { id: bigint; guildId: bigint; applicationId?: bigint }) => any; + integrationDelete: ( + bot: Bot, + payload: { id: bigint; guildId: bigint; applicationId?: bigint }, + ) => any; integrationUpdate: (bot: Bot, payload: { guildId: bigint }) => any; inviteCreate: (bot: Bot, invite: DiscordenoInvite) => any; inviteDelete: ( @@ -432,16 +463,28 @@ export interface EventHandlers { code: string; }, ) => any; - guildMemberAdd: (bot: Bot, member: DiscordenoMember, user: DiscordenoUser) => any; + guildMemberAdd: ( + bot: Bot, + member: DiscordenoMember, + user: DiscordenoUser, + ) => any; guildMemberRemove: (bot: Bot, user: DiscordenoUser, guildId: bigint) => any; - guildMemberUpdate: (bot: Bot, member: DiscordenoMember, user: DiscordenoUser) => any; + guildMemberUpdate: ( + bot: Bot, + member: DiscordenoMember, + user: DiscordenoUser, + ) => any; messageCreate: (bot: Bot, message: DiscordenoMessage) => any; messageDelete: ( bot: Bot, payload: { id: bigint; channelId: bigint; guildId?: bigint }, message?: DiscordenoMessage, ) => any; - messageUpdate: (bot: Bot, message: DiscordenoMessage, oldMessage?: DiscordenoMessage) => any; + messageUpdate: ( + bot: Bot, + message: DiscordenoMessage, + oldMessage?: DiscordenoMessage, + ) => any; reactionAdd: ( bot: Bot, payload: { @@ -480,8 +523,15 @@ export interface EventHandlers { guildId?: bigint; }, ) => any; - presenceUpdate: (bot: Bot, presence: DiscordenoPresence, oldPresence?: DiscordenoPresence) => any; - voiceServerUpdate: (bot: Bot, payload: { token: string; endpoint?: string; guildId: bigint }) => any; + presenceUpdate: ( + bot: Bot, + presence: DiscordenoPresence, + oldPresence?: DiscordenoPresence, + ) => any; + voiceServerUpdate: ( + bot: Bot, + payload: { token: string; endpoint?: string; guildId: bigint }, + ) => any; voiceStateUpdate: ( bot: Bot, voiceState: { @@ -502,7 +552,11 @@ export interface EventHandlers { }, ) => any; channelCreate: (bot: Bot, channel: DiscordenoChannel) => any; - dispatchRequirements: (bot: Bot, data: GatewayPayload, shardId: number) => any; + dispatchRequirements: ( + bot: Bot, + data: GatewayPayload, + shardId: number, + ) => any; voiceChannelLeave: ( bot: Bot, voiceState: DiscordenoVoiceState, @@ -510,7 +564,10 @@ export interface EventHandlers { channel?: DiscordenoChannel, ) => any; channelDelete: (bot: Bot, channel: DiscordenoChannel) => any; - channelPinsUpdate: (bot: Bot, data: { guildId?: bigint; channelId: bigint; lastPinTimestamp?: number }) => any; + channelPinsUpdate: ( + bot: Bot, + data: { guildId?: bigint; channelId: bigint; lastPinTimestamp?: number }, + ) => any; channelUpdate: (bot: Bot, channel: DiscordenoChannel) => any; stageInstanceCreate: ( bot: Bot, @@ -557,7 +614,10 @@ export interface EventHandlers { roleCreate: (bot: Bot, role: DiscordenoRole) => any; roleDelete: (bot: Bot, payload: { guildId: bigint; roleId: bigint }) => any; roleUpdate: (bot: Bot, role: DiscordenoRole) => any; - webhooksUpdate: (bot: Bot, payload: { channelId: bigint; guildId: bigint }) => any; + webhooksUpdate: ( + bot: Bot, + payload: { channelId: bigint; guildId: bigint }, + ) => any; botUpdate: (bot: Bot, user: DiscordenoUser) => any; typingStart: ( bot: Bot, @@ -647,14 +707,18 @@ export interface BotGatewayHandlerOptions { export function createBotGatewayHandlers( options: Partial, -): Record any> { +): Record< + GatewayDispatchEventNames | "GUILD_LOADED_DD", + (bot: Bot, data: GatewayPayload, shardId: number) => any +> { return { // misc READY: options.READY ?? handlers.handleReady, // channels CHANNEL_CREATE: options.CHANNEL_CREATE ?? handlers.handleChannelCreate, CHANNEL_DELETE: options.CHANNEL_DELETE ?? handlers.handleChannelDelete, - CHANNEL_PINS_UPDATE: options.CHANNEL_PINS_UPDATE ?? handlers.handleChannelPinsUpdate, + CHANNEL_PINS_UPDATE: options.CHANNEL_PINS_UPDATE ?? + handlers.handleChannelPinsUpdate, CHANNEL_UPDATE: options.CHANNEL_UPDATE ?? handlers.handleChannelUpdate, // THREAD_CREATE: options.THREAD_CREATE ?? handlers.handleThreadCreate, // THREAD_UPDATE: options.THREAD_UPDATE ?? handlers.handleThreadUpdate, @@ -662,9 +726,12 @@ export function createBotGatewayHandlers( // THREAD_LIST_SYNC: options.THREAD_LIST_SYNC ?? handlers.handleThreadListSync, // THREAD_MEMBER_UPDATE: options.THREAD_MEMBER_UPDATE ?? handlers.handleThreadMemberUpdate, // THREAD_MEMBERS_UPDATE: options.THREAD_MEMBERS_UPDATE ?? handlers.handleThreadMembersUpdate, - STAGE_INSTANCE_CREATE: options.STAGE_INSTANCE_CREATE ?? handlers.handleStageInstanceCreate, - STAGE_INSTANCE_UPDATE: options.STAGE_INSTANCE_UPDATE ?? handlers.handleStageInstanceUpdate, - STAGE_INSTANCE_DELETE: options.STAGE_INSTANCE_DELETE ?? handlers.handleStageInstanceDelete, + STAGE_INSTANCE_CREATE: options.STAGE_INSTANCE_CREATE ?? + handlers.handleStageInstanceCreate, + STAGE_INSTANCE_UPDATE: options.STAGE_INSTANCE_UPDATE ?? + handlers.handleStageInstanceUpdate, + STAGE_INSTANCE_DELETE: options.STAGE_INSTANCE_DELETE ?? + handlers.handleStageInstanceDelete, // guilds GUILD_BAN_ADD: options.GUILD_BAN_ADD ?? handlers.handleGuildBanAdd, @@ -672,50 +739,73 @@ export function createBotGatewayHandlers( GUILD_CREATE: options.GUILD_CREATE ?? handlers.handleGuildCreate, GUILD_LOADED_DD: options.GUILD_LOADED_DD ?? handlers.handleGuildLoaded, GUILD_DELETE: options.GUILD_DELETE ?? handlers.handleGuildDelete, - GUILD_EMOJIS_UPDATE: options.GUILD_EMOJIS_UPDATE ?? handlers.handleGuildEmojisUpdate, - GUILD_INTEGRATIONS_UPDATE: options.GUILD_INTEGRATIONS_UPDATE ?? handlers.handleGuildIntegrationsUpdate, + GUILD_EMOJIS_UPDATE: options.GUILD_EMOJIS_UPDATE ?? + handlers.handleGuildEmojisUpdate, + GUILD_INTEGRATIONS_UPDATE: options.GUILD_INTEGRATIONS_UPDATE ?? + handlers.handleGuildIntegrationsUpdate, GUILD_MEMBER_ADD: options.GUILD_MEMBER_ADD ?? handlers.handleGuildMemberAdd, - GUILD_MEMBER_REMOVE: options.GUILD_MEMBER_REMOVE ?? handlers.handleGuildMemberRemove, - GUILD_MEMBER_UPDATE: options.GUILD_MEMBER_UPDATE ?? handlers.handleGuildMemberUpdate, - GUILD_MEMBERS_CHUNK: options.GUILD_MEMBERS_CHUNK ?? handlers.handleGuildMembersChunk, - GUILD_ROLE_CREATE: options.GUILD_ROLE_CREATE ?? handlers.handleGuildRoleCreate, - GUILD_ROLE_DELETE: options.GUILD_ROLE_DELETE ?? handlers.handleGuildRoleDelete, - GUILD_ROLE_UPDATE: options.GUILD_ROLE_UPDATE ?? handlers.handleGuildRoleUpdate, + GUILD_MEMBER_REMOVE: options.GUILD_MEMBER_REMOVE ?? + handlers.handleGuildMemberRemove, + GUILD_MEMBER_UPDATE: options.GUILD_MEMBER_UPDATE ?? + handlers.handleGuildMemberUpdate, + GUILD_MEMBERS_CHUNK: options.GUILD_MEMBERS_CHUNK ?? + handlers.handleGuildMembersChunk, + GUILD_ROLE_CREATE: options.GUILD_ROLE_CREATE ?? + handlers.handleGuildRoleCreate, + GUILD_ROLE_DELETE: options.GUILD_ROLE_DELETE ?? + handlers.handleGuildRoleDelete, + GUILD_ROLE_UPDATE: options.GUILD_ROLE_UPDATE ?? + handlers.handleGuildRoleUpdate, GUILD_UPDATE: options.GUILD_UPDATE ?? handlers.handleGuildUpdate, // guild events - GUILD_SCHEDULED_EVENT_CREATE: options.GUILD_SCHEDULED_EVENT_CREATE ?? handlers.handleGuildScheduledEventCreate, - GUILD_SCHEDULED_EVENT_DELETE: options.GUILD_SCHEDULED_EVENT_DELETE ?? handlers.handleGuildScheduledEventDelete, - GUILD_SCHEDULED_EVENT_UPDATE: options.GUILD_SCHEDULED_EVENT_UPDATE ?? handlers.handleGuildScheduledEventUpdate, - GUILD_SCHEDULED_EVENT_USER_ADD: options.GUILD_SCHEDULED_EVENT_USER_ADD ?? handlers.handleGuildScheduledEventUserAdd, + GUILD_SCHEDULED_EVENT_CREATE: options.GUILD_SCHEDULED_EVENT_CREATE ?? + handlers.handleGuildScheduledEventCreate, + GUILD_SCHEDULED_EVENT_DELETE: options.GUILD_SCHEDULED_EVENT_DELETE ?? + handlers.handleGuildScheduledEventDelete, + GUILD_SCHEDULED_EVENT_UPDATE: options.GUILD_SCHEDULED_EVENT_UPDATE ?? + handlers.handleGuildScheduledEventUpdate, + GUILD_SCHEDULED_EVENT_USER_ADD: options.GUILD_SCHEDULED_EVENT_USER_ADD ?? + handlers.handleGuildScheduledEventUserAdd, GUILD_SCHEDULED_EVENT_USER_REMOVE: options.GUILD_SCHEDULED_EVENT_USER_REMOVE ?? handlers.handleGuildScheduledEventUserRemove, // interactions - INTERACTION_CREATE: options.INTERACTION_CREATE ?? handlers.handleInteractionCreate, + INTERACTION_CREATE: options.INTERACTION_CREATE ?? + handlers.handleInteractionCreate, // invites INVITE_CREATE: options.INVITE_CREATE ?? handlers.handleInviteCreate, INVITE_DELETE: options.INVITE_DELETE ?? handlers.handleInviteCreate, // messages MESSAGE_CREATE: options.MESSAGE_CREATE ?? handlers.handleMessageCreate, - MESSAGE_DELETE_BULK: options.MESSAGE_DELETE_BULK ?? handlers.handleMessageDeleteBulk, + MESSAGE_DELETE_BULK: options.MESSAGE_DELETE_BULK ?? + handlers.handleMessageDeleteBulk, MESSAGE_DELETE: options.MESSAGE_DELETE ?? handlers.handleMessageDelete, - MESSAGE_REACTION_ADD: options.MESSAGE_REACTION_ADD ?? handlers.handleMessageReactionAdd, - MESSAGE_REACTION_REMOVE_ALL: options.MESSAGE_REACTION_REMOVE_ALL ?? handlers.handleMessageReactionRemoveAll, - MESSAGE_REACTION_REMOVE_EMOJI: options.MESSAGE_REACTION_REMOVE_EMOJI ?? handlers.handleMessageReactionRemoveEmoji, - MESSAGE_REACTION_REMOVE: options.MESSAGE_REACTION_REMOVE ?? handlers.handleMessageReactionRemove, + MESSAGE_REACTION_ADD: options.MESSAGE_REACTION_ADD ?? + handlers.handleMessageReactionAdd, + MESSAGE_REACTION_REMOVE_ALL: options.MESSAGE_REACTION_REMOVE_ALL ?? + handlers.handleMessageReactionRemoveAll, + MESSAGE_REACTION_REMOVE_EMOJI: options.MESSAGE_REACTION_REMOVE_EMOJI ?? + handlers.handleMessageReactionRemoveEmoji, + MESSAGE_REACTION_REMOVE: options.MESSAGE_REACTION_REMOVE ?? + handlers.handleMessageReactionRemove, MESSAGE_UPDATE: options.MESSAGE_UPDATE ?? handlers.handleMessageUpdate, // presence PRESENCE_UPDATE: options.PRESENCE_UPDATE ?? handlers.handlePresenceUpdate, TYPING_START: options.TYPING_START ?? handlers.handleTypingStart, USER_UPDATE: options.USER_UPDATE ?? handlers.handleUserUpdate, // voice - VOICE_SERVER_UPDATE: options.VOICE_SERVER_UPDATE ?? handlers.handleVoiceServerUpdate, - VOICE_STATE_UPDATE: options.VOICE_STATE_UPDATE ?? handlers.handleVoiceStateUpdate, + VOICE_SERVER_UPDATE: options.VOICE_SERVER_UPDATE ?? + handlers.handleVoiceServerUpdate, + VOICE_STATE_UPDATE: options.VOICE_STATE_UPDATE ?? + handlers.handleVoiceStateUpdate, // webhooks WEBHOOKS_UPDATE: options.WEBHOOKS_UPDATE ?? handlers.handleWebhooksUpdate, // integrations - INTEGRATION_CREATE: options.INTEGRATION_CREATE ?? handlers.handleIntegrationCreate, - INTEGRATION_UPDATE: options.INTEGRATION_UPDATE ?? handlers.handleIntegrationUpdate, - INTEGRATION_DELETE: options.INTEGRATION_DELETE ?? handlers.handleIntegrationDelete, + INTEGRATION_CREATE: options.INTEGRATION_CREATE ?? + handlers.handleIntegrationCreate, + INTEGRATION_UPDATE: options.INTEGRATION_UPDATE ?? + handlers.handleIntegrationUpdate, + INTEGRATION_DELETE: options.INTEGRATION_DELETE ?? + handlers.handleIntegrationDelete, }; } @@ -723,5 +813,7 @@ export type RemoveFirstFromTuple = T["length"] extends 0 ? [] : ((...b: T) => void) extends (a: any, ...b: infer I) => void ? I : []; export type FinalHelpers = { - [K in keyof Helpers]: (...args: RemoveFirstFromTuple>) => ReturnType; + [K in keyof Helpers]: ( + ...args: RemoveFirstFromTuple> + ) => ReturnType; }; diff --git a/src/ws/gateway_manager.ts b/src/ws/gateway_manager.ts index 6510bdd6c..7f38056bc 100644 --- a/src/ws/gateway_manager.ts +++ b/src/ws/gateway_manager.ts @@ -6,7 +6,14 @@ import { handleOnMessage } from "./handleOnMessage.ts"; import { heartbeat } from "./heartbeat.ts"; import { identify } from "./identify.ts"; import { processGatewayQueue } from "./processGatewayQueue.ts"; -import { resharder } from "./resharder.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"; @@ -67,7 +74,14 @@ export function createGatewayManager( heartbeat: options.heartbeat ?? heartbeat, tellWorkerToIdentify, debug: options.debug || function () {}, - resharder: options.resharder ?? resharder, + 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, @@ -166,8 +180,21 @@ export interface GatewayManager { tellWorkerToIdentify: typeof tellWorkerToIdentify; /** Handle the different logs. Used for debugging. */ debug: (text: string, ...args: any[]) => unknown; - /** Handles resharding the bot when necessary. */ - resharder: typeof resharder; + /** 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. */ diff --git a/src/ws/resharder.ts b/src/ws/resharder.ts index 63444d4b9..4eab11175 100644 --- a/src/ws/resharder.ts +++ b/src/ws/resharder.ts @@ -1,26 +1,38 @@ import { GetGatewayBot } from "../types/gateway/getGatewayBot.ts"; -import { GatewayManager } from "./gateway_manager.ts"; +import { DiscordReady } from "../types/gateway/ready.ts"; +import { Collection } from "../util/collection.ts"; +import { createGatewayManager, GatewayManager } from "./gateway_manager.ts"; /** The handler to automatically reshard when necessary. */ -export async function resharder(gateway: GatewayManager) { - // 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())) as GetGatewayBot; +export async function resharder( + oldGateway: GatewayManager, + results: GetGatewayBot, +) { + oldGateway.debug("[Resharding] Starting the reshard process."); - 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; - } + const gateway = createGatewayManager({ + ...oldGateway, + // IGNORE EVENTS FOR NOW + 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); + } + }, + }); // 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.maxShards && gateway.useOptimalLargeBotSharding) { + gateway.debug("[Resharding] Using optimal large bot sharding solution."); gateway.maxShards = Math.ceil( gateway.maxShards / (results.sessionStartLimit.maxConcurrency === 1 ? 16 : results.sessionStartLimit.maxConcurrency), @@ -28,4 +40,100 @@ export async function resharder(gateway: GatewayManager) { } 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("[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 overriden 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.maxShards; 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 overriden 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 overriden 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("[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())) 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 overriden 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 SHARDID 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 }