seyfert
Version:
The most advanced framework for discord bots
548 lines (547 loc) • 22.5 kB
JavaScript
;
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.WorkerManager = void 0;
const node_cluster_1 = __importDefault(require("node:cluster"));
const node_crypto_1 = require("node:crypto");
const __1 = require("../..");
const cache_1 = require("../../cache");
const base_1 = require("../../client/base");
const common_1 = require("../../common");
const constants_1 = require("../constants");
const structures_1 = require("../structures");
const timeout_1 = require("../structures/timeout");
class WorkerManager extends Map {
static prepareSpaces(options, logger) {
logger?.info('Preparing buckets');
const chunks = structures_1.DynamicBucket.chunk(new Array(options.shardEnd - options.shardStart), options.shardsPerWorker);
chunks.forEach((shards, index) => {
for (let i = 0; i < shards.length; i++) {
const id = i + (index > 0 ? index * options.shardsPerWorker : 0) + options.shardStart;
chunks[index][i] = id;
}
});
logger?.info(`${chunks.length} buckets created`);
return chunks;
}
options;
debugger;
connectQueue;
workerQueue = [];
cacheAdapter;
promises = new Map();
rest;
reshardingWorkerQueue = [];
_info;
constructor(options) {
super();
this.options = options;
this.cacheAdapter = new cache_1.MemoryAdapter();
if (this.options.handleWorkerMessage) {
const oldFn = this.handleWorkerMessage.bind(this);
this.handleWorkerMessage = async (message) => {
await this.options.handleWorkerMessage(message);
return oldFn(message);
};
}
}
setCache(adapter) {
this.cacheAdapter = adapter;
}
setRest(rest) {
this.rest = rest;
}
get remaining() {
return this.options.info.session_start_limit.remaining;
}
get concurrency() {
return this.options.info.session_start_limit.max_concurrency;
}
get totalWorkers() {
return this.options.workers;
}
get totalShards() {
return this.options.totalShards ?? this.options.info.shards;
}
get shardStart() {
return this.options.shardStart ?? 0;
}
get shardEnd() {
return this.options.shardEnd ?? this.totalShards;
}
get shardsPerWorker() {
return this.options.shardsPerWorker;
}
async syncLatency({ shardId, workerId, }) {
if (typeof shardId !== 'number' && typeof workerId !== 'number') {
throw new Error('Undefined workerId and shardId');
}
const id = workerId ?? this.calculateWorkerId(shardId);
if (!this.has(id)) {
throw new Error(`Worker #${workerId} doesnt exist`);
}
const data = await this.getWorkerInfo(id);
return data.shards.reduce((acc, prv) => acc + prv.latency, 0) / data.shards.length;
}
calculateShardId(guildId) {
return Number((BigInt(guildId) >> 22n) % BigInt(this.totalShards));
}
calculateWorkerId(shardId) {
const workerId = Math.floor((shardId - this.shardStart) / this.shardsPerWorker);
if (workerId >= this.totalWorkers) {
throw new Error('Invalid shardId');
}
return workerId;
}
postMessage(id, body) {
const worker = this.get(id);
if (!worker)
return this.debugger?.error(`Worker ${id} does not exists.`);
switch (this.options.mode) {
case 'clusters':
worker.send(body);
break;
case 'threads':
worker.postMessage(body);
break;
case 'custom':
this.options.adapter.postMessage(id, body);
break;
}
}
prepareWorkers(shards, resharding = false) {
const worker_threads = (0, common_1.lazyLoadPackage)('node:worker_threads');
if (!worker_threads)
throw new Error('Cannot prepare workers without worker_threads.');
for (let i = 0; i < shards.length; i++) {
const workerExists = this.has(i);
if (resharding || !workerExists) {
this[resharding ? 'reshardingWorkerQueue' : 'workerQueue'].push(() => {
const worker = this.createWorker({
path: this.options.path,
debug: this.options.debug,
token: this.options.token,
shards: shards[i],
intents: this.options.intents,
workerId: i,
workerProxy: this.options.workerProxy,
totalShards: resharding ? this._info.shards : this.totalShards,
mode: this.options.mode,
resharding,
totalWorkers: shards.length,
info: {
...this.options.info,
shards: this.totalShards,
},
compress: this.options.compress,
});
this.set(i, worker);
});
}
}
}
createWorker(workerData) {
if (this.has(workerData.workerId)) {
if (workerData.resharding) {
this.postMessage(workerData.workerId, {
type: 'WORKER_ALREADY_EXISTS_RESHARDING',
});
}
const worker = this.get(workerData.workerId);
return worker;
}
const worker_threads = (0, common_1.lazyLoadPackage)('node:worker_threads');
if (!worker_threads)
throw new Error('Cannot create worker without worker_threads.');
const env = {
SEYFERT_SPAWNING: 'true',
};
if (workerData.resharding)
env.SEYFERT_WORKER_RESHARDING = 'true';
for (const i in workerData) {
const data = workerData[i];
env[`SEYFERT_WORKER_${i.toUpperCase()}`] = typeof data === 'object' && data ? JSON.stringify(data) : data;
}
switch (this.options.mode) {
case 'threads': {
const worker = new worker_threads.Worker(workerData.path, {
env,
});
worker.on('message', data => this.handleWorkerMessage(data));
return worker;
}
case 'clusters': {
node_cluster_1.default.setupPrimary({
exec: workerData.path,
});
const worker = node_cluster_1.default.fork(env);
worker.on('message', data => this.handleWorkerMessage(data));
return worker;
}
case 'custom': {
this.options.adapter.spawn(workerData, env);
return {
ready: false,
};
}
}
}
spawn(workerId, shardId, resharding = false) {
this.connectQueue.push(() => {
const worker = this.has(workerId);
if (!worker) {
this.debugger?.fatal(`Trying ${resharding ? 'reshard' : 'spawn'} with worker that doesn't exist`);
return;
}
this.postMessage(workerId, {
type: resharding ? 'ALLOW_CONNECT_RESHARDING' : 'ALLOW_CONNECT',
shardId,
presence: this.options.presence?.(shardId, workerId),
});
});
}
async handleWorkerMessage(message) {
switch (message.type) {
case 'WORKER_READY_RESHARDING':
{
this.get(message.workerId).resharded = true;
if (!this.reshardingWorkerQueue.length && [...this.values()].every(w => w.resharded)) {
for (const [id] of this.entries()) {
this.postMessage(id, {
type: 'DISCONNECT_ALL_SHARDS_RESHARDING',
});
}
this.forEach(w => {
delete w.resharded;
});
}
else {
const nextWorker = this.reshardingWorkerQueue.shift();
if (nextWorker) {
this.debugger?.info('Spawning next worker to reshard');
nextWorker();
}
else {
this.debugger?.info('No more workers to reshard left');
}
}
}
break;
case 'DISCONNECTED_ALL_SHARDS_RESHARDING':
{
this.get(message.workerId).disconnected = true;
if ([...this.values()].every(w => w.disconnected)) {
this.options.totalShards = this._info.shards;
this.options.shardEnd = this.options.totalShards = this.options.info.shards = this._info.shards;
this.options.workers = this.size;
delete this._info;
for (const [id] of this.entries()) {
this.postMessage(id, {
type: 'CONNECT_ALL_SHARDS_RESHARDING',
});
}
this.forEach(w => {
delete w.disconnected;
});
}
}
break;
case 'WORKER_START_RESHARDING':
{
this.postMessage(message.workerId, {
type: 'SPAWN_SHARDS_RESHARDING',
compress: this.options.compress ?? false,
info: {
...this.options.info,
shards: this._info.shards,
},
properties: {
...constants_1.properties,
...this.options.properties,
},
});
}
break;
case 'WORKER_START':
{
this.postMessage(message.workerId, {
type: 'SPAWN_SHARDS',
compress: this.options.compress ?? false,
info: {
...this.options.info,
shards: this.totalShards,
},
properties: {
...constants_1.properties,
...this.options.properties,
},
});
}
break;
case 'CONNECT_QUEUE_RESHARDING':
this.spawn(message.workerId, message.shardId, true);
break;
case 'CONNECT_QUEUE':
this.spawn(message.workerId, message.shardId);
break;
case 'CACHE_REQUEST':
{
const worker = this.has(message.workerId);
if (!worker) {
throw new Error('Invalid request from unavailable worker');
}
// @ts-expect-error
const result = await this.cacheAdapter[message.method](...message.args);
this.postMessage(message.workerId, {
type: 'CACHE_RESULT',
nonce: message.nonce,
result,
});
}
break;
case 'RECEIVE_PAYLOAD':
await this.options.handlePayload?.(message.shardId, message.workerId, message.payload);
break;
case 'RESULT_PAYLOAD':
{
const resultPayload = this.promises.get(message.nonce);
if (!resultPayload) {
return;
}
this.promises.delete(message.nonce);
clearTimeout(resultPayload.timeout);
resultPayload.resolve(true);
}
break;
case 'SHARD_INFO':
{
const { nonce, type, ...data } = message;
const shardInfo = this.promises.get(nonce);
if (!shardInfo) {
return;
}
this.promises.delete(nonce);
clearTimeout(shardInfo.timeout);
shardInfo.resolve(data);
}
break;
case 'WORKER_INFO':
{
const { nonce, type, ...data } = message;
const workerInfo = this.promises.get(nonce);
if (!workerInfo) {
return;
}
this.promises.delete(nonce);
clearTimeout(workerInfo.timeout);
workerInfo.resolve(data);
}
break;
case 'WORKER_READY':
{
this.get(message.workerId).ready = true;
if (this.size === this.totalWorkers && [...this.values()].every(w => w.ready)) {
this.postMessage(this.keys().next().value, {
type: 'BOT_READY',
});
this.forEach(w => {
delete w.ready;
});
}
}
break;
case 'WORKER_SHARDS_CONNECTED':
{
const nextWorker = this.workerQueue.shift();
if (nextWorker) {
this.debugger?.info('Spawning next worker');
nextWorker();
}
else {
this.debugger?.info('No more workers to spawn left');
}
}
break;
case 'WORKER_API_REQUEST':
{
if (this.options.mode === 'clusters' && message.requestOptions.files?.length) {
message.requestOptions.files.forEach(file => {
//@ts-expect-error
if (file.data.type === 'Buffer' && Array.isArray(file.data?.data))
//@ts-expect-error
file.data = new Uint8Array(file.data.data);
});
}
const response = await this.rest.request(message.method, message.url, message.requestOptions);
this.postMessage(message.workerId, {
nonce: message.nonce,
response,
type: 'API_RESPONSE',
});
}
break;
case 'EVAL_RESPONSE':
{
const { nonce, response } = message;
const evalResponse = this.promises.get(nonce);
if (!evalResponse) {
return;
}
this.promises.delete(nonce);
clearTimeout(evalResponse.timeout);
evalResponse.resolve(response);
}
break;
case 'EVAL_TO_WORKER':
{
const nonce = this.generateNonce();
this.postMessage(message.toWorkerId, {
nonce,
func: message.func,
type: 'EXECUTE_EVAL_TO_WORKER',
toWorkerId: message.toWorkerId,
vars: message.vars,
});
this.generateSendPromise(nonce, 'Eval timeout').then(val => this.postMessage(message.workerId, {
nonce: message.nonce,
response: val,
type: 'EVAL_RESPONSE',
}));
}
break;
}
}
generateNonce() {
const uuid = (0, node_crypto_1.randomUUID)();
if (this.promises.has(uuid))
return this.generateNonce();
return uuid;
}
generateSendPromise(nonce, message = 'Timeout') {
return new Promise((res, rej) => {
const timeout = setTimeout(() => {
this.promises.delete(nonce);
rej(new Error(message));
}, 60e3);
this.promises.set(nonce, { resolve: res, timeout });
});
}
async send(data, shardId) {
const workerId = this.calculateWorkerId(shardId);
const worker = this.has(workerId);
if (!worker) {
throw new Error(`Worker #${workerId} doesnt exist`);
}
const nonce = this.generateNonce();
this.postMessage(workerId, {
type: 'SEND_PAYLOAD',
shardId,
nonce,
...data,
});
return this.generateSendPromise(nonce, 'Shard send payload timeout');
}
async getShardInfo(shardId) {
const workerId = this.calculateWorkerId(shardId);
const worker = this.has(workerId);
if (!worker) {
throw new Error(`Worker #${workerId} doesnt exist`);
}
const nonce = this.generateNonce();
this.postMessage(workerId, { shardId, nonce, type: 'SHARD_INFO' });
return this.generateSendPromise(nonce, 'Get shard info timeout');
}
async getWorkerInfo(workerId) {
const worker = this.has(workerId);
if (!worker) {
throw new Error(`Worker #${workerId} doesnt exist`);
}
const nonce = this.generateNonce();
this.postMessage(workerId, { nonce, type: 'WORKER_INFO' });
return this.generateSendPromise(nonce, 'Get worker info timeout');
}
tellWorker(workerId, func, vars) {
const nonce = this.generateNonce();
this.postMessage(workerId, {
type: 'EXECUTE_EVAL',
func: func.toString(),
nonce,
vars: JSON.stringify(vars),
});
return this.generateSendPromise(nonce);
}
tellWorkers(func, vars) {
const promises = [];
for (const i of this.keys()) {
promises.push(this.tellWorker(i, func, vars));
}
return Promise.all(promises);
}
async start() {
const rc = (await this.options.getRC?.()) ??
(await base_1.BaseClient.prototype.getRC());
this.options.debug ||= rc.debug ?? false;
this.options.intents ||= rc.intents ?? 0;
this.options.token ??= rc.token;
this.rest ??= new __1.ApiHandler({
token: this.options.token,
baseUrl: 'api/v10',
domain: common_1.BASE_HOST,
debug: this.options.debug,
});
this.options.info ??= await this.rest.proxy.gateway.bot.get();
this.options.shardEnd ??= this.options.totalShards ?? this.options.info.shards;
this.options.totalShards ??= this.options.shardEnd;
this.options = (0, common_1.MergeOptions)(constants_1.WorkerManagerDefaults, this.options);
this.options.resharding.getInfo ??= () => this.rest.proxy.gateway.bot.get();
this.options.workers ??= Math.ceil(this.options.totalShards / this.options.shardsPerWorker);
this.connectQueue = new timeout_1.ConnectQueue(5.5e3, this.concurrency);
if (this.options.debug) {
this.debugger = new __1.Logger({
name: '[WorkerManager]',
});
}
if (this.totalShards / this.shardsPerWorker > this.totalWorkers) {
throw new Error(`Cannot create enough shards in the specified workers, minimum: ${Math.ceil(this.totalShards / this.shardsPerWorker)}`);
}
const spaces = WorkerManager.prepareSpaces({
shardStart: this.shardStart,
shardEnd: this.shardEnd,
shardsPerWorker: this.shardsPerWorker,
}, this.debugger);
this.prepareWorkers(spaces);
// Start workers queue
this.workerQueue.shift()();
await this.startResharding();
}
async startResharding() {
if (this.options.resharding.interval <= 0)
return;
if (this.shardStart !== 0 || this.shardEnd !== this.totalShards)
return this.debugger?.debug('Cannot start resharder');
setInterval(async () => {
this.debugger?.debug('Checking if reshard is needed');
const info = await this.options.resharding.getInfo();
if (info.shards <= this.totalShards)
return this.debugger?.debug('Resharding not needed');
//https://github.com/discordeno/discordeno/blob/6a5f446c0651b9fad9f1550ff1857fe7a026426b/packages/gateway/src/manager.ts#L106C8-L106C94
const percentage = (info.shards / ((this.totalShards * 2500) / 1000)) * 100;
if (percentage < this.options.resharding.percentage)
return this.debugger?.debug(`Percentage is not enough to reshard ${percentage}/${this.options.resharding.percentage}`);
this.debugger?.info(`Starting resharding process to ${info.shards}`);
this._info = info;
this.connectQueue.concurrency = info.session_start_limit.max_concurrency;
this.options.info.session_start_limit.max_concurrency = info.session_start_limit.max_concurrency;
const spaces = WorkerManager.prepareSpaces({
shardsPerWorker: this.shardsPerWorker,
shardEnd: info.shards,
shardStart: 0,
}, this.debugger);
this.prepareWorkers(spaces, true);
return this.reshardingWorkerQueue.shift()();
}, this.options.resharding.interval);
}
}
exports.WorkerManager = WorkerManager;