feat: backport handle zombie connection (#7626)

* feat: backport zombie connection fixes

* fix: enums

* fix: prettier

* feat: add zombie connection event to shard events

* Apply suggestions from code review

Co-authored-by: muchnameless <12682826+muchnameless@users.noreply.github.com>

* fix: prettier

* fix: handleZombieConnection

* feat: backport new logic of handling zombie connection

Co-authored-by: muchnameless <12682826+muchnameless@users.noreply.github.com>

packages/discord.js/src/client/websocket/WebSocketManager.js

@@ -220,13 +220,8 @@ class WebSocketManager extends EventEmitter {
 
         this.shardQueue.add(shard);
 
-        if (shard.sessionId) {
-          this.debug(`Session id is present, attempting an immediate reconnect...`, shard);
-          this.reconnect();
-        } else {
-          shard.destroy({ reset: true, emit: false, log: false });
-          this.reconnect();
-        }
+        if (shard.sessionId) this.debug(`Session id is present, attempting an immediate reconnect...`, shard);
+        this.reconnect();
       });
 
       shard.on(ShardEvents.InvalidSession, () => {

packages/discord.js/src/client/websocket/WebSocketShard.js

@@ -84,6 +84,13 @@ class WebSocketShard extends EventEmitter {
      */
     this.lastHeartbeatAcked = true;
 
+    /**
+     * Used to prevent calling {@link WebSocketShard#event:close} twice while closing or terminating the WebSocket.
+     * @type {boolean}
+     * @private
+     */
+    this.closeEmitted = false;
+
     /**
      * Contains the rate limit queue and metadata
      * @name WebSocketShard#ratelimit
@@ -129,6 +136,14 @@ class WebSocketShard extends EventEmitter {
      */
     Object.defineProperty(this, 'helloTimeout', { value: null, writable: true });
 
+    /**
+     * The WebSocket timeout.
+     * @name WebSocketShard#wsCloseTimeout
+     * @type {?NodeJS.Timeout}
+     * @private
+     */
+    Object.defineProperty(this, 'wsCloseTimeout', { value: null, writable: true });
+
     /**
      * If the manager attached its event handlers on the shard
      * @name WebSocketShard#eventsAttached
@@ -256,7 +271,8 @@ class WebSocketShard extends EventEmitter {
 
       this.connectedAt = Date.now();
 
-      const ws = (this.connection = WebSocket.create(gateway, wsQuery));
+      // Adding a handshake timeout to just make sure no zombie connection appears.
+      const ws = (this.connection = WebSocket.create(gateway, wsQuery, { handshakeTimeout: 30_000 }));
       ws.onopen = this.onOpen.bind(this);
       ws.onmessage = this.onMessage.bind(this);
       ws.onerror = this.onError.bind(this);
@@ -343,21 +359,35 @@ class WebSocketShard extends EventEmitter {
    * @private
    */
   onClose(event) {
+    this.closeEmitted = true;
     if (this.sequence !== -1) this.closeSequence = this.sequence;
     this.sequence = -1;
+    this.setHeartbeatTimer(-1);
+    this.setHelloTimeout(-1);
+    // Clearing the WebSocket close timeout as close was emitted.
+    this.setWsCloseTimeout(-1);
+    // If we still have a connection object, clean up its listeners
+    if (this.connection) this._cleanupConnection();
+    this.status = Status.DISCONNECTED;
+    this.emitClose(event);
+  }
 
+  /**
+   * This method is responsible to emit close event for this shard.
+   * This method helps the shard reconnect.
+   * @param {CloseEvent} [event] Close event that was received
+   */
+  emitClose(
+    event = {
+      code: 1011,
+      reason: 'INTERNAL_ERROR',
+      wasClean: false,
+    },
+  ) {
     this.debug(`[CLOSE]
     Event Code: ${event.code}
     Clean     : ${event.wasClean}
     Reason    : ${event.reason ?? 'No reason received'}`);
-
-    this.setHeartbeatTimer(-1);
-    this.setHelloTimeout(-1);
-    // If we still have a connection object, clean up its listeners
-    if (this.connection) this._cleanupConnection();
-
-    this.status = Status.Disconnected;
-
     /**
      * Emitted when a shard's WebSocket closes.
      * @private
@@ -366,7 +396,6 @@ class WebSocketShard extends EventEmitter {
      */
     this.emit(ShardEvents.Close, event);
   }
-
   /**
    * Called whenever a packet is received.
    * @param {Object} packet The received packet
@@ -526,6 +555,47 @@ class WebSocketShard extends EventEmitter {
     }, 20_000).unref();
   }
 
+  /**
+   * Sets the WebSocket Close timeout.
+   * This method is responsible for detecting any zombie connections if the WebSocket fails to close properly.
+   * @param {number} [time] If set to -1, it will clear the timeout
+   * @private
+   */
+  setWsCloseTimeout(time) {
+    if (this.wsCloseTimeout) {
+      this.debug('[WebSocket] Clearing the close timeout.');
+      clearTimeout(this.wsCloseTimeout);
+    }
+    if (time === -1) {
+      this.wsCloseTimeout = null;
+      return;
+    }
+    this.wsCloseTimeout = setTimeout(() => {
+      this.setWsCloseTimeout(-1);
+      this.debug(`[WebSocket] Close Emitted: ${this.closeEmitted}`);
+      // Check if close event was emitted.
+      if (this.closeEmitted) {
+        this.debug(
+          `[WebSocket] was closed. | WS State: ${
+            CONNECTION_STATE[this.connection?.readyState ?? WebSocket.CLOSED]
+          } | Close Emitted: ${this.closeEmitted}`,
+        );
+        // Setting the variable false to check for zombie connections.
+        this.closeEmitted = false;
+        return;
+      }
+
+      this.debug(
+        // eslint-disable-next-line max-len
+        `[WebSocket] did not close properly, assuming a zombie connection.\nEmitting close and reconnecting again.`,
+      );
+
+      this.emitClose();
+      // Setting the variable false to check for zombie connections.
+      this.closeEmitted = false;
+    }, time).unref();
+  }
+
   /**
    * Sets the heartbeat timer for this shard.
    * @param {number} time If -1, clears the interval, any other number sets an interval
@@ -567,7 +637,7 @@ class WebSocketShard extends EventEmitter {
     Connection State: ${this.connection ? CONNECTION_STATE[this.connection.readyState] : 'No Connection??'}`,
       );
 
-      this.destroy({ closeCode: 4009, reset: true });
+      this.destroy({ reset: true, closeCode: 4009 });
       return;
     }
 
@@ -716,11 +786,17 @@ class WebSocketShard extends EventEmitter {
     this.setHeartbeatTimer(-1);
     this.setHelloTimeout(-1);
 
+    this.debug(
+      `[WebSocket] Destroy: Attempting to close the WebSocket. | WS State: ${
+        CONNECTION_STATE[this.connection?.readyState ?? WebSocket.CLOSED]
+      }`,
+    );
     // Step 1: Close the WebSocket connection, if any, otherwise, emit DESTROYED
     if (this.connection) {
       // If the connection is currently opened, we will (hopefully) receive close
       if (this.connection.readyState === WebSocket.OPEN) {
         this.connection.close(closeCode);
+        this.debug(`[WebSocket] Close: Tried closing. | WS State: ${CONNECTION_STATE[this.connection.readyState]}`);
       } else {
         // Connection is not OPEN
         this.debug(`WS State: ${CONNECTION_STATE[this.connection.readyState]}`);
@@ -729,8 +805,13 @@ class WebSocketShard extends EventEmitter {
         // Attempt to close the connection just in case
         try {
           this.connection.close(closeCode);
-        } catch {
-          // No-op
+        } catch (err) {
+          this.debug(
+            `[WebSocket] Close: Something went wrong while closing the WebSocket: ${
+              err.message || err
+            }. Forcefully terminating the connection | WS State: ${CONNECTION_STATE[this.connection.readyState]}`,
+          );
+          this.connection.terminate();
         }
         // Emit the destroyed event if needed
         if (emit) this._emitDestroyed();
@@ -740,11 +821,20 @@ class WebSocketShard extends EventEmitter {
       this._emitDestroyed();
     }
 
+    if (this.connection?.readyState === WebSocket.CLOSING || this.connection?.readyState === WebSocket.CLOSED) {
+      this.closeEmitted = false;
+      this.debug(
+        `[WebSocket] Adding a WebSocket close timeout to ensure a correct WS reconnect.
+        Timeout: ${this.manager.client.options.closeTimeout}ms`,
+      );
+      this.setWsCloseTimeout(this.manager.client.options.closeTimeout);
+    }
+
     // Step 2: Null the connection object
     this.connection = null;
 
-    // Step 3: Set the shard status to Disconnected
-    this.status = Status.Disconnected;
+    // Step 3: Set the shard status to DISCONNECTED
+    this.status = Status.DISCONNECTED;
 
     // Step 4: Cache the old sequence (use to attempt a resume)
     if (this.sequence !== -1) this.closeSequence = this.sequence;

packages/discord.js/src/util/Options.js

@@ -17,6 +17,8 @@ const Transformers = require('./Transformers');
  * @property {number|number[]|string} [shards] The shard's id to run, or an array of shard ids. If not specified,
  * the client will spawn {@link ClientOptions#shardCount} shards. If set to `auto`, it will fetch the
  * recommended amount of shards from Discord and spawn that amount
+ * @property {number} [closeTimeout=1] The amount of time in milliseconds to wait for the close frame to be received
+ * from the WebSocket. Don't have this too high/low. Its best to have it between 2_000-6_000 ms.
  * @property {number} [shardCount=1] The total amount of shards used by all processes of this bot
  * (e.g. recommended shard count, shard count of the ShardingManager)
  * @property {CacheFactory} [makeCache] Function to create a cache.
@@ -72,6 +74,7 @@ class Options extends null {
    */
   static createDefault() {
     return {
+      closeTimeout: 5_000,
       waitGuildTimeout: 15_000,
       shardCount: 1,
       makeCache: this.cacheWithLimits(this.DefaultMakeCacheSettings),

packages/discord.js/typings/index.d.ts

@@ -2819,6 +2819,8 @@ export class WebSocketShard extends EventEmitter {
   private eventsAttached: boolean;
   private expectedGuilds: Set<Snowflake> | null;
   private readyTimeout: NodeJS.Timeout | null;
+  private closeEmitted: boolean;
+  private wsCloseTimeout: NodeJS.Timeout | null;
 
   public manager: WebSocketManager;
   public id: number;
@@ -2834,6 +2836,7 @@ export class WebSocketShard extends EventEmitter {
   private onPacket(packet: unknown): void;
   private checkReady(): void;
   private setHelloTimeout(time?: number): void;
+  private setWsCloseTimeout(time?: number): void;
   private setHeartbeatTimer(time: number): void;
   private sendHeartbeat(): void;
   private ackHeartbeat(): void;
@@ -2843,6 +2846,7 @@ export class WebSocketShard extends EventEmitter {
   private _send(data: unknown): void;
   private processQueue(): void;
   private destroy(destroyOptions?: { closeCode?: number; reset?: boolean; emit?: boolean; log?: boolean }): void;
+  private emitClose(event?: CloseEvent): void;
   private _cleanupConnection(): void;
   private _emitDestroyed(): void;
 
@@ -3920,6 +3924,7 @@ export interface ClientFetchInviteOptions {
 export interface ClientOptions {
   shards?: number | number[] | 'auto';
   shardCount?: number;
+  closeTimeout?: number;
   makeCache?: CacheFactory;
   allowedMentions?: MessageMentionOptions;
   partials?: Partials[];