Refactor Identify : batch guild queries and parallelise relation fetches

This commit is contained in:
Oh64
2026-03-27 01:44:49 +01:00
committed by Rory&
parent b11ac9b4b4
commit d32c75d164
3 changed files with 92 additions and 56 deletions
+89 -55
View File
@@ -291,35 +291,33 @@ export async function onIdentify(this: WebSocket, data: Payload) {
const userMetaQueryTime = taskSw.getElapsedAndReset();
const { result: memberGuilds, elapsed: queryGuildsTime } = await timePromise(() =>
Promise.all(
members.map((m) =>
Guild.findOneOrFail({
where: { id: m.guild_id },
select: Object.fromEntries(
getDatabase()!
.getMetadata(Guild)
.columns.map((x) => [x.propertyName, true]),
),
}),
),
),
);
const guildIds = memberGuilds.map((g) => g.id);
const memberGuildIds = members.map((m) => m.guild_id);
// select relations
const [
{ result: memberGuilds, elapsed: queryGuildsTime },
{ result: memberGuildChannels, elapsed: queryGuildChannelsTime },
{ result: memberGuildEmojis, elapsed: queryGuildEmojisTime },
{ result: memberGuildRoles, elapsed: queryGuildRolesTime },
{ result: memberGuildStickers, elapsed: queryGuildStickersTime },
{ result: memberGuildVoiceStates, elapsed: queryGuildVoiceStatesTime },
{ result: threadMembers, elapsed: threadMemberTime },
{ result: allThreadsRaw, elapsed: queryThreadsTime },
] = await Promise.all([
timePromise(() =>
Guild.find({
where: { id: In(memberGuildIds) },
select: Object.fromEntries(
getDatabase()!
.getMetadata(Guild)
.columns.map((x) => [x.propertyName, true]),
),
}),
),
timePromise(() =>
Channel.find({
where: {
guild_id: In(guildIds),
guild_id: In(memberGuildIds),
type: Not(In([ChannelType.GUILD_PUBLIC_THREAD, ChannelType.GUILD_PRIVATE_THREAD, ChannelType.GUILD_NEWS_THREAD])),
},
order: { guild_id: "ASC" },
@@ -328,34 +326,93 @@ export async function onIdentify(this: WebSocket, data: Payload) {
),
timePromise(() =>
Emoji.find({
where: { guild_id: In(guildIds) },
where: { guild_id: In(memberGuildIds) },
order: { guild_id: "ASC" },
}),
),
timePromise(() =>
Role.find({
where: { guild_id: In(guildIds) },
where: { guild_id: In(memberGuildIds) },
order: { guild_id: "ASC" },
}),
),
timePromise(() =>
Sticker.find({
where: { guild_id: In(guildIds) },
where: { guild_id: In(memberGuildIds) },
order: { guild_id: "ASC" },
}),
),
timePromise(() =>
VoiceState.find({
where: { guild_id: In(guildIds) },
where: { guild_id: In(memberGuildIds) },
order: { guild_id: "ASC" },
}),
),
timePromise(() =>
ThreadMember.find({
where: { member_idx: In(members.map(({ index }) => index)) },
}),
),
timePromise(() =>
Channel.find({
where: {
type: In([ChannelType.GUILD_NEWS_THREAD, ChannelType.GUILD_PUBLIC_THREAD]),
guild_id: In(memberGuildIds),
},
}),
),
]);
const guildIds = memberGuilds.map((g) => g.id);
const allThreads = allThreadsRaw.filter(({ thread_metadata }) => thread_metadata?.archived === false);
const threadMemberMap = new Map(threadMembers.map((member) => [member.id, member] as const));
const channelsByGuild = new Map<string, Channel[]>();
const emojisByGuild = new Map<string, Emoji[]>();
const rolesByGuild = new Map<string, Role[]>();
const stickersByGuild = new Map<string, Sticker[]>();
const voiceStatesByGuild = new Map<string, VoiceState[]>();
const threadsByGuild = new Map<string, Channel[]>();
for (const c of memberGuildChannels) {
const arr = channelsByGuild.get(c.guild_id!);
if (arr) arr.push(c);
else channelsByGuild.set(c.guild_id!, [c]);
}
for (const e of memberGuildEmojis) {
const arr = emojisByGuild.get(e.guild_id!);
if (arr) arr.push(e);
else emojisByGuild.set(e.guild_id!, [e]);
}
for (const r of memberGuildRoles) {
const arr = rolesByGuild.get(r.guild_id!);
if (arr) arr.push(r);
else rolesByGuild.set(r.guild_id!, [r]);
}
for (const s of memberGuildStickers) {
const arr = stickersByGuild.get(s.guild_id!);
if (arr) arr.push(s);
else stickersByGuild.set(s.guild_id!, [s]);
}
for (const v of memberGuildVoiceStates) {
const arr = voiceStatesByGuild.get(v.guild_id!);
if (arr) arr.push(v);
else voiceStatesByGuild.set(v.guild_id!, [v]);
}
for (const t of allThreads) {
const arr = threadsByGuild.get(t.guild_id!);
if (arr) arr.push(t);
else threadsByGuild.set(t.guild_id!, [t]);
}
const guildMap = new Map(memberGuilds.map((g) => [g.id, g]));
const mergeMemberGuildsTrace: TraceNode = {
micros: 0,
calls: [],
};
members.forEach((m) => {
const sw = Stopwatch.startNew();
const totalSw = Stopwatch.startNew();
@@ -364,32 +421,26 @@ export async function onIdentify(this: WebSocket, data: Payload) {
calls: [],
};
const g = memberGuilds.find((mg) => mg.id === m.guild_id);
const g = guildMap.get(m.guild_id);
if (g) {
m.guild = g;
trace.calls.push("findGuild", { micros: sw.getElapsedAndReset().totalMicroseconds });
//channels
g.channels = memberGuildChannels.filter((c) => c.guild_id === m.guild_id);
g.channels = channelsByGuild.get(m.guild_id) ?? [];
trace.calls.push(`filterChannels(${g.channels.length}/${memberGuildChannels.length})`, { micros: sw.getElapsedAndReset().totalMicroseconds });
//emojis
g.emojis = memberGuildEmojis.filter((e) => e.guild_id === m.guild_id);
g.emojis = emojisByGuild.get(m.guild_id) ?? [];
trace.calls.push(`filterEmojis(${g.emojis.length}/${memberGuildEmojis.length})`, { micros: sw.getElapsedAndReset().totalMicroseconds });
//roles
g.roles = memberGuildRoles.filter((r) => r.guild_id === m.guild_id);
g.roles = rolesByGuild.get(m.guild_id) ?? [];
trace.calls.push(`filterRoles(${g.roles.length}/${memberGuildRoles.length})`, { micros: sw.getElapsedAndReset().totalMicroseconds });
//stickers
g.stickers = memberGuildStickers.filter((s) => s.guild_id === m.guild_id);
g.stickers = stickersByGuild.get(m.guild_id) ?? [];
trace.calls.push(`filterStickers(${g.stickers.length}/${memberGuildStickers.length})`, { micros: sw.getElapsedAndReset().totalMicroseconds });
//voice states
g.voice_states = memberGuildVoiceStates.filter((v) => v.guild_id === m.guild_id);
g.voice_states = voiceStatesByGuild.get(m.guild_id) ?? [];
trace.calls.push(`filterVoiceStates(${g.voice_states.length}/${memberGuildVoiceStates.length})`, { micros: sw.getElapsedAndReset().totalMicroseconds });
//total
trace.micros = totalSw.elapsed().totalMicroseconds;
mergeMemberGuildsTrace.calls!.push(`guild_${m.guild_id}`, trace);
} else {
@@ -435,31 +486,15 @@ export async function onIdentify(this: WebSocket, data: Payload) {
];
});
const mergedMembersTime = taskSw.getElapsedAndReset();
const member_idx = members.map(({ index }) => index);
const threadMembers = await ThreadMember.find({
where: { member_idx: In(member_idx) },
});
const threadMemberMap = new Map(threadMembers.map((member) => [member.id, member] as const));
const threadMemberTime = taskSw.getElapsedAndReset();
// Populated with guilds 'unavailable' currently
// Just for bots
//TODO get this a better type
const pending_guilds: { id: string }[] = [];
const allThreads = (
await Channel.find({
where: {
type: In([ChannelType.GUILD_NEWS_THREAD, ChannelType.GUILD_PUBLIC_THREAD]),
guild_id: In(members.map(({ guild }) => guild.id)),
},
})
).filter(({ thread_metadata }) => thread_metadata?.archived === false);
// Generate guilds list ( make them unavailable if user is bot )
const guilds: GuildOrUnavailable[] = members.map((member) => {
member.guild.channels = member.guild.channels
member.guild.channels = (channelsByGuild.get(member.guild_id) ?? [])
/*
//TODO maybe implement this correctly, by causing create and delete events for users who can newly view and not view the channels, along with doing these checks correctly, as they don't currently take into account that the owner of the guild is always able to view channels, with potentially other issues
.filter((channel) => {
@@ -481,7 +516,7 @@ export async function onIdentify(this: WebSocket, data: Payload) {
})
.sort((a, b) => a.position - b.position);
const threads: Channel[] = allThreads.filter((_) => _.guild_id === member.guild_id);
const threads: Channel[] = threadsByGuild.get(member.guild_id) ?? [];
const guildjson = {
...member.guild.toJSON(),
@@ -512,10 +547,7 @@ export async function onIdentify(this: WebSocket, data: Payload) {
...DefaultUserGuildSettings,
...x.settings,
guild_id: x.guild_id,
channel_overrides: Object.entries(x.settings.channel_overrides ?? {}).map((y) => ({
...y[1],
channel_id: y[0],
})),
channel_overrides: x.settings.channel_overrides ? Object.entries(x.settings.channel_overrides).map(([k, v]) => ({ ...v, channel_id: k })) : [],
}));
const generateUserGuildSettingsTime = taskSw.getElapsedAndReset();
@@ -750,6 +782,8 @@ export async function onIdentify(this: WebSocket, data: Payload) {
queryGuildRolesTime,
queryGuildStickersTime,
queryGuildVoiceStatesTime,
threadMemberTime,
queryThreadsTime,
})) {
if (subvalue) {
val.calls.push(subkey, {
+2 -1
View File
@@ -31,7 +31,8 @@ export async function onRequestGuildMembers(this: WebSocket, { d }: Payload) {
if (d.user_ids && !Array.isArray(d.user_ids)) d.user_ids = [d.user_ids];
if (Config.get().offload.gateway.guildMembersUrl !== null) {
return await handleOffloadedGatewayRequest(this, Config.get().offload.gateway.guildMembersUrl!, d);
const guildIds: string[] = Array.isArray(d.guild_id) ? d.guild_id : [d.guild_id];
return await handleOffloadedGatewayRequest(this, Config.get().offload.gateway.guildMembersUrl!, guildIds);
}
check.call(this, RequestGuildMembersSchema, d);
+1
View File
@@ -71,6 +71,7 @@ export async function handleOffloadedGatewayRequest(socket: WebSocket, url: stri
Authorization: `Bearer ${socket.accessToken}`,
// because the session may not have an id in the token!
"X-Session-Id": socket.session_id,
"Content-Type": "application/json",
},
});