Add getRecommendedShards and automatic shard count in ShardingManager (#796)

* draft stuff

fix docstring for Client#token

Reorganise resolver

make env better for shards, clean up docs

Fix Gus' log messages

7

meh just gateway/bot not v7 :(

final changes, ready for mergin!

build docs

make default totalShards 'auto', fix docs for totalShards type

clean up docs more

run docs

* make consistancy real

* Update and rename getRecommendedShards.js to GetRecommendedShards.js

* Update GetRecommendedShards.js

* Update index.js

* Update RESTMethods.js

* Update Shard.js

* Update GetRecommendedShards.js

* Update ShardingManager.js

* run docs

src/client/Client.js

@@ -114,11 +114,15 @@ class Client extends EventEmitter {
      */
     this.presences = new Collection();
 
-    /**
-     * The authorization token for the logged in user/bot.
-     * @type {?string}
-     */
-    this.token = null;
+    if (!this.token && 'CLIENT_TOKEN' in process.env) {
+      /**
+       * The authorization token for the logged in user/bot.
+       * @type {?string}
+       */
+      this.token = process.env.CLIENT_TOKEN;
+    } else {
+      this.token = null;
+    }
 
     /**
      * The email, if there is one, for the logged in Client

src/client/rest/RESTMethods.js

@@ -14,7 +14,7 @@ class RESTMethods {
     this.rest = restManager;
   }
 
-  loginToken(token) {
+  loginToken(token = this.rest.client.token) {
     token = token.replace('Bot ', '');
     return new Promise((resolve, reject) => {
       this.rest.client.manager.connectToWebSocket(token, resolve, reject);
@@ -49,6 +49,10 @@ class RESTMethods {
     });
   }
 
+  getBotGateway() {
+    return this.rest.makeRequest('get', Constants.Endpoints.botGateway, true);
+  }
+
   sendMessage(channel, content, { tts, nonce, disableEveryone, split } = {}, file = null) {
     return new Promise((resolve, reject) => {
       if (typeof content !== 'undefined') content = this.rest.client.resolver.resolveString(content);

src/index.js

@@ -8,6 +8,7 @@ module.exports = {
   Collection: require('./util/Collection'),
   splitMessage: require('./util/SplitMessage'),
   escapeMarkdown: require('./util/EscapeMarkdown'),
+  getRecommendedShards: require('./util/GetRecommendedShards'),
 
   Channel: require('./structures/Channel'),
   ClientUser: require('./structures/ClientUser'),

src/sharding/Shard.js

@@ -10,7 +10,7 @@ class Shard {
   /**
    * @param {ShardingManager} manager The sharding manager
    * @param {number} id The ID of this shard
-   * @param {array} [spawnArgs=] Optional command line arguments to pass to the shard
+   * @param {array} [spawnArgs=[]] Command line arguments to pass to the shard
    */
   constructor(manager, id, spawnArgs = []) {
     /**
@@ -25,15 +25,22 @@ class Shard {
      */
     this.id = id;
 
+    /**
+     * The environment variables for the shard
+     * @type {Object}
+     */
+    this.env = {
+      SHARD_ID: this.id,
+      SHARD_COUNT: this.manager.totalShards,
+    };
+    if (this.manager.token) this.env.CLIENT_TOKEN = this.manager.token;
+
     /**
      * Process of the shard
      * @type {ChildProcess}
      */
     this.process = childProcess.fork(path.resolve(this.manager.file), spawnArgs, {
-      env: {
-        SHARD_ID: this.id,
-        SHARD_COUNT: this.manager.totalShards,
-      },
+      env: this.env,
     });
     this.process.on('message', this._handleMessage.bind(this));
     this.process.once('exit', () => {

src/sharding/ShardingManager.js

@@ -4,10 +4,12 @@ const EventEmitter = require('events').EventEmitter;
 const mergeDefault = require('../util/MergeDefault');
 const Shard = require('./Shard');
 const Collection = require('../util/Collection');
+const getRecommendedShards = require('../util/GetRecommendedShards');
 
 /**
  * This is a utility class that can be used to help you spawn shards of your Client. Each shard is completely separate
  * from the other. The Shard Manager takes a path to a file and spawns it under the specified amount of shards safely.
+ * If you do not select an amount of shards, the manager will automatically decide the best amount.
  * <warn>The Sharding Manager is still experimental</warn>
  * @extends {EventEmitter}
  */
@@ -18,8 +20,7 @@ class ShardingManager extends EventEmitter {
    */
   constructor(file, options = {}) {
     super();
-
-    options = mergeDefault({ totalShards: 1, respawn: true, spawnArgs: [] }, options);
+    options = mergeDefault({ totalShards: 'auto', respawn: true, spawnArgs: [], token: null }, options);
 
     /**
      * Path to the shard script file
@@ -31,17 +32,20 @@ class ShardingManager extends EventEmitter {
     const stats = fs.statSync(this.file);
     if (!stats.isFile()) throw new Error('File path does not point to a file.');
 
+    if (options.totalShards !== 'auto') {
+      if (typeof options.totalShards !== 'number' || isNaN(options.totalShards)) {
+        throw new TypeError('Amount of shards must be a number.');
+      }
+      if (options.totalShards < 1) throw new RangeError('Amount of shards must be at least 1.');
+      if (options.totalShards !== Math.floor(options.totalShards)) {
+        throw new RangeError('Amount of shards must be an integer.');
+      }
+    }
+
     /**
      * Amount of shards that this manager is going to spawn
-     * @type {number}
+     * @type {number|string}
      */
-    if (typeof options.totalShards !== 'number' || isNaN(options.totalShards)) {
-      throw new TypeError('Amount of shards must be a number.');
-    }
-    if (options.totalShards < 1) throw new RangeError('Amount of shards must be at least 1.');
-    if (options.totalShards !== Math.floor(options.totalShards)) {
-      throw new RangeError('Amount of shards must be an integer.');
-    }
     this.totalShards = options.totalShards;
 
     /**
@@ -52,10 +56,16 @@ class ShardingManager extends EventEmitter {
 
     /**
      * An array of arguments to pass to shards.
-     * @type {array}
+     * @type {string[]}
      */
     this.spawnArgs = options.spawnArgs;
 
+    /**
+     * Token that will be passed to the shards. Also used for auto shard count.
+     * @type {string}
+     */
+    this.token = options.token || null;
+
     /**
      * A collection of shards that this manager has spawned
      * @type {Collection<number, Shard>}
@@ -87,10 +97,28 @@ class ShardingManager extends EventEmitter {
    * @returns {Promise<Collection<number, Shard>>}
    */
   spawn(amount = this.totalShards, delay = 5500) {
-    if (typeof amount !== 'number' || isNaN(amount)) throw new TypeError('Amount of shards must be a number.');
-    if (amount < 1) throw new RangeError('Amount of shards must be at least 1.');
-    if (amount !== Math.floor(amount)) throw new RangeError('Amount of shards must be an integer.');
+    return new Promise((resolve, reject) => {
+      if (amount === 'auto') {
+        getRecommendedShards(this.token).then(count => {
+          resolve(this._spawn(count, delay));
+        }).catch(reject);
+      } else {
+        if (typeof amount !== 'number' || isNaN(amount)) throw new TypeError('Amount of shards must be a number.');
+        if (amount < 1) throw new RangeError('Amount of shards must be at least 1.');
+        if (amount !== Math.floor(amount)) throw new TypeError('Amount of shards must be an integer.');
+        resolve(this._spawn(amount, delay));
+      }
+    });
+  }
 
+  /**
+   * Actually spawns shards, unlike that poser above >:(
+   * @param {number} amount Number of shards to spawn
+   * @param {number} delay How long to wait in between spawning each shard (in milliseconds)
+   * @returns {Promise<Collection<number, Shard>>}
+   * @private
+   */
+  _spawn(amount, delay) {
     return new Promise(resolve => {
       if (this.shards.size >= amount) throw new Error(`Already spawned ${this.shards.size} shards.`);
       this.totalShards = amount;

src/util/Constants.js

@@ -72,6 +72,7 @@ const Endpoints = exports.Endpoints = {
   login: `${API}/auth/login`,
   logout: `${API}/auth/logout`,
   gateway: `${API}/gateway`,
+  botGateway: `${API}/gateway/bot`,
   invite: (id) => `${API}/invite/${id}`,
   inviteLink: (id) => `https://discord.gg/${id}`,
   CDN: 'https://cdn.discordapp.com',

src/util/GetRecommendedShards.js

@@ -0,0 +1,19 @@
+const superagent = require('superagent');
+const botGateway = require('./Constants').Endpoints.botGateway;
+
+/**
+ * Gets the recommended shard count from Discord
+ * @param {number} token Discord auth token
+ * @returns {Promise<number>} the recommended number of shards
+ */
+module.exports = function getRecommendedShards(token) {
+  return new Promise((resolve, reject) => {
+    if (!token) reject(new Error('A token must be provided.'));
+    superagent.get(botGateway)
+      .set('Authorization', `Bot ${token.replace('Bot ', '')}`)
+      .end((err, res) => {
+        if (err) reject(err);
+        resolve(res.body.shards);
+      });
+  });
+};