status-sharding
Version:
Welcome to Status Sharding! This package is designed to provide an efficient and flexible solution for sharding Discord bots, allowing you to scale your bot across multiple processes or workers.
297 lines (296 loc) • 13.9 kB
JavaScript
"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.Cluster = void 0;
// core/cluster.ts
const types_1 = require("../types");
const message_1 = require("../other/message");
const shardingUtils_1 = require("../other/shardingUtils");
const message_2 = require("../handlers/message");
const utils_1 = require("../other/utils");
const worker_1 = require("../classes/worker");
const child_1 = require("../classes/child");
const events_1 = __importDefault(require("events"));
const path_1 = __importDefault(require("path"));
/** A self-contained cluster created by the ClusterManager. */
class Cluster extends events_1.default {
manager;
id;
shardList;
/** Represents whether the cluster is ready. */
ready;
/** Exited. */
exited = false;
/** Represents the child process/worker of the cluster. */
thread;
/** Represents the last time the cluster received a heartbeat. */
lastHeartbeatReceived;
/** Message processor that handles messages from the child process/worker/manager. */
messageHandler;
/** Represents the environment data of the cluster. */
envData;
/** Creates an instance of Cluster. */
constructor(manager, id, shardList) {
super();
this.manager = manager;
this.id = id;
this.shardList = shardList;
this.ready = false;
this.thread = null;
this.envData = Object.assign({}, process.env, {
CLUSTER: this.id,
SHARD_LIST: this.shardList,
TOTAL_SHARDS: this.totalShards,
CLUSTER_COUNT: this.manager.options.totalClusters,
CLUSTER_QUEUE_MODE: this.manager.options.queueOptions?.mode ?? 'auto',
CLUSTER_MANAGER_MODE: this.manager.options.mode,
});
}
/** Count of shards assigned to this cluster. */
get totalShards() {
return this.manager.options.totalShards;
}
/** Count of clusters managed by the manager. */
get totalClusters() {
return this.manager.options.totalClusters;
}
/** Spawn function that spawns the cluster's child process/worker with proper event management. */
async spawn(spawnTimeout = -1) {
if (!this.manager.file)
throw new Error('NO_FILE_PROVIDED | Cluster ' + this.id + ' does not have a file provided.');
if (this.thread?.process)
return this.thread.process;
try {
const options = {
...this.manager.options.clusterOptions,
execArgv: this.manager.options.execArgv,
env: this.envData,
args: [...(this.manager.options.shardArgs || []), '--clusterId ' + this.id, `--shards [${this.shardList.join(', ').trim()}]`],
clusterData: { ...this.envData, ...this.manager.options.clusterData },
};
this.thread = this.manager.options.mode === 'process'
? new child_1.Child(path_1.default.resolve(this.manager.file), options)
: new worker_1.Worker(path_1.default.resolve(this.manager.file), options);
this.messageHandler = new message_2.ClusterHandler(this, this.thread);
const thread = this.thread.spawn();
this._setupEventListeners(thread);
this.emit('spawn', this, this.thread.process);
const shouldWaitForReady = spawnTimeout > 0 && spawnTimeout !== Infinity;
if (shouldWaitForReady) {
await new Promise((resolve, reject) => {
const cleanup = (removeListeners = false) => {
if (spawnTimeoutTimer)
clearTimeout(spawnTimeoutTimer);
if (removeListeners) {
this.off('ready', onReady);
this.off('death', onDeath);
}
};
const onReady = () => {
this.manager.emit('clusterReady', this);
cleanup(true);
resolve();
};
const onDeath = () => {
cleanup(true);
reject(new Error('CLUSTERING_READY_DIED | Cluster ' + this.id + ' died.'));
};
const onTimeout = () => {
cleanup();
reject(new Error('CLUSTERING_READY_TIMEOUT | Cluster ' + this.id + ' took too long to get ready.'));
};
const spawnTimeoutTimer = setTimeout(onTimeout, spawnTimeout);
this.once('ready', onReady);
this.once('death', onDeath);
});
}
return this.thread.process;
}
catch (error) {
console.error(`Failed to spawn cluster ${this.id}:`, error);
throw error;
}
}
_setupEventListeners(thread) {
if (!thread)
return;
if ((0, utils_1.isChildProcess)(thread)) {
// Child process.
thread.on('disconnect', this._handleDisconnect.bind(this));
thread.on('message', this._handleMessage.bind(this));
thread.on('error', this._handleError.bind(this));
thread.on('exit', this._handleExit.bind(this));
}
else {
// Worker thread.
thread.on('messageerror', this._handleError.bind(this));
thread.on('message', this._handleMessage.bind(this));
thread.on('error', this._handleError.bind(this));
thread.on('exit', this._handleExit.bind(this));
const healthCheck = setInterval(() => {
if (!this.thread?.process || ('threadId' in this.thread.process && !this.thread.process.threadId)) {
clearInterval(healthCheck);
this._handleUnexpectedExit();
}
}, 5000);
}
}
async kill(options) {
if (!this.thread) {
console.warn(`Cluster ${this.id} has no thread to kill.`);
return;
}
try {
const killResult = await this.thread.kill();
this.thread = null;
this.ready = false;
this.exited = true;
this.manager.heartbeat?.removeCluster(this.id);
this.manager._debug('[KILL] Cluster ' + this.id + ' killed with reason: ' + (options?.reason || 'Unknown reason.'));
if (!killResult)
console.warn(`Cluster ${this.id} kill operation completed but process may not have terminated cleanly.`);
}
catch (error) {
console.error(`Error killing cluster ${this.id}:`, error);
this.thread = null;
this.ready = false;
this.exited = true;
this.manager.heartbeat?.removeCluster(this.id);
}
}
/** Respawn function that respawns the cluster's child process/worker. */
async respawn(delay = this.manager.options.spawnOptions.delay || 5500, timeout = this.manager.options.spawnOptions.timeout || -1) {
this.ready = false;
this.exited = false;
if (this.thread)
await this.kill();
if (delay > 0)
await shardingUtils_1.ShardingUtils.delayFor(delay);
return this.spawn(timeout);
}
/** Send function that sends a message to the cluster's child process/worker. */
async send(message) {
if (!this.thread)
return Promise.reject(new Error('CLUSTERING_NO_CHILD_EXISTS | Cluster ' + this.id + ' does not have a child process/worker (#2).'));
this.manager._debug(`[IPC] [Cluster ${this.id}] Sending message to child.`);
return this.thread.send({
_type: types_1.MessageTypes.CustomMessage,
data: message,
});
}
/** Request function that sends a message to the cluster's child process/worker and waits for a response. */
async request(message, options = {}) {
if (!this.thread)
return Promise.reject(new Error('CLUSTERING_NO_CHILD_EXISTS | Cluster ' + this.id + ' does not have a child process/worker (#3).'));
const nonce = shardingUtils_1.ShardingUtils.generateNonce();
this.thread.send({
_type: types_1.MessageTypes.CustomRequest,
_nonce: nonce,
data: message,
});
return this.manager.promise.create(nonce, options.timeout);
}
/** Broadcast function that sends a message to all clusters. */
async broadcast(message, sendSelf = false) {
return await this.manager.broadcast(message, sendSelf ? undefined : [this.id]);
}
/** Eval function that evaluates a script on the current cluster. */
async eval(script, options) {
return eval(shardingUtils_1.ShardingUtils.parseInput(script, options?.context, this.manager.options.packageType, 'this'));
}
/** EvalOnClient function that evaluates a script on a specific cluster. */
async evalOnClient(script, options) {
if (!this.thread)
return Promise.reject(new Error('CLUSTERING_NO_CHILD_EXISTS | Cluster ' + this.id + ' does not have a child process/worker (#4).'));
const nonce = shardingUtils_1.ShardingUtils.generateNonce();
this.thread.send({
_type: types_1.MessageTypes.ClientEvalRequest,
_nonce: nonce,
data: {
script: shardingUtils_1.ShardingUtils.parseInput(script, options?.context),
options: options,
},
});
return this.manager.promise.create(nonce, options?.timeout);
}
/** EvalOnCluster function that evaluates a script on a specific cluster. */
async evalOnGuild(guildId, script, options) {
if (!this.thread)
return Promise.reject(new Error('CLUSTERING_NO_CHILD_EXISTS | Cluster ' + this.id + ' does not have a child process/worker (#5).'));
else if (this.manager.options.packageType !== 'discord.js')
return Promise.reject(new Error('CLUSTERING_EVAL_GUILD_UNSUPPORTED | evalOnGuild is only supported in discord.js package type.'));
return this.manager.evalOnGuild(guildId, script, options);
}
/** Function that allows you to construct your own BaseMessage and send it to the cluster. */
_sendInstance(message) {
if (!this.thread)
return Promise.reject(new Error('CLUSTERING_NO_CHILD_EXISTS | Cluster ' + this.id + ' does not have a child process/worker (#6).'));
this.emit('debug', `[IPC] [Child ${this.id}] Sending message to cluster.`);
return this.thread.send(message);
}
/** Message handler function that handles messages from the cluster's child process/worker/manager. */
_handleMessage(message) {
if (!message || '_data' in message)
return this.manager.broker.handleMessage(message);
else if (!this.messageHandler)
throw new Error('CLUSTERING_NO_MESSAGE_HANDLER | Cluster ' + this.id + ' does not have a message handler.');
if (this.manager.options.advanced?.logMessagesInDebug) {
this.manager._debug(`[IPC] [Cluster ${this.id}] Received message from child.`);
}
this.messageHandler.handleMessage(message);
if ([types_1.MessageTypes.CustomMessage, types_1.MessageTypes.CustomRequest].includes(message._type)) {
const ipcMessage = new message_1.ProcessMessage(this, message);
if (message._type === types_1.MessageTypes.CustomRequest) {
this.manager.emit('clientRequest', ipcMessage);
}
this.emit('message', ipcMessage);
this.manager.emit('message', ipcMessage);
}
}
/** Exit handler function that handles the cluster's child process/worker exiting. */
_handleExit(exitCode, signal) {
this.manager._debug(`[Cluster ${this.id}] Process exited with code ${exitCode}, signal ${signal}`);
if (!this.exited)
this.emit('death', this, this.thread?.process || null);
this.ready = false;
this.exited = true;
this.thread = null;
this.manager.heartbeat?.removeCluster(this.id);
if (!this.manager.heartbeat) {
if (this.manager.options.respawn && exitCode !== 0 && exitCode !== null) {
this.respawn().catch((err) => {
this.manager._debug(`[Cluster ${this.id}] Failed to respawn: ${err.message}`);
});
}
}
}
/** Error handler function that handles errors from the cluster's child process/worker/manager. */
_handleError(error) {
this.manager.emit('error', error);
}
/** Handle unexpected disconnection. */
_handleDisconnect() {
this.manager._debug(`[Cluster ${this.id}] Process disconnected unexpectedly.`);
this._handleUnexpectedExit();
}
/** Handle unexpected exit/crash. */
_handleUnexpectedExit() {
if (!this.exited && this.ready) {
this.manager._debug(`[Cluster ${this.id}] Detected unexpected exit/crash.`);
this.emit('death', this, this.thread?.process || null);
this.ready = false;
this.exited = true;
this.thread = null;
this.manager.heartbeat?.removeCluster(this.id);
if (this.manager.options.respawn) {
this.manager._debug(`[Cluster ${this.id}] Scheduling respawn after crash.`);
this.respawn().catch((err) => {
this.manager._debug(`[Cluster ${this.id}] Failed to respawn after crash: ${err.message}`);
});
}
}
}
}
exports.Cluster = Cluster;