@biggy1337/discord-sbjs 1.3.11

package/src/client/websocket/WebSocketManager.js

@@ -0,0 +1,439 @@
+'use strict';
+
+const EventEmitter = require('node:events');
+const { setImmediate } = require('node:timers');
+const { setTimeout: sleep } = require('node:timers/promises');
+const { Collection } = require('@discordjs/collection');
+const DispatchTable = require('./DispatchTable');
+const WebSocketShard = require('./WebSocketShard');
+const { Error } = require('../../errors');
+const { Events, ShardEvents, Status, WSCodes, WSEvents } = require('../../util/Constants');
+const FastQueue = require('../../util/FastQueue');
+const { hasListener } = require('../../util/ListenerUtil');
+
+const BeforeReadyWhitelist = new Set([
+  WSEvents.READY,
+  WSEvents.RESUMED,
+  WSEvents.GUILD_CREATE,
+  WSEvents.GUILD_DELETE,
+  WSEvents.GUILD_MEMBERS_CHUNK,
+  WSEvents.GUILD_MEMBER_ADD,
+  WSEvents.GUILD_MEMBER_REMOVE,
+]);
+
+const UNRECOVERABLE_CLOSE_CODES = new Set([4004, 4010, 4011, 4012, 4013, 4014, 4015, 4016]);
+const UNRESUMABLE_CLOSE_CODES = new Set([4007, 4009]);
+const getCloseCodeLabel = code => WSCodes[code] ?? `UNKNOWN_CLOSE_CODE_${code}`;
+
+/**
+ * The WebSocket manager for this client.
+ * <info>This class forwards raw dispatch events,
+ * read more about it here {@link https://discord.com/developers/docs/topics/gateway}</info>
+ * @extends {EventEmitter}
+ */
+class WebSocketManager extends EventEmitter {
+  constructor(client) {
+    super();
+
+    this._listenerCountsFast = Object.create(null);
+    this.on('newListener', this._onNewListener.bind(this));
+    this.on('removeListener', this._onRemoveListener.bind(this));
+
+    /**
+     * The client that instantiated this WebSocketManager
+     * @type {Client}
+     * @readonly
+     * @name WebSocketManager#client
+     */
+    Object.defineProperty(this, 'client', { value: client });
+
+    /**
+     * The gateway this manager uses
+     * @type {?string}
+     */
+    this.gateway = null;
+
+    /**
+     * The amount of shards this manager handles
+     * @private
+     * @type {number}
+     */
+    this.totalShards = this.client.options.shards.length;
+
+    /**
+     * A collection of all shards this manager handles
+     * @type {Collection<number, WebSocketShard>}
+     */
+    this.shards = new Collection();
+
+    /**
+     * An array of shards to be connected or that need to reconnect
+     * @type {Set<WebSocketShard>}
+     * @private
+     * @name WebSocketManager#shardQueue
+     */
+    Object.defineProperty(this, 'shardQueue', { value: new Set(), writable: true });
+
+    /**
+     * An array of queued events before this WebSocketManager became ready
+     * @type {Object[]}
+     * @private
+     * @name WebSocketManager#packetQueue
+     */
+    Object.defineProperty(this, 'packetQueue', { value: new FastQueue() });
+    this._processingQueue = false;
+
+    /**
+     * The current status of this WebSocketManager
+     * @type {Status}
+     */
+    this.status = Status.IDLE;
+
+    /**
+     * If this manager was destroyed. It will prevent shards from reconnecting
+     * @type {boolean}
+     * @private
+     */
+    this.destroyed = false;
+
+    /**
+     * If this manager is currently reconnecting one or multiple shards
+     * @type {boolean}
+     * @private
+     */
+    this.reconnecting = false;
+  }
+
+  /**
+   * The average ping of all WebSocketShards
+   * @type {number}
+   * @readonly
+   */
+  get ping() {
+    let total = 0;
+    for (const shard of this.shards.values()) {
+      total += shard.ping;
+    }
+    return this.shards.size ? total / this.shards.size : 0;
+  }
+
+  _onNewListener(event) {
+    if (event === 'newListener' || event === 'removeListener') return;
+    this._listenerCountsFast[event] = (this._listenerCountsFast[event] ?? 0) + 1;
+  }
+
+  _onRemoveListener(event) {
+    if (event === 'newListener' || event === 'removeListener') return;
+    const current = this._listenerCountsFast[event] ?? 0;
+    if (current <= 1) {
+      delete this._listenerCountsFast[event];
+    } else {
+      this._listenerCountsFast[event] = current - 1;
+    }
+  }
+
+  hasListenerFast(event) {
+    return (this._listenerCountsFast[event] ?? 0) > 0;
+  }
+
+  /**
+   * Emits a debug message.
+   * @param {string} message The debug message
+   * @param {?WebSocketShard} [shard] The shard that emitted this message, if any
+   * @private
+   */
+  debug(message, shard) {
+    if (!hasListener(this.client, Events.DEBUG)) return;
+    this.client.emit(Events.DEBUG, `[WS => ${shard ? `Shard ${shard.id}` : 'Manager'}] ${message}`);
+  }
+
+  /**
+   * Connects this manager to the gateway.
+   * @private
+   */
+  async connect() {
+    let gatewayURL = 'wss://gateway.discord.gg';
+    await this.client.api.gateway
+      .get({ auth: false })
+      .then(r => (gatewayURL = r.url))
+      .catch(() => {});
+
+    const total = Infinity;
+    const remaining = Infinity;
+    const recommendedShards = 1;
+
+    this.debug(`Fetched Gateway Information
+    URL: ${gatewayURL}
+    Recommended Shards: ${recommendedShards}`);
+
+    this.debug(`Session Limit Information
+    Total: ${total}
+    Remaining: ${remaining}`);
+
+    this.gateway = `${gatewayURL}/`;
+
+    let { shards } = this.client.options;
+
+    if (shards === 'auto') {
+      this.debug(`Using the recommended shard count provided by Discord: ${recommendedShards}`);
+      this.totalShards = this.client.options.shardCount = recommendedShards;
+      shards = this.client.options.shards = Array.from({ length: recommendedShards }, (_, i) => i);
+    }
+
+    this.totalShards = shards.length;
+    this.debug(`Spawning shards: ${shards.join(', ')}`);
+    this.shardQueue = new Set(shards.map(id => new WebSocketShard(this, id)));
+
+    return this.createShards();
+  }
+
+  /**
+   * Handles the creation of a shard.
+   * @returns {Promise<boolean>}
+   * @private
+   */
+  async createShards() {
+    // If we don't have any shards to handle, return
+    if (!this.shardQueue.size) return false;
+
+    const [shard] = this.shardQueue;
+
+    this.shardQueue.delete(shard);
+
+    if (!shard.eventsAttached) {
+      shard.on(ShardEvents.ALL_READY, unavailableGuilds => {
+        /**
+         * Emitted when a shard turns ready.
+         * @event Client#shardReady
+         * @param {number} id The shard id that turned ready
+         * @param {?Set<Snowflake>} unavailableGuilds Set of unavailable guild ids, if any
+         */
+        this.client.emit(Events.SHARD_READY, shard.id, unavailableGuilds);
+
+        if (!this.shardQueue.size) this.reconnecting = false;
+        this.checkShardsReady();
+      });
+
+      shard.on(ShardEvents.CLOSE, event => {
+        if (event.code === 1_000 ? this.destroyed : UNRECOVERABLE_CLOSE_CODES.has(event.code)) {
+          /**
+           * Emitted when a shard's WebSocket disconnects and will no longer reconnect.
+           * @event Client#shardDisconnect
+           * @param {CloseEvent} event The WebSocket close event
+           * @param {number} id The shard id that disconnected
+           */
+          this.client.emit(Events.SHARD_DISCONNECT, event, shard.id);
+          this.debug(getCloseCodeLabel(event.code), shard);
+          return;
+        }
+
+        if (UNRESUMABLE_CLOSE_CODES.has(event.code)) {
+          // These event codes cannot be resumed
+          shard.sessionId = null;
+        }
+
+        /**
+         * Emitted when a shard is attempting to reconnect or re-identify.
+         * @event Client#shardReconnecting
+         * @param {number} id The shard id that is attempting to reconnect
+         */
+        this.client.emit(Events.SHARD_RECONNECTING, shard.id);
+
+        this.shardQueue.add(shard);
+
+        if (shard.sessionId) this.debug(`Session id is present, attempting an immediate reconnect...`, shard);
+        this.reconnect();
+      });
+
+      shard.on(ShardEvents.INVALID_SESSION, () => {
+        this.client.emit(Events.SHARD_RECONNECTING, shard.id);
+      });
+
+      shard.on(ShardEvents.DESTROYED, () => {
+        this.debug('Shard was destroyed but no WebSocket connection was present! Reconnecting...', shard);
+
+        this.client.emit(Events.SHARD_RECONNECTING, shard.id);
+
+        this.shardQueue.add(shard);
+        this.reconnect();
+      });
+
+      shard.eventsAttached = true;
+    }
+
+    this.shards.set(shard.id, shard);
+
+    try {
+      await shard.connect();
+    } catch (error) {
+      if (error?.code && UNRECOVERABLE_CLOSE_CODES.has(error.code)) {
+        throw new Error(getCloseCodeLabel(error.code));
+        // Undefined if session is invalid, error event for regular closes
+      } else if (!error || error.code) {
+        this.debug('Failed to connect to the gateway, requeueing...', shard);
+        this.shardQueue.add(shard);
+      } else {
+        throw error;
+      }
+    }
+    // If we have more shards, add a 5s delay
+    if (this.shardQueue.size) {
+      this.debug(`Shard Queue Size: ${this.shardQueue.size}; continuing in 5 seconds...`);
+      await sleep(5_000);
+      return this.createShards();
+    }
+
+    return true;
+  }
+
+  /**
+   * Handles reconnects for this manager.
+   * @private
+   * @returns {Promise<boolean>}
+   */
+  async reconnect() {
+    if (this.reconnecting || this.status !== Status.READY) return false;
+    this.reconnecting = true;
+    try {
+      await this.createShards();
+    } catch (error) {
+      this.debug(`Couldn't reconnect or fetch information about the gateway. ${error}`);
+      if (error.httpStatus !== 401) {
+        this.debug(`Possible network error occurred. Retrying in 5s...`);
+        await sleep(5_000);
+        this.reconnecting = false;
+        return this.reconnect();
+      }
+      // If we get an error at this point, it means we cannot reconnect anymore
+      if (hasListener(this.client, Events.INVALIDATED)) {
+        /**
+         * Emitted when the client's session becomes invalidated.
+         * You are expected to handle closing the process gracefully and preventing a boot loop
+         * if you are listening to this event.
+         * @event Client#invalidated
+         */
+        this.client.emit(Events.INVALIDATED);
+        // Destroy just the shards. This means you have to handle the cleanup yourself
+        this.destroy();
+      } else {
+        this.client.destroy();
+      }
+    } finally {
+      this.reconnecting = false;
+    }
+    return true;
+  }
+
+  /**
+   * Broadcasts a packet to every shard this manager handles.
+   * @param {Object} packet The packet to send
+   * @private
+   */
+  broadcast(packet) {
+    for (const shard of this.shards.values()) shard.send(packet);
+  }
+
+  /**
+   * Destroys this manager and all its shards.
+   * @private
+   */
+  destroy() {
+    if (this.destroyed) return;
+    this.debug(`Manager was destroyed. Called by:\n${new Error('MANAGER_DESTROYED').stack}`);
+    this.destroyed = true;
+    this.shardQueue.clear();
+    for (const shard of this.shards.values()) shard.destroy({ closeCode: 1_000, reset: true, emit: false, log: false });
+  }
+
+  _dispatchPacket(packet, shard) {
+    if (!packet) return;
+
+    const handler = DispatchTable[packet.t];
+    if (handler) {
+      handler(this.client, packet, shard);
+    } else {
+      /**
+       * Emitted whenever a packet isn't handled.
+       * @event Client#unhandledPacket
+       * @param {Object} packet The packet (t: EVENT_NAME, d: any)
+       * @param {Number} shard The shard that received the packet (Shard 0)
+       */
+      this.client.emit(Events.UNHANDLED_PACKET, packet, shard);
+    }
+  }
+
+  _schedulePacketQueue() {
+    if (this._processingQueue || this.status !== Status.READY || !this.packetQueue.length) return;
+    this._processingQueue = true;
+    setImmediate(() => {
+      while (this.packetQueue.length && this.status === Status.READY) {
+        const queued = this.packetQueue.shift();
+        if (!queued) break;
+        const { packet, shard } = queued;
+        this._dispatchPacket(packet, shard);
+      }
+      this._processingQueue = false;
+      if (this.packetQueue.length && this.status === Status.READY) {
+        this._schedulePacketQueue();
+      }
+    }).unref();
+  }
+
+  /**
+   * Processes a packet and queues it if this WebSocketManager is not ready.
+   * @param {Object} [packet] The packet to be handled
+   * @param {WebSocketShard} [shard] The shard that will handle this packet
+   * @returns {boolean}
+   * @private
+   */
+  handlePacket(packet, shard) {
+    if (packet && this.status !== Status.READY && !BeforeReadyWhitelist.has(packet.t)) {
+      this.packetQueue.push({ packet, shard });
+      return false;
+    }
+
+    if (packet) {
+      this._dispatchPacket(packet, shard);
+    }
+
+    if (this.status === Status.READY && this.packetQueue.length) {
+      this._schedulePacketQueue();
+    }
+
+    return true;
+  }
+
+  /**
+   * Checks whether the client is ready to be marked as ready.
+   * @private
+   */
+  checkShardsReady() {
+    if (this.status === Status.READY) return;
+    if (this.shards.size !== this.totalShards || this.shards.some(s => s.status !== Status.READY)) {
+      return;
+    }
+
+    this.triggerClientReady();
+  }
+
+  /**
+   * Causes the client to be marked as ready and emits the ready event.
+   * @private
+   */
+  triggerClientReady() {
+    this.status = Status.READY;
+
+    this.client.readyAt = new Date();
+
+    /**
+     * Emitted when the client becomes ready to start working.
+     * @event Client#ready
+     * @param {Client} client The client
+     */
+    this.client.emit(Events.CLIENT_READY, this.client);
+
+    this.handlePacket();
+  }
+}
+
+module.exports = WebSocketManager;