From 372dc9988bcc20d7ef30add6dbaf680f7f6b0a73 Mon Sep 17 00:00:00 2001 From: ITOH <72305210+itohatweb@users.noreply.github.com> Date: Thu, 11 Mar 2021 14:09:59 +0000 Subject: [PATCH] fix(controllers/READY): reimplement guild cache mechanism (#647) * add lastShardID * fix ready event controller * forgot to push this file * move ready to its own file * some changes * Update READY.ts * some changes idk if they are good * Update options.ts * Update READY.ts * Update guilds.ts --- src/api/controllers/READY.ts | 104 ++++++++++++++++++++++++++++++++++ src/api/controllers/guilds.ts | 17 +++++- src/api/controllers/misc.ts | 66 +-------------------- src/api/controllers/mod.ts | 2 +- src/bot.ts | 6 +- src/types/options.ts | 4 ++ src/ws/shard.ts | 7 ++- 7 files changed, 134 insertions(+), 72 deletions(-) create mode 100644 src/api/controllers/READY.ts diff --git a/src/api/controllers/READY.ts b/src/api/controllers/READY.ts new file mode 100644 index 000000000..ac3e0ade0 --- /dev/null +++ b/src/api/controllers/READY.ts @@ -0,0 +1,104 @@ +import { + eventHandlers, + lastShardID, + setApplicationID, + setBotID, +} from "../../bot.ts"; +import { DiscordPayload, ReadyPayload } from "../../types/discord.ts"; +import { cache } from "../../util/cache.ts"; +import { delay } from "../../util/utils.ts"; +import { allowNextShard, basicShards } from "../../ws/mod.ts"; +import { initialMemberLoadQueue } from "../structures/guild.ts"; +import { structures } from "../structures/mod.ts"; +import { cacheHandlers } from "./cache.ts"; + +/** This function is the internal handler for the ready event. Users can override this with controllers if desired. */ +export async function handleInternalReady( + data: DiscordPayload, + shardID: number, +) { + // The bot has already started, the last shard is resumed, however. + if (cache.isReady) return; + + const payload = data.d as ReadyPayload; + setBotID(payload.user.id); + setApplicationID(payload.application.id); + + // Triggered on each shard + eventHandlers.shardReady?.(shardID); + // Save when the READY event was received to prevent infinite load loops + const now = Date.now(); + + const shard = basicShards.get(shardID); + if (!shard) return; + + // Set ready to false just to go sure + shard.ready = false; + // All guilds are unavailable at first + shard.unavailableGuildIDs = new Set(payload.guilds.map((g) => g.id)); + + // Start ready check in 2 seconds + setTimeout(() => checkReady(payload, shardID, now), 2000); + + // Wait 5 seconds to spawn next shard + await delay(5000); + allowNextShard(); +} + +// Don't pass the shard itself because unavailableGuilds won't be updated by the GUILD_CREATE event +/** This function checks if the shard is fully loaded */ +function checkReady(payload: ReadyPayload, shardID: number, now: number) { + const shard = basicShards.get(shardID); + if (!shard) return; + + // Check if all guilds were loaded + if (shard.unavailableGuildIDs.size) { + if (Date.now() - now > 10000) { + eventHandlers.shardFailedToLoad?.(shardID, shard.unavailableGuildIDs); + // Force execute the loaded function to prevent infinite loop + loaded(shardID); + } else { + // Not all guilds were loaded but 10 seconds haven't passed so check again + setTimeout(() => checkReady(payload, shardID, now), 2000); + } + } else { + // All guilds were loaded + loaded(shardID); + } +} + +async function loaded(shardID: number) { + const shard = basicShards.get(shardID); + if (!shard) return; + + shard.ready = true; + + // If it is the last shard we can go full ready + if (shardID === lastShardID - 1) { + // Still some shards are loading so wait another 2 seconds for them + if (basicShards.some((shard) => !shard.ready)) { + setTimeout(() => loaded(shardID), 2000); + } else { + cache.isReady = true; + eventHandlers.ready?.(); + + // All the members that came in on guild creates should now be processed 1 by 1 + for (const [guildID, members] of initialMemberLoadQueue.entries()) { + await Promise.allSettled( + members.map(async (member) => { + const memberStruct = await structures.createMemberStruct( + member, + guildID, + ); + + return cacheHandlers.set( + "members", + memberStruct.id, + memberStruct, + ); + }), + ); + } + } + } +} diff --git a/src/api/controllers/guilds.ts b/src/api/controllers/guilds.ts index 478dfb636..e828dc361 100644 --- a/src/api/controllers/guilds.ts +++ b/src/api/controllers/guilds.ts @@ -9,7 +9,8 @@ import { } from "../../types/mod.ts"; import { cache } from "../../util/cache.ts"; import { Collection } from "../../util/collection.ts"; -import { Member, structures } from "../structures/mod.ts"; +import { basicShards } from "../../ws/mod.ts"; +import { structures } from "../structures/mod.ts"; import { cacheHandlers } from "./cache.ts"; export async function handleInternalGuildCreate( @@ -26,15 +27,22 @@ export async function handleInternalGuildCreate( ); await cacheHandlers.set("guilds", guildStruct.id, guildStruct); - if (await cacheHandlers.has("unavailableGuilds", payload.id)) { + const shard = basicShards.get(shardID); + + if (shard?.unavailableGuildIDs.has(payload.id)) { await cacheHandlers.delete("unavailableGuilds", payload.id); + + shard.unavailableGuildIDs.delete(payload.id); } if (!cache.isReady) return eventHandlers.guildLoaded?.(guildStruct); eventHandlers.guildCreate?.(guildStruct); } -export async function handleInternalGuildDelete(data: DiscordPayload) { +export async function handleInternalGuildDelete( + data: DiscordPayload, + shardID: number, +) { const payload = data.d as GuildDeletePayload; cacheHandlers.forEach("messages", (message) => { if (message.guildID === payload.id) { @@ -62,6 +70,9 @@ export async function handleInternalGuildDelete(data: DiscordPayload) { }); if (payload.unavailable) { + const shard = basicShards.get(shardID); + if (shard) shard.unavailableGuildIDs.add(payload.id); + return cacheHandlers.set("unavailableGuilds", payload.id, Date.now()); } diff --git a/src/api/controllers/misc.ts b/src/api/controllers/misc.ts index de00f0f5a..a5c34d293 100644 --- a/src/api/controllers/misc.ts +++ b/src/api/controllers/misc.ts @@ -1,4 +1,4 @@ -import { eventHandlers, setApplicationID, setBotID } from "../../bot.ts"; +import { eventHandlers } from "../../bot.ts"; import { DiscordPayload, IntegrationCreateUpdateEvent, @@ -6,76 +6,14 @@ import { InviteCreateEvent, InviteDeleteEvent, PresenceUpdatePayload, - ReadyPayload, TypingStartPayload, UserPayload, VoiceStateUpdatePayload, WebhookUpdatePayload, } from "../../types/mod.ts"; -import { cache } from "../../util/cache.ts"; -import { delay } from "../../util/utils.ts"; -import { allowNextShard } from "../../ws/shard_manager.ts"; -import { initialMemberLoadQueue } from "../structures/guild.ts"; import { structures } from "../structures/mod.ts"; import { cacheHandlers } from "./cache.ts"; -/** This function is the internal handler for the ready event. Users can override this with controllers if desired. */ -export async function handleInternalReady( - data: DiscordPayload, - shardID: number, -) { - const payload = data.d as ReadyPayload; - setBotID(payload.user.id); - setApplicationID(payload.application.id); - - // Triggered on each shard - eventHandlers.shardReady?.(shardID); - if (payload.shard && shardID === payload.shard[1] - 1) { - const loadedAllGuilds = async () => { - const guildsMissing = async () => { - for (const g of payload.guilds) { - if (!(await cacheHandlers.has("guilds", g.id))) return true; - } - return false; - }; - - if (await guildsMissing()) { - setTimeout(loadedAllGuilds, 2000); - } else { - // The bot has already started, the last shard is resumed, however. - if (cache.isReady) return; - - cache.isReady = true; - eventHandlers.ready?.(); - - // All the members that came in on guild creates should now be processed 1 by 1 - for (const [guildID, members] of initialMemberLoadQueue.entries()) { - await Promise.all( - members.map(async (member) => { - const memberStruct = await structures.createMemberStruct( - member, - guildID, - ); - - return cacheHandlers.set( - "members", - memberStruct.id, - memberStruct, - ); - }), - ); - } - } - }; - - setTimeout(loadedAllGuilds, 2000); - } - - // Wait 5 seconds to spawn next shard - await delay(5000); - allowNextShard(); -} - /** This function is the internal handler for the presence update event. Users can override this with controllers if desired. */ export async function handleInternalPresenceUpdate(data: DiscordPayload) { const payload = data.d as PresenceUpdatePayload; @@ -233,7 +171,7 @@ export function handleInternalIntegrationDelete(data: DiscordPayload) { export function handleInternalInviteCreate(payload: DiscordPayload) { if (payload.t !== "INVITE_CREATE") return; - + //TODO: replace with tocamelcase const { channel_id: channelID, created_at: createdAt, diff --git a/src/api/controllers/mod.ts b/src/api/controllers/mod.ts index e0ad5a0ef..f306356b7 100644 --- a/src/api/controllers/mod.ts +++ b/src/api/controllers/mod.ts @@ -40,7 +40,6 @@ import { handleInternalInviteCreate, handleInternalInviteDelete, handleInternalPresenceUpdate, - handleInternalReady, handleInternalTypingStart, handleInternalUserUpdate, handleInternalVoiceStateUpdate, @@ -52,6 +51,7 @@ import { handleInternalMessageReactionRemoveAll, handleInternalMessageReactionRemoveEmoji, } from "./reactions.ts"; +import { handleInternalReady } from "./READY.ts"; import { handleInternalGuildRoleCreate, handleInternalGuildRoleDelete, diff --git a/src/bot.ts b/src/bot.ts index 111f58426..2dce72716 100644 --- a/src/bot.ts +++ b/src/bot.ts @@ -18,6 +18,7 @@ export let eventHandlers: EventHandlers = {}; export let botGatewayData: DiscordBotGatewayData; export let proxyWSURL = `wss://gateway.discord.gg`; +export let lastShardID = 0; export const identifyPayload: DiscordIdentify = { token: "", @@ -60,9 +61,10 @@ export async function startBot(config: BotConfig) { (bits, next) => (bits |= typeof next === "string" ? Intents[next] : next), 0, ); - identifyPayload.shard = [0, botGatewayData.shards]; + lastShardID = botGatewayData.shards; + identifyPayload.shard = [0, lastShardID]; - await spawnShards(botGatewayData, identifyPayload, 0, botGatewayData.shards); + await spawnShards(botGatewayData, identifyPayload, 0, lastShardID); } /** Allows you to dynamically update the event handlers by passing in new eventHandlers */ diff --git a/src/types/options.ts b/src/types/options.ts index 005365a49..1894fdb9c 100644 --- a/src/types/options.ts +++ b/src/types/options.ts @@ -219,6 +219,10 @@ export interface EventHandlers { roleGained?: (guild: Guild, member: Member, roleID: string) => unknown; roleLost?: (guild: Guild, member: Member, roleID: string) => unknown; shardReady?: (shardID: number) => unknown; + shardFailedToLoad?: ( + shardID: number, + guildIDs: Set, + ) => unknown; /** Sent when a user starts typing in a channel. */ typingStart?: (data: TypingStartPayload) => unknown; voiceChannelJoin?: (member: Member, channelID: string) => unknown; diff --git a/src/ws/shard.ts b/src/ws/shard.ts index f6627d17a..a384d72f5 100644 --- a/src/ws/shard.ts +++ b/src/ws/shard.ts @@ -6,13 +6,12 @@ import { DiscordPayload, FetchMembersOptions, GatewayOpcode, - GatewayStatusUpdatePayload, ReadyPayload, } from "../types/mod.ts"; +import { Collection } from "../util/collection.ts"; import { delay } from "../util/utils.ts"; import { decompressWith } from "./deps.ts"; import { handleDiscordPayload } from "./shard_manager.ts"; -import { Collection } from "../util/collection.ts"; export const basicShards = new Collection(); const heartbeating = new Map(); @@ -27,6 +26,8 @@ export interface BasicShard { sessionID: string; previousSequenceNumber: number | null; needToResume: boolean; + ready: boolean; + unavailableGuildIDs: Set; } interface RequestMemberQueuedRequest { @@ -53,6 +54,8 @@ export function createShard( sessionID: oldShard?.sessionID || "", previousSequenceNumber: oldShard?.previousSequenceNumber || 0, needToResume: false, + ready: false, + unavailableGuildIDs: new Set(), }; basicShards.set(basicShard.id, basicShard);