From ed9c63fb1158e3c7a9ba73397a1c922865687ead Mon Sep 17 00:00:00 2001 From: MARCROCK22 Date: Sat, 18 Jan 2025 20:30:48 -0400 Subject: [PATCH] feat: shard#requestGuildMember --- src/cache/index.ts | 26 +++++++++- src/client/workerclient.ts | 4 +- src/websocket/discord/shard.ts | 89 ++++++++++++++++++++++++++++++-- src/websocket/discord/sharder.ts | 12 ++--- src/websocket/discord/shared.ts | 3 ++ 5 files changed, 122 insertions(+), 12 deletions(-) diff --git a/src/cache/index.ts b/src/cache/index.ts index 0efe9c9bf..00b412250 100644 --- a/src/cache/index.ts +++ b/src/cache/index.ts @@ -92,7 +92,8 @@ export type CachedEvents = | 'VOICE_STATE_UPDATE' | 'STAGE_INSTANCE_CREATE' | 'STAGE_INSTANCE_UPDATE' - | 'STAGE_INSTANCE_DELETE'; + | 'STAGE_INSTANCE_DELETE' + | 'GUILD_MEMBERS_CHUNK'; export type DisabledCache = { [P in NonGuildBased | GuildBased | GuildRelated | SeyfertBased]?: boolean; @@ -554,6 +555,29 @@ export class Cache { ); } break; + case 'GUILD_MEMBERS_CHUNK': { + const data: Parameters[0] = []; + + if (this.members) { + for (const member of event.d.members) { + data.push( + [CacheFrom.Gateway, 'members', member, member.user.id, event.d.guild_id], + [CacheFrom.Gateway, 'users', member.user, member.user.id], + ); + } + } + + if (this.presences && event.d.presences) { + for (const presence of event.d.presences) { + data.push([CacheFrom.Gateway, 'presences', presence, presence.user.id, event.d.guild_id]); + } + } + + if (data.length) { + await this.bulkSet(data); + } + break; + } case 'GUILD_MEMBER_ADD': case 'GUILD_MEMBER_UPDATE': if (event.d.user) await this.members?.set(CacheFrom.Gateway, event.d.user.id, event.d.guild_id, event.d); diff --git a/src/client/workerclient.ts b/src/client/workerclient.ts index 107cef6b7..1772fa4ac 100644 --- a/src/client/workerclient.ts +++ b/src/client/workerclient.ts @@ -4,7 +4,7 @@ import { WorkerAdapter } from '../cache'; import { type DeepPartial, LogLevels, type MakeRequired, type When, lazyLoadPackage } from '../common'; import { EventHandler } from '../events'; import type { GatewayDispatchPayload, GatewaySendPayload } from '../types'; -import { Shard, type ShardManagerOptions, type WorkerData, properties } from '../websocket'; +import { Shard, type ShardManagerOptions, ShardSocketCloseCodes, type WorkerData, properties } from '../websocket'; import type { WorkerDisconnectedAllShardsResharding, WorkerMessages, @@ -350,7 +350,7 @@ export class WorkerClient extends BaseClient { case 'DISCONNECT_ALL_SHARDS_RESHARDING': { for (const i of this.shards.values()) { - await i.disconnect(); + await i.disconnect(ShardSocketCloseCodes.Resharding); } this.postMessage({ type: 'DISCONNECTED_ALL_SHARDS_RESHARDING', diff --git a/src/websocket/discord/shard.ts b/src/websocket/discord/shard.ts index 324894c78..2ce1568e8 100644 --- a/src/websocket/discord/shard.ts +++ b/src/websocket/discord/shard.ts @@ -1,13 +1,19 @@ import { inflateSync } from 'node:zlib'; import { LogLevels, Logger, type MakeRequired, MergeOptions, hasIntent } from '../../common'; import { + type APIGuildMember, GatewayCloseCodes, GatewayDispatchEvents, type GatewayDispatchPayload, + type GatewayGuildMembersChunkPresence, GatewayOpcodes, type GatewayReceivePayload, type GatewaySendPayload, } from '../../types'; +import type { + GatewayRequestGuildMembersDataWithQuery, + GatewayRequestGuildMembersDataWithUserIds, +} from '../../types/gateway'; import { properties } from '../constants'; import { DynamicBucket } from '../structures'; import { ConnectTimeout } from '../structures/timeout'; @@ -42,6 +48,18 @@ export class Shard { pendingGuilds = new Set(); options: MakeRequired; isReady = false; + private requestGuildMembersChunk = new Map< + string, + { + members: APIGuildMember[]; + presences: GatewayGuildMembersChunkPresence[]; + resolve: (value: { + members: APIGuildMember[]; + presences: GatewayGuildMembersChunkPresence[]; + }) => void; + reject: (reason?: any) => void; + } + >(); constructor( public id: number, @@ -198,14 +216,14 @@ export class Shard { ); } - disconnect() { + disconnect(code = ShardSocketCloseCodes.Shutdown) { this.debugger?.info(`[Shard #${this.id}] Disconnecting`); - this.close(ShardSocketCloseCodes.Shutdown, 'Shard down request'); + this.close(code, 'Shard down request'); } async reconnect() { this.debugger?.info(`[Shard #${this.id}] Reconnecting`); - this.disconnect(); + this.disconnect(ShardSocketCloseCodes.Reconnect); await this.connect(); } @@ -300,6 +318,29 @@ export class Shard { this.options.handlePayload(this.id, packet); } break; + case GatewayDispatchEvents.GuildMembersChunk: + { + if (!packet.d.nonce) { + this.options.handlePayload(this.id, packet); + break; + } + const guildMemberChunk = this.requestGuildMembersChunk.get(packet.d.nonce); + if (!guildMemberChunk) { + this.options.handlePayload(this.id, packet); + break; + } + guildMemberChunk.members.push(...packet.d.members); + if (packet.d.presences) guildMemberChunk.presences.push(...packet.d.presences); + if (packet.d.chunk_index + 1 === packet.d.chunk_count) { + this.requestGuildMembersChunk.delete(packet.d.nonce); + guildMemberChunk.resolve({ + members: guildMemberChunk.members, + presences: guildMemberChunk.presences, + }); + } + this.options.handlePayload(this.id, packet); + } + break; default: this.options.handlePayload(this.id, packet); break; @@ -309,6 +350,48 @@ export class Shard { } } + async requestGuildMember( + options: + | Omit + | Omit, + ) { + const nonce = Date.now().toString() + Math.random().toString(36); + + let resolve: (value: { + members: APIGuildMember[]; + presences: GatewayGuildMembersChunkPresence[]; + }) => void = () => { + // + }; + let reject: (reason?: any) => void = () => { + // + }; + + const promise = new Promise<{ + members: APIGuildMember[]; + presences: GatewayGuildMembersChunkPresence[]; + }>((res, rej) => { + resolve = res; + reject = rej; + }); + this.requestGuildMembersChunk.set(nonce, { + members: [], + presences: [], + reject, + resolve, + }); + + this.send(false, { + op: GatewayOpcodes.RequestGuildMembers, + d: { + ...options, + nonce, + }, + }); + + return promise; + } + protected async handleClosed(close: { code: number; reason: string }) { this.isReady = false; clearInterval(this.heart.nodeInterval); diff --git a/src/websocket/discord/sharder.ts b/src/websocket/discord/sharder.ts index 5735544d6..eb1f7b97a 100644 --- a/src/websocket/discord/sharder.ts +++ b/src/websocket/discord/sharder.ts @@ -19,7 +19,7 @@ import { ShardManagerDefaults } from '../constants'; import { DynamicBucket } from '../structures'; import { ConnectQueue } from '../structures/timeout'; import { Shard } from './shard'; -import type { ShardData, ShardManagerOptions, WorkerData } from './shared'; +import { type ShardData, type ShardManagerOptions, ShardSocketCloseCodes, type WorkerData } from './shared'; let parentPort: import('node:worker_threads').MessagePort; let workerData: WorkerData; @@ -159,7 +159,7 @@ export class ShardManager extends Map { handlePayload = () => { // }; - this.disconnectAll(); + this.disconnectAll(ShardSocketCloseCodes.Resharding); this.clear(); this.options.totalShards = this.options.shardEnd = this.options.info.shards = info.shards; @@ -220,14 +220,14 @@ export class ShardManager extends Map { return this.create(shardId).identify(); } - disconnect(shardId: number) { + disconnect(shardId: number, code?: ShardSocketCloseCodes) { this.debugger?.info(`Shard #${shardId} force disconnect`); - return this.get(shardId)?.disconnect(); + return this.get(shardId)?.disconnect(code); } - disconnectAll() { + disconnectAll(code = ShardSocketCloseCodes.ShutdownAll) { this.debugger?.info('Disconnect all shards'); - this.forEach(shard => shard.disconnect()); + this.forEach(shard => shard.disconnect(code)); } setShardPresence(shardId: number, payload: GatewayUpdatePresence['d']) { diff --git a/src/websocket/discord/shared.ts b/src/websocket/discord/shared.ts index 0a9345e6d..ac3850442 100644 --- a/src/websocket/discord/shared.ts +++ b/src/websocket/discord/shared.ts @@ -133,6 +133,9 @@ export interface ShardOptions extends ShardDetails { export enum ShardSocketCloseCodes { Shutdown = 3000, ZombiedConnection = 3010, + Reconnect = 3020, + Resharding = 3030, + ShutdownAll = 3040, } export interface WorkerData {