From 272ddac08b99a0bd19b8db1ec771aed6a28fe457 Mon Sep 17 00:00:00 2001
From: MARCROCK22 <marcos22dev@gmail.com>
Date: Sat, 22 Feb 2025 20:47:25 -0400
Subject: [PATCH] fix: retry connection when failed

---
 src/websocket/discord/shard.ts   | 29 +++++++++++++++++++++++++----
 src/websocket/discord/sharder.ts |  2 ++
 src/websocket/discord/shared.ts  |  5 +++++
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/src/websocket/discord/shard.ts b/src/websocket/discord/shard.ts
index d125af2da..c3607913f 100644
--- a/src/websocket/discord/shard.ts
+++ b/src/websocket/discord/shard.ts
@@ -1,5 +1,5 @@
 import { inflateSync } from 'node:zlib';
-import { LogLevels, Logger, type MakeRequired, MergeOptions, hasIntent } from '../../common';
+import { LogLevels, Logger, type MakeRequired, MergeOptions, delay, hasIntent } from '../../common';
 import {
 	type APIGuildMember,
 	GatewayCloseCodes,
@@ -46,8 +46,11 @@ export class Shard {
 	bucket: DynamicBucket;
 	offlineSendQueue: ((_?: unknown) => void)[] = [];
 	pendingGuilds = new Set<string>();
-	options: MakeRequired<ShardOptions, 'properties' | 'ratelimitOptions'>;
+	options: MakeRequired<ShardOptions, 'properties' | 'ratelimitOptions' | 'reconnectTimeout' | 'connectionTimeout'>;
 	isReady = false;
+
+	connectionTimeout?: NodeJS.Timeout;
+
 	private requestGuildMembersChunk = new Map<
 		string,
 		{
@@ -72,6 +75,8 @@ export class Shard {
 					rateLimitResetInterval: 60_000,
 					maxRequestsPerRateLimitTick: 120,
 				},
+				reconnectTimeout: 10e3,
+				connectionTimeout: 30e3,
 			} as ShardOptions,
 			options,
 		);
@@ -127,6 +132,11 @@ export class Shard {
 
 		this.debugger?.debug(`[Shard #${this.id}] Connecting to ${this.currentGatewayURL}`);
 
+		this.connectionTimeout = setTimeout(
+			() => this.reconnect(ShardSocketCloseCodes.Timeout),
+			this.options.connectionTimeout,
+		);
+
 		// @ts-expect-error Use native websocket when using Bun
 		// biome-ignore lint/correctness/noUndeclaredVariables: /\
 		this.websocket = new BaseSocket(typeof Bun === 'undefined' ? 'ws' : 'bun', this.currentGatewayURL);
@@ -217,13 +227,16 @@ export class Shard {
 	}
 
 	disconnect(code = ShardSocketCloseCodes.Shutdown) {
+		clearTimeout(this.connectionTimeout);
+		this.connectionTimeout = undefined;
 		this.debugger?.info(`[Shard #${this.id}] Disconnecting`);
 		this.close(code, 'Shard down request');
 	}
 
-	async reconnect() {
+	async reconnect(code = ShardSocketCloseCodes.Reconnect) {
 		this.debugger?.info(`[Shard #${this.id}] Reconnecting`);
-		this.disconnect(ShardSocketCloseCodes.Reconnect);
+		this.disconnect(code);
+		await delay(this.options.reconnectTimeout);
 		await this.connect();
 	}
 
@@ -236,6 +249,8 @@ export class Shard {
 
 		switch (packet.op) {
 			case GatewayOpcodes.Hello: {
+				clearTimeout(this.connectionTimeout);
+				this.connectionTimeout = undefined;
 				clearInterval(this.heart.nodeInterval);
 
 				this.heart.interval = packet.d.heartbeat_interval;
@@ -275,12 +290,16 @@ export class Shard {
 					switch (packet.t) {
 						case GatewayDispatchEvents.Resumed:
 							{
+								clearTimeout(this.connectionTimeout);
+								this.connectionTimeout = undefined;
 								this.isReady = true;
 								this.offlineSendQueue.map(resolve => resolve());
 								this.options.handlePayload(this.id, packet);
 							}
 							break;
 						case GatewayDispatchEvents.Ready: {
+							clearTimeout(this.connectionTimeout);
+							this.connectionTimeout = undefined;
 							if (hasIntent(this.options.intents, 'Guilds')) {
 								for (let i = 0; i < packet.d.guilds.length; i++) {
 									this.pendingGuilds.add(packet.d.guilds.at(i)!.id);
@@ -411,6 +430,8 @@ export class Shard {
 			case GatewayCloseCodes.UnknownOpcode:
 			case GatewayCloseCodes.InvalidSeq:
 			case GatewayCloseCodes.SessionTimedOut:
+			// shard failed to connect, try connecting from scratch
+			case ShardSocketCloseCodes.Timeout:
 				{
 					this.data.resume_seq = 0;
 					this.data.session_id = undefined;
diff --git a/src/websocket/discord/sharder.ts b/src/websocket/discord/sharder.ts
index eb1f7b97a..3637dd150 100644
--- a/src/websocket/discord/sharder.ts
+++ b/src/websocket/discord/sharder.ts
@@ -100,6 +100,8 @@ export class ShardManager extends Map<number, Shard> {
 			debugger: this.debugger,
 			compress: this.options.compress ?? false,
 			presence: this.options.presence?.(shardId, -1),
+			connectionTimeout: this.options.connectionTimeout,
+			reconnectTimeout: this.options.reconnectTimeout,
 		});
 
 		this.set(shardId, shard);
diff --git a/src/websocket/discord/shared.ts b/src/websocket/discord/shared.ts
index ac3850442..223e4c7b6 100644
--- a/src/websocket/discord/shared.ts
+++ b/src/websocket/discord/shared.ts
@@ -46,6 +46,8 @@ export interface ShardManagerOptions extends ShardDetails {
 		interval: number;
 		percentage: number;
 	};
+	reconnectTimeout?: number;
+	connectionTimeout?: number;
 }
 
 export interface CustomManagerAdapter {
@@ -128,6 +130,8 @@ export interface ShardOptions extends ShardDetails {
 	debugger?: Logger;
 	compress: boolean;
 	presence?: GatewayPresenceUpdateData;
+	reconnectTimeout?: number;
+	connectionTimeout?: number;
 }
 
 export enum ShardSocketCloseCodes {
@@ -136,6 +140,7 @@ export enum ShardSocketCloseCodes {
 	Reconnect = 3020,
 	Resharding = 3030,
 	ShutdownAll = 3040,
+	Timeout = 3050,
 }
 
 export interface WorkerData {