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

@@ -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,
+            );
+          }),
+        );
+      }
+    }
+  }
+}

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());
   }
 

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,

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,

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 */

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<string>,
+  ) => unknown;
   /** Sent when a user starts typing in a channel. */
   typingStart?: (data: TypingStartPayload) => unknown;
   voiceChannelJoin?: (member: Member, channelID: string) => unknown;

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<number, BasicShard>();
 const heartbeating = new Map<number, boolean>();
@@ -27,6 +26,8 @@ export interface BasicShard {
   sessionID: string;
   previousSequenceNumber: number | null;
   needToResume: boolean;
+  ready: boolean;
+  unavailableGuildIDs: Set<string>;
 }
 
 interface RequestMemberQueuedRequest {
@@ -53,6 +54,8 @@ export function createShard(
     sessionID: oldShard?.sessionID || "",
     previousSequenceNumber: oldShard?.previousSequenceNumber || 0,
     needToResume: false,
+    ready: false,
+    unavailableGuildIDs: new Set<string>(),
   };
 
   basicShards.set(basicShard.id, basicShard);