From 107ef59fd2a69d35c06011135432748d35f2b96a Mon Sep 17 00:00:00 2001 From: Skillz Date: Thu, 16 Jul 2020 20:41:25 -0400 Subject: [PATCH] fix WS member& bot status stuff for basic sharding --- module/basicShard.ts | 130 ++++++++++++++++++++++++++++++++++++++ module/client.ts | 2 +- module/shardingManager.ts | 26 +++++++- types/options.ts | 1 + utils/utils.ts | 8 +++ 5 files changed, 163 insertions(+), 4 deletions(-) diff --git a/module/basicShard.ts b/module/basicShard.ts index 55dfa2f31..087cdad7b 100644 --- a/module/basicShard.ts +++ b/module/basicShard.ts @@ -17,6 +17,8 @@ import { import { DiscordHeartbeatPayload } from "../types/discord.ts"; import { logRed } from "../utils/logger.ts"; import { handleDiscordPayload } from "./shardingManager.ts"; +import { FetchMembersOptions } from "../types/guild.ts"; +import { BotStatusRequest } from "../utils/utils.ts"; const basicShards = new Map(); @@ -29,6 +31,16 @@ export interface BasicShard { needToResume: boolean; } +const RequestMembersQueue: RequestMemberQueuedRequest[] = []; +let processQueue = false; + +interface RequestMemberQueuedRequest { + guildID: string; + shardID: number; + nonce: string; + options?: FetchMembersOptions; +} + export async function createBasicShard( data: DiscordBotGatewayData, identifyPayload: IdentifyPayload, @@ -163,6 +175,8 @@ async function heartbeat( interval: number, ) { await delay(interval); + if (shard.socket.isClosed) return; + shard.socket.send( JSON.stringify( { op: GatewayOpcode.Heartbeat, d: shard.previousSequenceNumber }, @@ -201,3 +215,119 @@ async function resumeConnection( await delay(1000 * 15); if (shard.needToResume) resumeConnection(botGatewayData, payload); } + +export function requestGuildMembers( + guildID: string, + shardID: number, + nonce: string, + options?: FetchMembersOptions, + queuedRequest = false, +) { + const shard = basicShards.get(shardID); + + // This request was not from this queue so we add it to queue first + if (!queuedRequest) { + RequestMembersQueue.push({ + guildID, + shardID, + nonce, + options, + }); + + if (!processQueue) { + processQueue = true; + processGatewayQueue(); + } + return; + } + + // If its closed add back to queue to redo on resume + if (shard?.socket.isClosed) { + requestGuildMembers(guildID, shardID, nonce, options); + return; + } + + shard?.socket.send(JSON.stringify({ + op: GatewayOpcode.RequestGuildMembers, + d: { + guild_id: guildID, + query: options?.query || "", + limit: options?.query || 0, + presences: options?.presences || false, + user_ids: options?.userIDs, + nonce, + }, + })); +} + +async function processGatewayQueue() { + if (!RequestMembersQueue.length) { + processQueue = false; + return; + } + + basicShards.forEach((shard) => { + const index = RequestMembersQueue.findIndex((q) => q.shardID === shard.id); + // 2 events per second is the rate limit. + const request = RequestMembersQueue[index]; + if (request) { + requestGuildMembers( + request.guildID, + request.shardID, + request.nonce, + request.options, + true, + ); + // Remove item from queue + RequestMembersQueue.splice(index, 1); + + const secondIndex = RequestMembersQueue.findIndex((q) => + q.shardID === shard.id + ); + const secondRequest = RequestMembersQueue[secondIndex]; + if (secondRequest) { + requestGuildMembers( + request.guildID, + request.shardID, + secondRequest.nonce, + secondRequest.options, + true, + ); + // Remove item from queue + RequestMembersQueue.splice(secondIndex, 1); + } + } + }); + + await delay(1500); + + eventHandlers.debug?.( + { + type: "requestMembersProcessing", + data: { + remaining: RequestMembersQueue.length, + first: RequestMembersQueue[0], + }, + }, + ); + processGatewayQueue(); +} + +export function botGatewayStatusRequest(payload: BotStatusRequest) { + basicShards.forEach((shard) => { + shard.socket.send(JSON.stringify({ + op: GatewayOpcode.StatusUpdate, + d: { + since: null, + game: payload.game.name + ? { + name: payload.game.name, + type: payload.game.type, + } + : null, + status: payload.status, + afk: false, + }, + })); + }); +} diff --git a/module/client.ts b/module/client.ts index 46accd5a7..4dfeb4c5c 100644 --- a/module/client.ts +++ b/module/client.ts @@ -49,7 +49,7 @@ export const createClient = async (data: ClientOptions) => { (bits, next) => (bits |= next), 0, ); - identifyPayload.shard = [0, botGatewayData.shards] + identifyPayload.shard = [0, botGatewayData.shards]; spawnShards(botGatewayData, identifyPayload); }; diff --git a/module/shardingManager.ts b/module/shardingManager.ts index 020ce4b76..85946b71e 100644 --- a/module/shardingManager.ts +++ b/module/shardingManager.ts @@ -55,9 +55,15 @@ import { } from "../types/message.ts"; import { createMessage } from "../structures/message.ts"; import { GuildUpdateChange } from "../types/options.ts"; -import { createBasicShard } from "./basicShard.ts"; +import { + createBasicShard, + requestGuildMembers, + botGatewayStatusRequest, +} from "./basicShard.ts"; +import { BotStatusRequest } from "../utils/utils.ts"; let shardCounter = 0; +let basicSharding = false; export interface FetchAllMembersRequest { resolve: Function; @@ -114,6 +120,7 @@ export const spawnShards = async ( createNextShard = false; if (data.shards >= 25) createShardWorker(); else { + basicSharding = true; createBasicShard(data, payload, false, id - 1); } spawnShards(data, payload, id + 1); @@ -136,9 +143,11 @@ export async function handleDiscordPayload( return eventHandlers.heartbeat?.(); case GatewayOpcode.Dispatch: if (data.t === "READY") { - setBotID((data.d as ReadyPayload).user.id); + const payload = data.d as ReadyPayload; + setBotID(payload.user.id); // Triggered on each shard - eventHandlers.ready?.(); + eventHandlers.shardReady?.(shardID); + if (payload.shard && shardID === payload.shard[1] - 1) eventHandlers.ready?.() // Wait 5 seconds to spawn next shard await delay(5000); createNextShard = true; @@ -643,6 +652,10 @@ export async function requestAllMembers( const nonce = Math.random().toString(); fetchAllMembersProcessingRequests.set(nonce, resolve); + if (basicSharding) { + return requestGuildMembers(guild.id, guild.shardID, nonce, options); + } + shards[guild.shardID].postMessage({ type: "FETCH_MEMBERS", guildID: guild.id, @@ -652,6 +665,13 @@ export async function requestAllMembers( } export function sendGatewayCommand(type: "EDIT_BOTS_STATUS", payload: object) { + if (basicSharding) { + if (type === "EDIT_BOTS_STATUS") { + botGatewayStatusRequest(payload as BotStatusRequest); + } + + return; + } shards.forEach((shard) => { shard.postMessage({ type, diff --git a/types/options.ts b/types/options.ts index 20c7a75c1..86fdbf233 100644 --- a/types/options.ts +++ b/types/options.ts @@ -125,6 +125,7 @@ export interface EventHandlers { roleUpdate?: (guild: Guild, role: Role, cachedRole: Role) => unknown; roleGained?: (guild: Guild, member: Member, roleID: string) => unknown; roleLost?: (guild: Guild, member: Member, roleID: string) => unknown; + shardReady?: (shardID: number) => unknown; typingStart?: (data: TypingStartPayload) => unknown; voiceChannelJoin?: (member: Member, channelID: string) => unknown; voiceChannelLeave?: (member: Member, channelID: string) => unknown; diff --git a/utils/utils.ts b/utils/utils.ts index e9c19cbdf..8b2cf2ce8 100644 --- a/utils/utils.ts +++ b/utils/utils.ts @@ -7,6 +7,14 @@ export const sleep = (timeout: number) => { return new Promise((resolve) => setTimeout(resolve, timeout)); }; +export interface BotStatusRequest { + status: StatusType; + game: { + name?: string; + type: ActivityType; + }; +} + export function editBotsStatus( status: StatusType, name?: string,