galactic-bot/src/middleware/shard/ShardingManager.ts

/* Adopted from Discord.js v13 */

// const { EventEmitter } = require('events');
// const { Collection } = require('@discordjs/collection');
// const { Util } = require('../../utilities');
// const fs = require('fs');
// const path = require('path');

// const Shard = require('./Shard.js');

import { EventEmitter } from 'node:events';
import { Collection } from '@discordjs/collection';
import { Util } from '../../utilities';
import fs from 'node:fs';
import path from 'node:path';

import Shard from './Shard';
import { IPCMessage } from '../../@types/Shared';
import { BroadcastEvalOptions, ShardingOptions, ShardMethod } from '../../@types/Shard';

class ShardingManager extends EventEmitter {

    #file: string;
    #shardList: 'auto' | number[];
    #totalShards: 'auto' | number;
    #mode: 'worker' | 'process';
    #respawn: boolean;
    #shardArgs: string[];
    #execArgv: string[];
    #token: string | null;
    #shards: Collection<number, Shard>;

    constructor (file: string, options: ShardingOptions = {}) {
      
        super();

        options = Util.mergeDefault(
            {
                totalShards: 'auto',
                mode: 'process',
                respawn: true,
                shardArgs: [],
                execArgv: [],
                token: process.env.DISCORD_TOKEN
            },
            options
        );

        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.');

        this.#shardList = options.shardList ?? 'auto';
        if (this.#shardList !== 'auto') {
            if (!Array.isArray(this.#shardList)) {
                throw new TypeError('[shardmanager] ShardList must be an array.');
            }
            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' || isNaN(shardId) || !Number.isInteger(shardId) || shardId < 0)) {
                throw new TypeError('[shardmanager] ShardList must be an array of positive integers.');
            }
        }

        this.#totalShards = options.totalShards || 'auto';
        if (this.#totalShards !== 'auto') {
            if (typeof this.#totalShards !== 'number' || isNaN(this.#totalShards)) {
                throw new TypeError('[shardmanager] totalShards must be an integer.');
            }
            if (this.#totalShards < 1) 
                throw new RangeError('[shardmanager] totalShards must be at least one.');
            if (!Number.isInteger(this.#totalShards)) {
                throw new RangeError('[shardmanager] totalShards must be an integer.');
            }
        }

        this.#mode = options.mode as 'worker' | 'process';
        if (this.#mode !== 'process' && this.#mode !== 'worker') {
            throw new RangeError('[shardmanager] Mode must be either \'worker\' or \'process\'.');
        }

        this.#respawn = options.respawn as boolean;
        this.#shardArgs = options.shardArgs as string[];
        this.#execArgv = options.execArgv as string[];

        this.#token = options.token?.replace(/^Bot\s*/iu, '') ?? null;

        this.#shards = new Collection();

        process.env.SHARDING_MANAGER = true as unknown as string;
        process.env.SHARDING_MANAGER_MODE = this.#mode;
        process.env.DISCORD_TOKEN = this.#token || '';

    }

    get respawn () {
        return this.#respawn;
    }

    get shardArgs () {
        return this.#shardArgs;
    }

    get execArgv () {
        return this.#execArgv;
    }

    get totalShards () {
        return this.#totalShards as number;
    }

    get token () {
        return this.#token || '';
    }

    get mode () {
        return this.#mode;
    }

    get file () {
        return this.#file;
    }

    get shards () {
        return this.#shards.clone();
    }

    createShard (id = this.#shards.size) {
        const shard = new Shard(this, id);
        this.#shards.set(id, shard);
        this.emit('shardCreate', shard);
        return shard;
    }

    async spawn ({ amount = this.#totalShards, delay = 5500, timeout = 30000 } = {}) {
        if (amount === 'auto') {
            if (!this.#token) 
                throw new Error('[shardmanager] Must supply token for auto shard amount');
            amount = await Util.fetchRecommendedShards(this.#token);
        } else {
            if (typeof amount !== 'number' || isNaN(amount)) {
                throw new TypeError('[shardmanager] Amount of shards must be a number.');
            }
            if (amount < 1) 
                throw new RangeError('[shardmanager] Amount of shards must be at least one.');
            if (!Number.isInteger(amount)) {
                throw new TypeError('[shardmanager] Amount of shards must be an integer.');
            }
        }

        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() ];
        }
        if (this.#totalShards === 'auto' || this.#totalShards !== amount) {
            this.#totalShards = amount;
        }

        if (this.#shardList.some((shardId) => shardId >= (amount as number))) {
            throw new RangeError('[shardmanager] Amount of shards cannot be larger than the highest shard ID.');
        }

        for (const shardId of this.#shardList) {
            const shard = this.createShard(shardId);
            await shard.spawn(timeout);
            if (delay > 0 && this.#shards.size !== this.#shardList.length)
                await Util.delayFor(delay);
        }

        return this.#shards;
    }

    broadcast (message: IPCMessage) {
        const promises = [];
        for (const shard of this.#shards.values()) 
            promises.push(shard.send(message));
        return Promise.all(promises);
    }

    // eslint-disable-next-line @typescript-eslint/ban-types
    broadcastEval (script: string | Function, options: BroadcastEvalOptions = {}) {
        if (typeof script !== 'function') 
            return Promise.reject(new TypeError('[shardmanager] Provided eval must be a function.'));
        return this._performOnShards('eval', [ `(${script})(this, ${JSON.stringify(options.context)})` ], options.shard);
    }

    fetchClientValues (prop: string, shard: number) {
        return this._performOnShards('fetchClientValue', [ prop ], shard);
    }

    _performOnShards (method: ShardMethod, args: [string, object?], shard?: number): Promise<unknown> {
        if (this.#shards.size === 0) 
            return Promise.reject(new Error('[shardmanager] No shards available.'));

        if (typeof shard === 'number') {
            if (!this.#shards.has(shard))
                Promise.reject(new Error('[shardmanager] Shard not found.'));
            
            const s = this.#shards.get(shard) as Shard;
            if (method === 'eval')
                return s.eval(...args);
            else if (method === 'fetchClientValue')
                return s.eval(args[0]);
        }

        if (this.#shards.size !== this.#shardList.length) 
            return Promise.reject(new Error('[shardmanager] Sharding in progress.'));

        const promises = [];
        for (const sh of this.#shards.values()) {
            if (method === 'eval')
                promises.push(sh.eval(...args));
            else if (method === 'fetchClientValue')
                promises.push(sh.eval(args[0]));
        }
        return Promise.all(promises);
    }

    async respawnAll ({ shardDelay = 5000, respawnDelay = 500, timeout = 30000 } = {}) {
        let s = 0;
        for (const shard of this.#shards.values()) {
            const promises: Promise<unknown>[] = [ shard.respawn({ delay: respawnDelay, timeout }) ];
            if (++s < this.#shards.size && shardDelay > 0) 
                promises.push(Util.delayFor(shardDelay));
            await Promise.all(promises); // eslint-disable-line no-await-in-loop
        }
        return this.#shards;
    }

}

// module.exports = ShardingManager;
export default ShardingManager;