fix: identify proper rate limits

packages/gateway/src/Shard.ts

@@ -6,10 +6,10 @@ import type {
   DiscordHello,
   DiscordMember,
   DiscordReady,
-  RequestGuildMembers
+  RequestGuildMembers,
 } from '@discordeno/types'
 import { GatewayCloseEventCodes, GatewayIntents, GatewayOpcodes } from '@discordeno/types'
-import { camelize, Collection, delay, LeakyBucket, logger } from '@discordeno/utils'
+import { Collection, LeakyBucket, camelize, delay, logger } from '@discordeno/utils'
 import { inflateSync } from 'node:zlib'
 import WebSocket from 'ws'
 import type { RequestMemberRequest } from './manager.js'
@@ -70,6 +70,7 @@ export class DiscordenoShard {
     }
 
     if (options.requestIdentify) this.requestIdentify = options.requestIdentify
+    if (options.shardIsReady) this.shardIsReady = options.shardIsReady
 
     this.bucket = new LeakyBucket({
       max: this.calculateSafeRequests(),
@@ -192,6 +193,8 @@ export class DiscordenoShard {
     return await new Promise((resolve) => {
       this.resolves.set('READY', () => {
         this.events.identified?.(this)
+        // Tells the manager that this shard is ready
+        this.shardIsReady();
         resolve()
       })
       // When identifying too fast,
@@ -267,7 +270,7 @@ export class DiscordenoShard {
     // Else bucket and token wait time just get wasted.
     await this.checkOffline(highPriority)
 
-    await this.bucket.acquire(highPriority)
+    await this.bucket.acquire(this.id, highPriority)
 
     // It's possible, that the shard went offline after a token has been acquired from the bucket.
     await this.checkOffline(highPriority)
@@ -378,7 +381,7 @@ export class DiscordenoShard {
         this.startHeartbeating(interval)
 
         if (this.state !== ShardState.Resuming) {
-          const currentQueue = [...this.bucket.queue];
+          const currentQueue = [...this.bucket.queue]
           // 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.
@@ -389,7 +392,7 @@ export class DiscordenoShard {
           })
 
           // Queue should not be lost on a re-identify.
-          this.bucket.queue.unshift(...currentQueue);
+          this.bucket.queue.unshift(...currentQueue)
         }
 
         this.events.hello?.(this)
@@ -507,6 +510,9 @@ export class DiscordenoShard {
   /** 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. */
   async requestIdentify(): Promise<void> {}
 
+  /** This function communicates with the management process, in order to tell it can identify the next shard. */
+  async shardIsReady(): Promise<void> {}
+
   /** Start sending heartbeat payloads to Discord in the provided interval. */
   startHeartbeating(interval: number): void {
     //   gateway.debug("GW HEARTBEATING_STARTED", { shardId, interval });
@@ -753,6 +759,8 @@ export interface ShardCreateOptions {
   events: ShardEvents
   /** The handler to request a space to make an identify request. */
   requestIdentify?: () => Promise<void>
+  /** The handler to alert the gateway manager that this shard has received a READY event. */
+  shardIsReady?: () => Promise<void>
 }
 
 export default DiscordenoShard

packages/gateway/src/manager.ts

@@ -1,5 +1,5 @@
 import type { AtLeastOne, BigString, Camelize, DiscordGetGatewayBot, DiscordMember, RequestGuildMembers } from '@discordeno/types'
-import { Collection, delay, LeakyBucket, logger } from '@discordeno/utils'
+import { Collection, delay, logger } from '@discordeno/utils'
 import Shard from './Shard.js'
 import type { ShardEvents, StatusUpdate, UpdateVoiceState } from './types.js'
 
@@ -82,9 +82,7 @@ export function createGatewayManager(options: CreateGatewayManagerOptions): Gate
         logger.debug(`[Gateway] Preparing buckets for concurrency: ${i}`)
         gateway.buckets.set(i, {
           workers: [],
-          leak: new LeakyBucket({
-            refillInterval: gateway.spawnShardDelay,
-          }),
+          identifyRequests: [],
         })
       }
 
@@ -159,17 +157,30 @@ export function createGatewayManager(options: CreateGatewayManagerOptions): Gate
           },
           events: options.events,
           requestIdentify: async () => {
-            await gateway.buckets.get(shardId % gateway.connection.sessionStartLimit.maxConcurrency)!.leak.acquire()
+            // TODO: remove this entire concept, its much easier to do in reality.
+            // await gateway.buckets.get(shardId % gateway.connection.sessionStartLimit.maxConcurrency)!.leak.acquire(shardId)
+          },
+          shardIsReady: async () => {
+            logger.debug(`[Shard] Shard #${shardId} is ready`)
+            await delay(gateway.spawnShardDelay)
+            logger.debug(`[Shard] Resolving shard identify request`)
+            gateway.buckets.get(shardId % gateway.connection.sessionStartLimit.maxConcurrency)!.identifyRequests.shift()?.()
           },
         })
 
         this.shards.set(shardId, shard)
       }
 
-      logger.debug(`[Gateway] requesting to identify shard #(${shardId}) from bucket.`)
-      await gateway.requestIdentify(shard.id)
-      logger.debug(`[Gateway] Identify request successful for shard #(${shardId}).`)
-      return await shard.identify()
+      const bucket = gateway.buckets.get(shardId % gateway.connection.sessionStartLimit.maxConcurrency)
+      if (!bucket) return
+
+      return await new Promise((resolve) => {
+        // Mark that we are making an identify request so another is not made.
+        bucket.identifyRequests.push(resolve)
+        logger.debug(`[Gateway] identifying shard #(${shardId}).`)
+        // This will trigger identify and when READY is received it will resolve the above request.
+        shard?.identify()
+      })
     },
     async kill(shardId: number) {
       const shard = this.shards.get(shardId)
@@ -184,7 +195,12 @@ export function createGatewayManager(options: CreateGatewayManagerOptions): Gate
 
     async requestIdentify(shardId: number) {
       logger.debug(`[Gateway] requesting identify`)
-      await gateway.buckets.get(shardId % gateway.connection.sessionStartLimit.maxConcurrency)!.leak.acquire()
+      const bucket = gateway.buckets.get(shardId % gateway.connection.sessionStartLimit.maxConcurrency)
+      if (!bucket) return
+
+      return await new Promise((resolve) => {
+        bucket.identifyRequests.push(resolve)
+      })
     },
 
     // Helpers methods below this
@@ -353,7 +369,8 @@ export interface GatewayManager extends Required<CreateGatewayManagerOptions> {
     number,
     {
       workers: Array<{ id: number; queue: number[] }>
-      leak: LeakyBucket
+      /** Requests to identify shards are made based on whether it is available to be made. */
+      identifyRequests: Array<(value: void | PromiseLike<void>) => void>
     }
   >
   /** The shards that are created. */
@@ -374,7 +391,7 @@ export interface GatewayManager extends Required<CreateGatewayManagerOptions> {
   identify: (shardId: number) => Promise<void>
   /** Kill a shard. Close a shards connection to Discord's gateway (if any) and remove it from the manager. */
   kill: (shardId: number) => Promise<void>
-  /** This function communicates with the parent manager, in order to know whether this manager is allowed to identify a new shard. */
+  /** This function makes sure that the bucket is allowed to make the next identify request. */
   requestIdentify: (shardId: number) => Promise<void>
   /** Calculates the number of shards based on the guild id and total shards. */
   calculateShardId: (guildId: BigString, totalShards?: number) => number

packages/utils/src/bucket.ts

@@ -9,7 +9,7 @@ export class LeakyBucket implements LeakyBucketOptions {
   /** The amount of requests that have been used up already. */
   used: number = 0
   /** The queue of requests to acquire an available request. Mapped by <shardId, resolve()> */
-  queue: Map<number, (value: void | PromiseLike<void>) => void> = new Map()
+  queue: Array<{ shardId: number; resolve: (value: void | PromiseLike<void>) => void }> = []
   /** Whether or not the queue is already processing. */
   processing: boolean = false
   /** The timeout id for the timer to reduce the used amount by the refill amount.  */
@@ -42,7 +42,7 @@ export class LeakyBucket implements LeakyBucketOptions {
       if (this.remaining) {
         logger.debug(`[LeakyBucket] Processing queue. Remaining: ${this.remaining} Length: ${this.queue.length}`)
         // Resolves the promise allowing the paused execution of this request to resolve and continue.
-        this.queue.shift()?.()
+        this.queue.shift()?.resolve()
         // A request can be made
         this.used++
 
@@ -74,16 +74,16 @@ export class LeakyBucket implements LeakyBucketOptions {
     }
 
     // Loop has ended mark false so it can restart later when needed
-    this.processing = false;
+    this.processing = false
   }
 
   /** Pauses the execution until the request is available to be made. */
-  async acquire(highPriority?: boolean): Promise<void> {
+  async acquire(shardId: number, highPriority?: boolean): Promise<void> {
     return await new Promise((resolve) => {
       // High priority requests get added to the start of the queue
-      if (highPriority) this.queue.unshift(resolve)
+      if (highPriority) this.queue.unshift({ shardId, resolve })
       // All other requests get pushed to the end.
-      else this.queue.push(resolve)
+      else this.queue.push({ shardId, resolve })
 
       // Each request should trigger the queue to be processesd.
       void this.processQueue()