shard stuff

This commit is contained in:
Erik 2021-06-09 02:45:50 +03:00
parent a584f38f38
commit 5bf30707f4
No known key found for this signature in database
GPG Key ID: 7E862371D3409F16
2 changed files with 178 additions and 124 deletions

View File

@ -1,7 +1,7 @@
/*Adopted from Discord.js */
const path = require('path');
const EventEmitter = require('events');
const path = require('path');
const { Util } = require('../util/');
@ -21,9 +21,11 @@ class Shard extends EventEmitter {
this.id = id;
this.args = manager.shardArgs || [];
this.execArgv = manager.execArgv;
this.env = { ...process.env, SHARDING_MANAGER: true,
this.env = {
...process.env,
SHARDING_MANAGER: true,
SHARDS: this.id,
TOTAL_SHARD_COUNT: this.manager.totalShards,
SHARD_COUNT: this.manager.totalShards,
DISCORD_TOKEN: this.manager.token
};
@ -38,35 +40,64 @@ class Shard extends EventEmitter {
}
async spawn(waitForReady = true) {
async spawn(spawnTimeout = 30000) {
if (this.process) throw new Error(`[shard${this.id}] Sharding process already exists.`);
if (this.worker) throw new Error(`[shard${this.id}] Sharding worker already exists.`);
if (this.manager.mode === 'process') {
this.process = childProcess.fork(path.resolve(this.manager.file), this.args, {
this.process = childProcess
.fork(path.resolve(this.manager.file), this.args, {
env: this.env,
execArgv: this.execArgv
}).
on('message', this._handleMessage.bind(this)).
on('exit', this._exitListener);
})
.on('message', this._handleMessage.bind(this))
.on('exit', this._exitListener);
} else if (this.manager.mode === 'worker') {
this.worker = new Worker(path.resolve(this.manager.file), { workerData: this.env }).
on('message', this._handleMessage.bind(this)).
on('exit', this._exitListener);
this.worker = new Worker(path.resolve(this.manager.file), { workerData: this.env })
.on('message', this._handleMessage.bind(this))
.on('exit', this._exitListener);
}
this._evals.clear();
this._fetches.clear();
this.emit('spawn', this.process || this.worker);
if(!waitForReady) return this.process || this.worker;
if (spawnTimeout === -1 || spawnTimeout === Infinity) return this.process || this.worker;
await new Promise((resolve, reject) => {
this.once('ready', resolve);
this.once('disconnect', () => reject(new Error(`[shard${this.id}] Shard disconnected while readying.`)));
this.once('death', () => reject(new Error(`[shard${this.id}] Shard died while readying.`)));
setTimeout(() => reject(new Error(`[shard${this.id}] Shard timed out while readying.`)), 30000);
const cleanup = () => {
clearTimeout(spawnTimeoutTimer);
this.off('ready', onReady);
this.off('disconnect', onDisconnect);
this.off('death', onDeath);
};
const onReady = () => {
cleanup();
resolve();
};
const onDisconnect = () => {
cleanup();
reject(new Error(`[shard${this.id}] Shard disconnected while readying.`));
};
const onDeath = () => {
cleanup();
reject(new Error(`[shard${this.id}] Shard died while readying.`));
};
const onTimeout = () => {
cleanup();
reject(new Error(`[shard${this.id}] Shard timed out while readying.`));
};
const spawnTimeoutTimer = setTimeout(onTimeout, spawnTimeout);
this.once('ready', onReady);
this.once('disconnect', onDisconnect);
this.once('death', onDeath);
});
return this.process || this.worker;
}
kill() {
@ -79,20 +110,20 @@ class Shard extends EventEmitter {
}
this._handleExit(false);
}
async respawn(delay = 500, waitForReady = true) {
async respawn(delay = 500, spawnTimeout) {
this.kill();
if (delay > 0) await Util.delayFor(delay);
return this.spawn(waitForReady);
return this.spawn(spawnTimeout);
}
send(message) {
return new Promise((resolve, reject) => {
if (this.process) {
this.process.send(message, (error) => {
if(error) reject(error); else resolve(this);
this.process.send(message, err => {
if (err) reject(err);
else resolve(this);
});
} else {
this.worker.postMessage(message);
@ -102,6 +133,10 @@ class Shard extends EventEmitter {
}
fetchClientValue(prop) {
// Shard is dead (maybe respawning), don't cache anything and error immediately
if (!this.process && !this.worker) return Promise.reject(new Error(`[shard${this.id}] Shard process missing.`));
// Cached promise from previous call
if (this._fetches.has(prop)) return this._fetches.get(prop);
const promise = new Promise((resolve, reject) => {
@ -124,11 +159,13 @@ class Shard extends EventEmitter {
this._fetches.set(prop, promise);
return promise;
}
eval(script) {
// Shard is dead (maybe respawning), don't cache anything and error immediately
if (!this.process && !this.worker) return Promise.reject(new Error(`[shard${this.id}] Shard process missing.`));
// Cached promise from previous call
if (this._evals.has(script)) return this._evals.get(script);
const promise = new Promise((resolve, reject) => {
@ -138,7 +175,8 @@ class Shard extends EventEmitter {
if (!message || message._eval !== script) return;
child.removeListener('message', listener);
this._evals.delete(script);
if(!message._error) resolve(message._result); else reject(new Error(message._error.stack));
if (!message._error) resolve(message._result);
else reject(Util.makePlainError(message._error));
};
child.on('message', listener);
@ -152,50 +190,62 @@ class Shard extends EventEmitter {
this._evals.set(script, promise);
return promise;
}
_handleMessage(message) {
if (message) {
if(message._ready) { //Shard ready
// Shard is ready
if (message._ready) {
this.ready = true;
this.emit('ready');
return;
}
if(message._disconnect) { //Shard disconnected
// Shard has disconnected
if (message._disconnect) {
this.ready = false;
this.emit('disconnect');
return;
}
if(message._reconnecting) { //Shard attempting to reconnect
// Shard is attempting to reconnect
if (message._reconnecting) {
this.ready = false;
this.emit('reconnecting');
return;
}
if(message._sFetchProp) { //Shard requesting property fetch
this.manager.fetchClientValues(message._sFetchProp).then(
(results) => this.send({ _sFetchProp: message._sFetchProp, _result: results }),
(err) => this.send({ _sFetchProp: message._sFetchProp, _error: Util.makePlainError(err) })
// Shard is requesting a property fetch
if (message._sFetchProp) {
const resp = { _sFetchProp: message._sFetchProp, _sFetchPropShard: message._sFetchPropShard };
this.manager.fetchClientValues(message._sFetchProp, message._sFetchPropShard).then(
(results) => this.send({ ...resp, _result: results }),
(err) => this.send({ ...resp, _error: Util.makePlainError(err) })
);
return;
}
if(message._sEval) { //Shard requesting eval broadcast
this.manager.broadcastEval(message._sEval).then(
(results) => this.send({ _sEval: message._sEval, _result: results }),
(err) => this.send({ _sEval: message._sEval, _error: Util.makePlainError(err) })
// Shard is requesting an eval broadcast
if (message._sEval) {
const resp = { _sEval: message._sEval, _sEvalShard: message._sEvalShard };
this.manager.broadcastEval(message._sEval, message._sEvalShard).then(
(results) => this.send({ ...resp, _result: results }),
(err) => this.send({ ...resp, _error: Util.makePlainError(err) })
);
return;
}
if(message._sRespawnAll) { //Shard requesting to respawn all shards.
const { shardDelay, respawnDelay, waitForReady } = message._sRespawnAll;
this.manager.respawnAll(shardDelay, respawnDelay, waitForReady).catch(() => { //eslint-disable-line no-empty-function
// Shard is requesting a respawn of all shards
if (message._sRespawnAll) {
const { shardDelay, respawnDelay, spawnTimeout } = message._sRespawnAll;
this.manager.respawnAll(shardDelay, respawnDelay, spawnTimeout).catch(() => {
// Do nothing
});
return;
}
}
this.manager.emit('message', this, message);
}
_handleExit(respawn = this.manager.respawn) {
@ -208,9 +258,7 @@ class Shard extends EventEmitter {
this._fetches.clear();
if (respawn) this.spawn().catch((err) => this.emit('error', err));
}
}
module.exports = Shard;

View File

@ -1,13 +1,13 @@
/*Adopted from Discord.js */
const path = require('path');
const fs = require('fs');
const EventEmitter = require('events');
const fs = require('fs');
const path = require('path');
const Shard = require('./Shard.js');
const Shard = require('./Shard');
const { Util, Collection } = require('../util/');
class ShardManager extends EventEmitter {
class ShardingManager extends EventEmitter {
constructor(file, options = {}) {
@ -25,7 +25,6 @@ class ShardManager extends EventEmitter {
this.file = file;
if (!file) throw new Error('[shardmanager] File must be specified.');
if (!path.isAbsolute(file)) this.file = path.resolve(process.cwd(), file);
const stats = fs.statSync(this.file);
if (!stats.isFile()) throw new Error('[shardmanager] File path does not point to a valid file.');
@ -36,7 +35,7 @@ class ShardManager extends EventEmitter {
}
this.shardList = [...new Set(this.shardList)];
if (this.shardList.length < 1) throw new RangeError('[shardmanager] ShardList must have one ID.');
if(this.shardList.some((shardID) => typeof shardID !== 'number' ||
if (this.shardList.some(shardID => typeof shardID !== 'number' ||
isNaN(shardID) ||
!Number.isInteger(shardID) ||
shardID < 0)
@ -64,13 +63,12 @@ class ShardManager extends EventEmitter {
this.respawn = options.respawn;
this.shardArgs = options.shardArgs;
this.execArgv = options.execArgv;
this.token = options.token;
this.token = options.token ? options.token.replace(/^Bot\s*/i, '') : null;
this.shards = new Collection();
process.env.SHARDING_MANAGER = true;
process.env.SHARDING_MANAGER_MODE = this.mode;
process.env.DISCORD_TOKEN = this.token;
}
createShard(id = this.shards.size) {
@ -81,8 +79,7 @@ class ShardManager extends EventEmitter {
return shard;
}
async spawn(amount = this.totalShards, delay = 5500, waitForReady = true) {
async spawn(amount = this.totalShards, delay = 5500, spawnTimeout) {
if (amount === 'auto') {
amount = await Util.fetchRecommendedShards(this.token);
} else {
@ -95,6 +92,7 @@ class ShardManager extends EventEmitter {
}
}
// Make sure this many shards haven't already been spawned
if (this.shards.size >= amount) throw new Error('[shardmanager] Already spawned all necessary shards.');
if (this.shardList === 'auto' || this.totalShards === 'auto' || this.totalShards !== amount) {
this.shardList = [...Array(amount).keys()];
@ -102,20 +100,21 @@ class ShardManager extends EventEmitter {
if (this.totalShards === 'auto' || this.totalShards !== amount) {
this.totalShards = amount;
}
if(this.shardList.some((id) => id >= amount)) {
if (this.shardList.some((shardID) => shardID >= amount)) {
throw new RangeError('[shardmanager] Amount of shards cannot be larger than the highest shard ID.');
}
// Spawn the shards
for (const shardID of this.shardList) {
const promises = [];
const shard = this.createShard(shardID);
promises.push(shard.spawn(waitForReady));
if(delay > 0 && this.shards.size !== this.shardList.length - 1) promises.push(Util.delayFor(delay));
await Promise.all(promises);
promises.push(shard.spawn(spawnTimeout));
if (delay > 0 && this.shards.size !== this.shardList.length) promises.push(Util.delayFor(delay));
await Promise.all(promises); // eslint-disable-line no-await-in-loop
}
return this.shards;
}
broadcast(message) {
@ -124,30 +123,37 @@ class ShardManager extends EventEmitter {
return Promise.all(promises);
}
broadcastEval(script) {
const promises = [];
for(const shard of this.shards.values()) promises.push(shard.eval(script));
return Promise.all(promises);
broadcastEval(script, shard) {
return this._performOnShards('eval', [script], shard);
}
fetchClientValues(prop) {
fetchClientValues(prop, shard) {
return this._performOnShards('fetchClientValue', [prop], shard);
}
_performOnShards(method, args, shard) {
if (this.shards.size === 0) return Promise.reject(new Error('[shardmanager] No shards available.'));
if(this.shards.size !== this.totalShards) return Promise.reject(new Error('[shardmanager] Sharding in progress.'));
if (this.shards.size !== this.shardList.length) return Promise.reject(new Error('[shardmanager] Sharding in progress.'));
if (typeof shard === 'number') {
if (this.shards.has(shard)) return this.shards.get(shard)[method](...args);
return Promise.reject(new Error(`[shardmanager] Shard ${shard} not found.`));
}
const promises = [];
for(const shard of this.shards.values()) promises.push(shard.fetchClientValue(prop));
for (const sh of this.shards.values()) promises.push(sh[method](...args));
return Promise.all(promises);
}
async respawnAll(shardDelay = 5000, respawnDelay = 500, waitForReady = true) {
async respawnAll(shardDelay = 5000, respawnDelay = 500, spawnTimeout) {
let s = 0;
for (const shard of this.shards.values()) {
const promises = [shard.respawn(respawnDelay, waitForReady)];
const promises = [shard.respawn(respawnDelay, spawnTimeout)];
if (++s < this.shards.size && shardDelay > 0) promises.push(Util.delayFor(shardDelay));
await Promise.all(promises);
await Promise.all(promises); // eslint-disable-line no-await-in-loop
}
return this.shards;
}
}
module.exports = ShardManager;
module.exports = ShardingManager;