fix(gateway): rate-limit handling (#2079)

* fix: Closes Gateway Ratelimiting  #1886

* fix: fmt

* fix: invalid import

gateway/gateway_manager.ts

@@ -1,5 +1,6 @@
 import { GatewayIntents, GatewayPayload, StatusUpdate } from "../types/mod.ts";
 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";
@@ -18,7 +19,7 @@ import { resume } from "./resume.ts";
 import { sendShardMessage } from "./sendShardMessage.ts";
 import { prepareBuckets, spawnShards } from "./spawnShards.ts";
 import { tellWorkerToIdentify } from "./tellWorkerToIdentify.ts";
-import { DiscordenoShard } from "./ws.ts";
+import { DiscordenoShard } from "./shard.ts";
 
 /** Create a new Gateway Manager.
  *
@@ -30,6 +31,8 @@ export function createGatewayManager(
   options: Partial<GatewayManager> & Pick<GatewayManager, "handleDiscordPayload">,
 ): GatewayManager {
   return {
+    queueResetInterval: 60000,
+    maxRequestsPerInterval: 120,
     cache: {
       guildIds: new Set(),
       loadingGuildIds: new Set(),
@@ -87,6 +90,7 @@ export function createGatewayManager(
     closeWS: options.closeWS ?? closeWS,
     sendShardMessage: options.sendShardMessage ?? sendShardMessage,
     resume: options.resume ?? resume,
+    safeRequestsPerShard: options.safeRequestsPerShard ?? safeRequestsPerShard,
     handleDiscordPayload: options.handleDiscordPayload,
   };
 }
@@ -155,6 +159,10 @@ export interface GatewayManager {
     }
   >;
   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<bigint>;
@@ -205,4 +213,6 @@ export interface GatewayManager {
   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;
 }

gateway/handleOnMessage.ts

@@ -47,6 +47,8 @@ export async function handleOnMessage(gateway: GatewayManager, message: any, sha
       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 (gateway.shards.has(shardId)) {

gateway/identify.ts

@@ -36,6 +36,8 @@ export function identify(gateway: GatewayManager, shardId: number, maxShards: nu
     processingQueue: false,
     queueStartedAt: Date.now(),
     queueCounter: 0,
+    // BY DEFAULT SET TO 120. EDIT IN HELLO
+    safeRequestsPerShard: 120,
   });
 
   socket.onopen = () => {

gateway/mod.ts

@@ -10,5 +10,6 @@ export * from "./spawnShards.ts";
 export * from "./sendShardMessage.ts";
 export * from "./startGatewayOptions.ts";
 export * from "./tellWorkerToIdentify.ts";
-export * from "./ws.ts";
+export * from "./shard.ts";
 export * from "./gateway_manager.ts";
+export * from "./safeRequestsPerShard.ts";

gateway/processGatewayQueue.ts

@@ -15,7 +15,7 @@ export async function processGatewayQueue(gateway: GatewayManager, id: number) {
     }
 
     const now = Date.now();
-    if (now - shard.queueStartedAt >= 60000) {
+    if (now - shard.queueStartedAt >= gateway.queueResetInterval) {
       shard.queueStartedAt = now;
       shard.queueCounter = 0;
     }
@@ -28,16 +28,20 @@ export async function processGatewayQueue(gateway: GatewayManager, id: number) {
 
     shard.ws.send(JSON.stringify(request));
 
-    // Counter is useful for preventing 120/m requests.
+    // Counter is useful for preventing max requests.
     shard.queueCounter++;
 
     // Handle if the requests have been maxed
-    if (shard.queueCounter >= 116) {
-      gateway.debug("GW MAX_REQUESTS", {
-        message: "Max gateway requests per minute reached setting timeout for one minute",
-        shardId: shard.id,
-      });
-      await delay(60000);
+    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;
     }

gateway/resume.ts

@@ -43,6 +43,7 @@ export function resume(gateway: GatewayManager, shardId: number) {
     processingQueue: false,
     queueStartedAt: Date.now(),
     queueCounter: 0,
+    safeRequestsPerShard: oldShard.safeRequestsPerShard || 120,
   });
 
   // Resume on open

gateway/safeRequestsPerShard.ts

@@ -0,0 +1,9 @@
+import { GatewayManager } from "./gateway_manager.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;
+}

gateway/sendShardMessage.ts

@@ -1,5 +1,5 @@
 import { GatewayManager } from "./gateway_manager.ts";
-import { DiscordenoShard, WebSocketRequest } from "./ws.ts";
+import { DiscordenoShard, WebSocketRequest } from "./shard.ts";
 
 export function sendShardMessage(
   gateway: GatewayManager,

gateway/shard.ts

@@ -40,13 +40,11 @@ export interface DiscordenoShard {
   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;
-  // guildId: bigint;
-  // shardId: number;
-  // nonce?: bigint;
-  // options?: Record<string, unknown>;
 }