From 196cf7652e1a9c9009578a0c532bf6f0ddd7451e Mon Sep 17 00:00:00 2001
From: Schuyler Cebulskie <Gawdl3y@gawdl3y.com>
Date: Thu, 16 Nov 2017 22:49:38 -0500
Subject: [PATCH] Add Shard#ready property and related events

---
 src/sharding/Shard.js           | 70 ++++++++++++++++++++++++++++-----
 src/sharding/ShardClientUtil.js |  9 +++--
 src/sharding/ShardingManager.js |  4 +-
 3 files changed, 69 insertions(+), 14 deletions(-)

diff --git a/src/sharding/Shard.js b/src/sharding/Shard.js
index e67bece78..76da23527 100644
--- a/src/sharding/Shard.js
+++ b/src/sharding/Shard.js
@@ -8,9 +8,9 @@ const { Error } = require('../errors');
  */
 class Shard {
   /**
-   * @param {ShardingManager} manager The sharding manager
-   * @param {number} id The ID of this shard
-   * @param {Array} [args=[]] Command line arguments to pass to the script
+   * @param {ShardingManager} manager Manager that is spawning this shard
+   * @param {number} id ID of this shard
+   * @param {string[]} [args=[]] Command line arguments to pass to the script
    */
   constructor(manager, id, args = []) {
     /**
@@ -26,7 +26,7 @@ class Shard {
     this.id = id;
 
     /**
-     * The environment variables for the shard
+     * Environment variables for the shard's process
      * @type {Object}
      */
     this.env = Object.assign({}, process.env, {
@@ -41,14 +41,33 @@ class Shard {
      */
     this.process = childProcess.fork(path.resolve(this.manager.file), args, {
       env: this.env,
-    });
-    this.process.on('message', this._handleMessage.bind(this));
-    this.process.once('exit', () => {
-      if (this.manager.respawn) this.manager.createShard(this.id);
-    });
+    }).on('message', this._handleMessage.bind(this));
 
+    /**
+     * Whether the shard's {@link Client} is ready
+     * @type {boolean}
+     */
+    this.ready = false;
+
+    /**
+     * Ongoing promises for calls to {@link Shard#eval}, mapped by the `script` they were called with
+     * @type {Map<string, Promise>}
+     * @private
+     */
     this._evals = new Map();
+
+    /**
+     * Ongoing promises for calls to {@link Shard#fetchClientValue}, mapped by the `prop` they were called with
+     * @type {Map<string, Promise>}
+     * @private
+     */
     this._fetches = new Map();
+
+    // Handle the death of the process
+    this.process.once('exit', () => {
+      this.ready = false;
+      if (this.manager.respawn) this.manager.createShard(this.id).catch(err => { this.manager.emit('error', err); });
+    });
   }
 
   /**
@@ -134,6 +153,39 @@ class Shard {
    */
   _handleMessage(message) {
     if (message) {
+      // Shard is ready
+      if (message._ready) {
+        this.ready = true;
+        /**
+         * Emitted upon the shard's {@link Client#ready} event.
+         * @event Shard#ready
+         */
+        this.emit('ready');
+        return;
+      }
+
+      // Shard has disconnected
+      if (message._disconnect) {
+        this.ready = false;
+        /**
+         * Emitted upon the shard's {@link Client#disconnect} event.
+         * @event Shard#disconnect
+         */
+        this.emit('disconnect');
+        return;
+      }
+
+      // Shard is attempting to reconnect
+      if (message._reconnecting) {
+        this.ready = false;
+        /**
+         * Emitted upon the shard's {@link Client#reconnecting} event.
+         * @event Shard#reconnecting
+         */
+        this.emit('reconnecting');
+        return;
+      }
+
       // Shard is requesting a property fetch
       if (message._sFetchProp) {
         this.manager.fetchClientValues(message._sFetchProp).then(
diff --git a/src/sharding/ShardClientUtil.js b/src/sharding/ShardClientUtil.js
index aceb85f2d..70e625824 100644
--- a/src/sharding/ShardClientUtil.js
+++ b/src/sharding/ShardClientUtil.js
@@ -7,11 +7,14 @@ const { Error } = require('../errors');
  */
 class ShardClientUtil {
   /**
-   * @param {Client} client The client of the current shard
+   * @param {Client} client Client of the current shard
    */
   constructor(client) {
     this.client = client;
     process.on('message', this._handleMessage.bind(this));
+    client.on('ready', () => { process.send({ _ready: true }); });
+    client.on('disconnect', () => { process.send({ _disconnect: true }); });
+    client.on('reconnecting', () => { process.send({ _reconnecting: true }); });
   }
 
   /**
@@ -49,7 +52,7 @@ class ShardClientUtil {
   /**
    * Fetches a client property value of each shard.
    * @param {string} prop Name of the client property to get, using periods for nesting
-   * @returns {Promise<Array>}
+   * @returns {Promise<Array<*>>}
    * @example
    * client.shard.fetchClientValues('guilds.size')
    *   .then(results => {
@@ -76,7 +79,7 @@ class ShardClientUtil {
   /**
    * Evaluates a script on all shards, in the context of the Clients.
    * @param {string} script JavaScript to run on each shard
-   * @returns {Promise<Array>} Results of the script execution
+   * @returns {Promise<Array<*>>} Results of the script execution
    */
   broadcastEval(script) {
     return new Promise((resolve, reject) => {
diff --git a/src/sharding/ShardingManager.js b/src/sharding/ShardingManager.js
index 7a18898be..117fca4f9 100644
--- a/src/sharding/ShardingManager.js
+++ b/src/sharding/ShardingManager.js
@@ -168,7 +168,7 @@ class ShardingManager extends EventEmitter {
   /**
    * Evaluates a script on all shards, in the context of the Clients.
    * @param {string} script JavaScript to run on each shard
-   * @returns {Promise<Array>} Results of the script execution
+   * @returns {Promise<Array<*>>} Results of the script execution
    */
   broadcastEval(script) {
     const promises = [];
@@ -179,7 +179,7 @@ class ShardingManager extends EventEmitter {
   /**
    * Fetches a client property value of each shard.
    * @param {string} prop Name of the client property to get, using periods for nesting
-   * @returns {Promise<Array>}
+   * @returns {Promise<Array<*>>}
    * @example
    * manager.fetchClientValues('guilds.size')
    *   .then(results => {