From 8d7a1ad8b05b621192e145f3a468635cf5141394 Mon Sep 17 00:00:00 2001
From: Saya <davexdme@gmail.com>
Date: Sun, 22 Dec 2024 18:38:07 +0800
Subject: [PATCH] feat: test identify server before starting the client

---
 src/client/ShardClient.ts | 48 +++++++++++++++++++++++++++++++++++----
 1 file changed, 43 insertions(+), 5 deletions(-)

diff --git a/src/client/ShardClient.ts b/src/client/ShardClient.ts
index 4cffa6d..c52389d 100644
--- a/src/client/ShardClient.ts
+++ b/src/client/ShardClient.ts
@@ -1,6 +1,6 @@
 import type { Client, ClientOptions as DiscordJsClientOptions } from 'discord.js';
 import { Indomitable } from '../Indomitable';
-import { EnvProcessData, ClientEvents, ClientEventData } from '../Util';
+import {EnvProcessData, ClientEvents, ClientEventData, Delay} from '../Util';
 import { ShardClientUtil } from './ShardClientUtil';
 import { ConcurrencyClient } from '../concurrency/ConcurrencyClient';
 
@@ -9,18 +9,40 @@ export interface PartialInternalEvents {
     data: { clusterId: number, shardId?: number, replayed?: number, event?: CloseEvent, ipcId?: string }
 }
 
+/**
+ * A little helper on binding Indomitable to Discord.JS client
+ */
 export class ShardClient {
-    public readonly manager: Indomitable;
-    public readonly client: Client;
-    public readonly clusterId: number;
+    /**
+     * Indomitable instance (non master client)
+     * @private
+     */
+    private readonly manager: Indomitable;
+    /**
+     * Discord.JS client
+     * @private
+     */
+    private readonly client: Client;
+    /**
+     * Cluster Id of this helper
+     * @private
+     */
+    private readonly clusterId: number;
+    /**
+     * Concurrency client for this instance, if available
+     * @private
+     */
+    private readonly concurrency?: ConcurrencyClient;
+
     public constructor(manager: Indomitable) {
         this.manager = manager;
         const clientOptions = manager.clientOptions as DiscordJsClientOptions || {};
         clientOptions.shards = EnvProcessData.shardIds;
         clientOptions.shardCount = EnvProcessData.shardCount;
         if (manager.handleConcurrency) {
+            this.concurrency = new ConcurrencyClient();
             if (!clientOptions.ws) clientOptions.ws = {};
-            clientOptions.ws.buildIdentifyThrottler = () => Promise.resolve(new ConcurrencyClient());
+            clientOptions.ws.buildIdentifyThrottler = () => Promise.resolve(this.concurrency!);
         }   
         this.client = new manager.client(clientOptions);
         // @ts-expect-error: Override shard client util with indomitable shard client util
@@ -28,7 +50,18 @@ export class ShardClient {
         this.clusterId = Number(EnvProcessData.clusterId);
     }
 
+    /**
+     * Starts this client
+     * @param token
+     */
     public async start(token: string): Promise<void> {
+        if (this.concurrency) {
+            // tests the server if it's accessible first before starting the client
+            this.client.emit('debug', '[Indomitable]: Handle concurrency enabled! Testing the identify server before logging in...');
+            await this.concurrency.waitForIdentify(0, new AbortSignal())
+            this.client.emit('debug', '[Indomitable]: Identify server responded and is working, waiting 5s before starting...');
+            await Delay(5000);
+        }
         // attach listeners 
         this.client.once('ready', () => this.send({ op: ClientEvents.READY, data: { clusterId: this.clusterId }}));
         this.client.on('shardReady', (shardId: number) => this.send({ op: ClientEvents.SHARD_READY, data: { clusterId: this.clusterId, shardId }}));
@@ -39,6 +72,11 @@ export class ShardClient {
         await this.client.login(token);
     }
 
+    /**
+     * A helper send code to make myself less miserable
+     * @param partial
+     * @private
+     */
     private send(partial: PartialInternalEvents): void {
         // @ts-ignore -- our own class
         const shardClientUtil = this.client.shard as ShardClientUtil;