discordeno/examples/bigbot/src/gateway/gatewayManager.ts

import type { Worker } from 'node:worker_threads';
import { createGatewayManager, createLogger, createRestManager } from '@discordeno/bot';
import { DISCORD_TOKEN, GATEWAY_INTENTS, REST_AUTHORIZATION, REST_URL, SHARDS_PER_WORKER, TOTAL_SHARDS, TOTAL_WORKERS } from '../config.js';
import { promiseWithResolvers } from '../util.js';
import { createWorker } from './worker/createWorker.js';
import type { ManagerMessage, WorkerMessage } from './worker/types.js';

export const workers = new Map<number, Worker>();
export const logger = createLogger({ name: 'GATEWAY' });

const restManager = createRestManager({
  token: DISCORD_TOKEN,
  proxy: {
    baseUrl: REST_URL,
    authorization: REST_AUTHORIZATION,
  },
});

const gatewayBotConfig = await restManager.getGatewayBot();

const gatewayManager = createGatewayManager({
  token: DISCORD_TOKEN,
  intents: GATEWAY_INTENTS,
  connection: gatewayBotConfig,
  shardsPerWorker: SHARDS_PER_WORKER,
  totalShards: TOTAL_SHARDS,
  totalWorkers: TOTAL_WORKERS,
  resharding: {
    getSessionInfo: restManager.getGatewayBot,
  },
});

gatewayManager.resharding.tellWorkerToPrepare = async (workerId, shardId, bucketId) => {
  logger.info(`Tell worker to prepare, workerId: ${workerId}, shardId: ${shardId}, bucketId: ${bucketId}`);

  let worker = workers.get(workerId);
  if (!worker) {
    worker = createWorker(workerId);
    workers.set(workerId, worker);
  }

  worker.postMessage({
    type: 'PrepareShard',
    shardId,
    totalShards: gatewayManager.totalShards,
  } satisfies WorkerMessage);

  const { promise, resolve } = promiseWithResolvers<void>();

  const waitForShardPrepared = (message: ManagerMessage) => {
    if (message.type === 'ShardPrepared' && message.shardId === shardId) {
      resolve();
    }
  };

  worker.on('message', waitForShardPrepared);

  await promise;

  worker.off('message', waitForShardPrepared);
};

gatewayManager.resharding.onReshardingSwitch = async () => {
  logger.info('Resharding switch triggered, telling workers to switch the shards');

  for (const worker of workers.values()) {
    worker.postMessage({
      type: 'SwitchShards',
    } satisfies WorkerMessage);
  }
};

gatewayManager.tellWorkerToIdentify = async (workerId, shardId, bucketId) => {
  logger.info(`Tell worker to identify, workerId: ${workerId}, shardId: ${shardId}, bucketId: ${bucketId}`);

  const worker = workers.get(workerId) ?? createWorker(workerId);
  workers.set(workerId, worker);

  worker.postMessage({
    type: 'IdentifyShard',
    shardId,
  } satisfies WorkerMessage);

  const { promise, resolve } = promiseWithResolvers<void>();

  const waitForShardIdentified = (message: ManagerMessage) => {
    if (message.type === 'ShardIdentified' && message.shardId === shardId) {
      resolve();
    }
  };

  worker.on('message', waitForShardIdentified);

  await promise;

  worker.off('message', waitForShardIdentified);
};

gatewayManager.sendPayload = async (shardId, payload) => {
  const workerId = gatewayManager.calculateWorkerId(shardId);
  const worker = workers.get(workerId);

  if (!worker) return;

  worker.postMessage({
    type: 'ShardPayload',
    shardId,
    payload,
  } satisfies WorkerMessage);
};

gatewayManager.editBotStatus = async (payload) => {
  const workersArray = Array.from(workers.values());

  for (const worker of workersArray) {
    worker.postMessage({
      type: 'EditShardsPresence',
      payload,
    } satisfies WorkerMessage);
  }
};

export default gatewayManager;