'use strict'; const EventEmitter = require('node:events'); const fs = require('node:fs'); const path = require('node:path'); const process = require('node:process'); const { setTimeout: sleep } = require('node:timers/promises'); const { Collection } = require('@discordjs/collection'); const Shard = require('./Shard'); const { Error, TypeError, RangeError } = require('../errors'); const Util = require('../util/Util'); /** * This is a utility class that makes multi-process sharding of a bot an easy and painless experience. * It works by spawning a self-contained {@link ChildProcess} or {@link Worker} for each individual shard, each * containing its own instance of your bot's {@link Client}. They all have a line of communication with the master * process, and there are several useful methods that utilise it in order to simplify tasks that are normally difficult * with sharding. It can spawn a specific number of shards or the amount that Discord suggests for the bot, and takes a * path to your main bot script to launch for each one. * @extends {EventEmitter} */ class ShardingManager extends EventEmitter { /** * The mode to spawn shards with for a {@link ShardingManager}. Can be either one of: * * 'process' to use child processes * * 'worker' to use [Worker threads](https://nodejs.org/api/worker_threads.html) * @typedef {string} ShardingManagerMode */ /** * The options to spawn shards with for a {@link ShardingManager}. * @typedef {Object} ShardingManagerOptions * @property {string|number} [totalShards='auto'] Number of total shards of all shard managers or "auto" * @property {string|number[]} [shardList='auto'] List of shards to spawn or "auto" * @property {ShardingManagerMode} [mode='process'] Which mode to use for shards * @property {boolean} [respawn=true] Whether shards should automatically respawn upon exiting * @property {string[]} [shardArgs=[]] Arguments to pass to the shard script when spawning * (only available when mode is set to 'process') * @property {string[]} [execArgv=[]] Arguments to pass to the shard script executable when spawning * (only available when mode is set to 'process') * @property {string} [token] Token to use for automatic shard count and passing to shards */ /** * @param {string} file Path to your shard script file * @param {ShardingManagerOptions} [options] Options for the sharding manager */ constructor(file, options = {}) { super(); options = Util.mergeDefault( { totalShards: 'auto', mode: 'process', respawn: true, shardArgs: [], execArgv: [], token: process.env.DISCORD_TOKEN, }, options, ); /** * Path to the shard script file * @type {string} */ this.file = file; if (!file) throw new Error('CLIENT_INVALID_OPTION', 'File', 'specified.'); if (!path.isAbsolute(file)) this.file = path.resolve(process.cwd(), file); const stats = fs.statSync(this.file); if (!stats.isFile()) throw new Error('CLIENT_INVALID_OPTION', 'File', 'a file'); /** * List of shards this sharding manager spawns * @type {string|number[]} */ this.shardList = options.shardList ?? 'auto'; if (this.shardList !== 'auto') { if (!Array.isArray(this.shardList)) { throw new TypeError('CLIENT_INVALID_OPTION', 'shardList', 'an array.'); } this.shardList = [...new Set(this.shardList)]; if (this.shardList.length < 1) throw new RangeError('CLIENT_INVALID_OPTION', 'shardList', 'at least 1 id.'); if ( this.shardList.some( shardId => typeof shardId !== 'number' || isNaN(shardId) || !Number.isInteger(shardId) || shardId < 0, ) ) { throw new TypeError('CLIENT_INVALID_OPTION', 'shardList', 'an array of positive integers.'); } } /** * Amount of shards that all sharding managers spawn in total * @type {number} */ this.totalShards = options.totalShards || 'auto'; if (this.totalShards !== 'auto') { if (typeof this.totalShards !== 'number' || isNaN(this.totalShards)) { throw new TypeError('CLIENT_INVALID_OPTION', 'Amount of shards', 'a number.'); } if (this.totalShards < 1) throw new RangeError('CLIENT_INVALID_OPTION', 'Amount of shards', 'at least 1.'); if (!Number.isInteger(this.totalShards)) { throw new RangeError('CLIENT_INVALID_OPTION', 'Amount of shards', 'an integer.'); } } /** * Mode for shards to spawn with * @type {ShardingManagerMode} */ this.mode = options.mode; if (this.mode !== 'process' && this.mode !== 'worker') { throw new RangeError('CLIENT_INVALID_OPTION', 'Sharding mode', '"process" or "worker"'); } /** * Whether shards should automatically respawn upon exiting * @type {boolean} */ this.respawn = options.respawn; /** * An array of arguments to pass to shards (only when {@link ShardingManager#mode} is `process`) * @type {string[]} */ this.shardArgs = options.shardArgs; /** * An array of arguments to pass to the executable (only when {@link ShardingManager#mode} is `process`) * @type {string[]} */ this.execArgv = options.execArgv; /** * Token to use for obtaining the automatic shard count, and passing to shards * @type {?string} */ this.token = options.token?.replace(/^Bot\s*/i, '') ?? null; /** * A collection of shards that this manager has spawned * @type {Collection} */ this.shards = new Collection(); process.env.SHARDING_MANAGER = true; process.env.SHARDING_MANAGER_MODE = this.mode; process.env.DISCORD_TOKEN = this.token; } /** * Creates a single shard. * Using this method is usually not necessary if you use the spawn method. * @param {number} [id=this.shards.size] Id of the shard to create * This is usually not necessary to manually specify. * @returns {Shard} Note that the created shard needs to be explicitly spawned using its spawn method. */ createShard(id = this.shards.size) { const shard = new Shard(this, id); this.shards.set(id, shard); /** * Emitted upon creating a shard. * @event ShardingManager#shardCreate * @param {Shard} shard Shard that was created */ this.emit('shardCreate', shard); return shard; } /** * Options used to spawn multiple shards. * @typedef {Object} MultipleShardSpawnOptions * @property {number|string} [amount=this.totalShards] Number of shards to spawn * @property {number} [delay=5500] How long to wait in between spawning each shard (in milliseconds) * @property {number} [timeout=30000] The amount in milliseconds to wait until the {@link Client} has become ready */ /** * Spawns multiple shards. * @param {MultipleShardSpawnOptions} [options] Options for spawning shards * @returns {Promise>} */ async spawn({ amount = this.totalShards, delay = 5500, timeout = 30_000 } = {}) { // Obtain/verify the number of shards to spawn if (amount === 'auto') { amount = await Util.fetchRecommendedShards(this.token); } else { if (typeof amount !== 'number' || isNaN(amount)) { throw new TypeError('CLIENT_INVALID_OPTION', 'Amount of shards', 'a number.'); } if (amount < 1) throw new RangeError('CLIENT_INVALID_OPTION', 'Amount of shards', 'at least 1.'); if (!Number.isInteger(amount)) { throw new TypeError('CLIENT_INVALID_OPTION', 'Amount of shards', 'an integer.'); } } // Make sure this many shards haven't already been spawned if (this.shards.size >= amount) throw new Error('SHARDING_ALREADY_SPAWNED', this.shards.size); if (this.shardList === 'auto' || this.totalShards === 'auto' || this.totalShards !== amount) { this.shardList = [...Array(amount).keys()]; } if (this.totalShards === 'auto' || this.totalShards !== amount) { this.totalShards = amount; } if (this.shardList.some(shardId => shardId >= amount)) { throw new RangeError( 'CLIENT_INVALID_OPTION', 'Amount of shards', 'bigger than the highest shardId in the shardList option.', ); } // Spawn the shards for (const shardId of this.shardList) { const promises = []; const shard = this.createShard(shardId); promises.push(shard.spawn(timeout)); if (delay > 0 && this.shards.size !== this.shardList.length) promises.push(sleep(delay)); await Promise.all(promises); // eslint-disable-line no-await-in-loop } return this.shards; } /** * Sends a message to all shards. * @param {*} message Message to be sent to the shards * @returns {Promise} */ broadcast(message) { const promises = []; for (const shard of this.shards.values()) promises.push(shard.send(message)); return Promise.all(promises); } /** * Options for {@link ShardingManager#broadcastEval} and {@link ShardClientUtil#broadcastEval}. * @typedef {Object} BroadcastEvalOptions * @property {number} [shard] Shard to run script on, all if undefined * @property {*} [context] The JSON-serializable values to call the script with */ /** * Evaluates a script on all shards, or a given shard, in the context of the {@link Client}s. * @param {Function} script JavaScript to run on each shard * @param {BroadcastEvalOptions} [options={}] The options for the broadcast * @returns {Promise<*|Array<*>>} Results of the script execution */ broadcastEval(script, options = {}) { if (typeof script !== 'function') return Promise.reject(new TypeError('SHARDING_INVALID_EVAL_BROADCAST')); return this._performOnShards('eval', [`(${script})(this, ${JSON.stringify(options.context)})`], options.shard); } /** * Fetches a client property value of each shard, or a given shard. * @param {string} prop Name of the client property to get, using periods for nesting * @param {number} [shard] Shard to fetch property from, all if undefined * @returns {Promise<*|Array<*>>} * @example * manager.fetchClientValues('guilds.cache.size') * .then(results => console.log(`${results.reduce((prev, val) => prev + val, 0)} total guilds`)) * .catch(console.error); */ fetchClientValues(prop, shard) { return this._performOnShards('fetchClientValue', [prop], shard); } /** * Runs a method with given arguments on all shards, or a given shard. * @param {string} method Method name to run on each shard * @param {Array<*>} args Arguments to pass through to the method call * @param {number} [shard] Shard to run on, all if undefined * @returns {Promise<*|Array<*>>} Results of the method execution * @private */ _performOnShards(method, args, shard) { if (this.shards.size === 0) return Promise.reject(new Error('SHARDING_NO_SHARDS')); if (typeof shard === 'number') { if (this.shards.has(shard)) return this.shards.get(shard)[method](...args); return Promise.reject(new Error('SHARDING_SHARD_NOT_FOUND', shard)); } if (this.shards.size !== this.shardList.length) return Promise.reject(new Error('SHARDING_IN_PROCESS')); const promises = []; for (const sh of this.shards.values()) promises.push(sh[method](...args)); return Promise.all(promises); } /** * Options used to respawn all shards. * @typedef {Object} MultipleShardRespawnOptions * @property {number} [shardDelay=5000] How long to wait between shards (in milliseconds) * @property {number} [respawnDelay=500] How long to wait between killing a shard's process and restarting it * (in milliseconds) * @property {number} [timeout=30000] The amount in milliseconds to wait for a shard to become ready before * continuing to another (`-1` or `Infinity` for no wait) */ /** * Kills all running shards and respawns them. * @param {MultipleShardRespawnOptions} [options] Options for respawning shards * @returns {Promise>} */ async respawnAll({ shardDelay = 5_000, respawnDelay = 500, timeout = 30_000 } = {}) { let s = 0; for (const shard of this.shards.values()) { const promises = [shard.respawn({ delay: respawnDelay, timeout })]; if (++s < this.shards.size && shardDelay > 0) promises.push(sleep(shardDelay)); await Promise.all(promises); // eslint-disable-line no-await-in-loop } return this.shards; } } module.exports = ShardingManager;