diff --git a/packages/gateway/src/Shard.ts b/packages/gateway/src/Shard.ts index 607117bb0..b9b69e953 100644 --- a/packages/gateway/src/Shard.ts +++ b/packages/gateway/src/Shard.ts @@ -72,6 +72,8 @@ export class DiscordenoShard { acknowledged: false, interval: 45000, } + + if (options.requestIdentify) this.requestIdentify = options.requestIdentify } /** The gateway configuration which is used to connect to Discord. */ @@ -153,7 +155,7 @@ export class DiscordenoShard { // A new identify has been requested even though there is already a connection open. // Therefore we need to close the old connection and heartbeating before creating a new one. if (this.isOpen()) { - console.log(`CLOSING EXISTING SHARD: #${this.id}`) + logger.debug(`CLOSING EXISTING SHARD: #${this.id}`) this.close(ShardSocketCloseCodes.ReIdentifying, 'Re-identifying closure of old connection.') } @@ -499,10 +501,7 @@ 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 { - // TODO: how to handle this - // return await options.requestIdentify(this.id) - } + async requestIdentify(): Promise {} /** Start sending heartbeat payloads to Discord in the provided interval. */ startHeartbeating(interval: number): void { @@ -748,6 +747,8 @@ export interface ShardCreateOptions { connection: ShardGatewayConfig /** The event handlers for events on the shard. */ events: ShardEvents + /** The handler to request a space to make an identify request. */ + requestIdentify?: () => Promise } export default DiscordenoShard diff --git a/packages/gateway/src/manager.ts b/packages/gateway/src/manager.ts index e7533cb19..c1f5921ec 100644 --- a/packages/gateway/src/manager.ts +++ b/packages/gateway/src/manager.ts @@ -162,11 +162,17 @@ export function createGatewayManager(options: CreateGatewayManagerOptions): Gate version: this.version, }, events: options.events, + requestIdentify: async () => { + await gateway.buckets.get(shardId % gateway.connection.sessionStartLimit.maxConcurrency)!.leak.acquire(1) + }, }) 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() }, async kill(shardId: number) { @@ -180,8 +186,9 @@ export function createGatewayManager(options: CreateGatewayManagerOptions): Gate return await shard.shutdown() }, - async requestIdentify() { + async requestIdentify(shardId: number) { logger.debug(`[Gateway] requesting identify`) + await gateway.buckets.get(shardId % gateway.connection.sessionStartLimit.maxConcurrency)!.leak.acquire(1) }, // Helpers methods below this @@ -372,7 +379,7 @@ export interface GatewayManager extends Required { /** Kill a shard. Close a shards connection to Discord's gateway (if any) and remove it from the manager. */ kill: (shardId: number) => Promise /** This function communicates with the parent manager, in order to know whether this manager is allowed to identify a new shard. */ - requestIdentify: () => Promise + requestIdentify: (shardId: number) => Promise /** Calculates the number of shards based on the guild id and total shards. */ calculateShardId: (guildId: BigString, totalShards?: number) => number /**