@dressed/ws
Version:
Communicate with the Discord Gateway
111 lines • 5.52 kB
JavaScript
import { GatewayDispatchEvents, GatewayIntentBits, } from "discord-api-types/v10";
import { getGatewayBot } from "dressed";
import { botEnv } from "dressed/utils";
import { env } from "node:process";
import { createCache } from "./cache/index.js";
import { startAutoResharder } from "./resharder.js";
import { Worker } from "node:worker_threads";
/**
* Establish a connection with the Discord Gateway
* @returns Functions for interacting with the connection and resharding
*/
export function createConnection(config = {}) {
const { intents = [], token = botEnv.DISCORD_TOKEN, shards = {}, ...rest } = config;
const listeners = new Map();
const workers = [];
let bot;
env.DISCORD_TOKEN = token;
const connection = {
...Object.fromEntries(Object.keys(GatewayDispatchEvents).map((k) => [
`on${k}`,
(callback, config = {}) => {
const id = crypto.randomUUID();
const eventName = GatewayDispatchEvents[k];
if (!listeners.has(eventName))
listeners.set(eventName, new Map());
listeners.get(eventName).set(id, { callback, config });
return () => { var _a; return (_a = listeners.get(eventName)) === null || _a === void 0 ? void 0 : _a.delete(id); };
},
])),
emit(op, d) {
for (const worker of workers) {
worker.postMessage({ type: "emit", op, d });
}
},
shards: {
workers,
numShards: 0,
isResharding: false,
cache: createCache({ getGatewayBot }),
async reshard(newShardCount) {
var _a, _b;
if (connection.shards.isResharding) {
throw new Error("Attempted to reshard while already resharding");
}
connection.shards.isResharding = true;
bot !== null && bot !== void 0 ? bot : (bot = await connection.shards.cache.getGatewayBot());
newShardCount !== null && newShardCount !== void 0 ? newShardCount : (newShardCount = bot.shards);
const prevShardCount = connection.shards.numShards;
const shardsPerWorker = (_a = shards.shardsPerWorker) !== null && _a !== void 0 ? _a : 100;
if (newShardCount === prevShardCount)
return;
const newWorkers = Math.ceil(newShardCount / shardsPerWorker);
while (newWorkers < workers.length)
(_b = workers.pop()) === null || _b === void 0 ? void 0 : _b.terminate();
const maxConcurrency = bot.session_start_limit.max_concurrency;
const numBuckets = Math.floor(newShardCount / maxConcurrency);
for (let bucketId = 0; bucketId < numBuckets; ++bucketId) {
for (let i = 0; i < maxConcurrency; ++i) {
let worker = workers[Math.floor((bucketId * maxConcurrency + i) / shardsPerWorker)];
if (!worker) {
worker = new Worker(new URL("./worker.js", import.meta.url));
worker.on("message", ({ type, t, d, shard }) => {
if (type === "dispatch" &&
t &&
shard[1] === connection.shards.numShards) {
const eventListeners = listeners.get(t);
if (eventListeners) {
for (const [id, callback] of eventListeners) {
callback.callback(d);
if (callback.config.once) {
eventListeners.delete(id);
}
}
}
}
});
workers.push(worker);
}
worker.postMessage({
type: "addShard",
config: {
token,
intents: intents.reduce((p, intent) => p | GatewayIntentBits[intent], 0),
...rest,
bot,
shard: [bucketId * maxConcurrency + i, newShardCount],
},
});
}
if (bucketId < numBuckets - 1) {
await new Promise((res) => setTimeout(res, 5000));
}
}
connection.shards.numShards = newShardCount;
for (let i = 0; i < prevShardCount; ++i) {
const worker = workers[Math.floor(i / prevShardCount)];
if (!worker)
continue;
worker.postMessage({
type: "removeShard",
shardId: `${i},${prevShardCount}`,
});
}
connection.shards.isResharding = false;
},
},
};
connection.shards.reshardInterval = startAutoResharder(connection, shards.reshardInterval, shards.shardCapacity);
return connection;
}
//# sourceMappingURL=gateway.js.map