Attempt to reduce OOM chance by factoring out member->guild queries

This commit is contained in:
Rory& 2025-09-09 04:10:54 +02:00
parent 72abd86a9e
commit 2cb838d0e3

View File

@ -16,15 +16,7 @@
along with this program. If not, see <https://www.gnu.org/licenses/>. along with this program. If not, see <https://www.gnu.org/licenses/>.
*/ */
import { import { CLOSECODES, Capabilities, OPCODES, Payload, Send, WebSocket, setupListener } from "@spacebar/gateway";
CLOSECODES,
Capabilities,
OPCODES,
Payload,
Send,
WebSocket,
setupListener,
} from "@spacebar/gateway";
import { import {
Application, Application,
Config, Config,
@ -104,8 +96,7 @@ export async function onIdentify(this: WebSocket, data: Payload) {
const userQueryTime = taskSw.getElapsedAndReset(); const userQueryTime = taskSw.getElapsedAndReset();
// Check intents // Check intents
if (!identify.intents) if (!identify.intents) identify.intents = 0b11011111111111111111111111111111111n; // TODO: what is this number?
identify.intents = 0b11011111111111111111111111111111111n; // TODO: what is this number?
this.intents = new Intents(identify.intents); this.intents = new Intents(identify.intents);
// TODO: actually do intent things. // TODO: actually do intent things.
@ -115,17 +106,9 @@ export async function onIdentify(this: WebSocket, data: Payload) {
this.shard_id = identify.shard[0]; this.shard_id = identify.shard[0];
this.shard_count = identify.shard[1]; this.shard_count = identify.shard[1];
if ( if (this.shard_count == null || this.shard_id == null || this.shard_id > this.shard_count || this.shard_id < 0 || this.shard_count <= 0) {
this.shard_count == null ||
this.shard_id == null ||
this.shard_id > this.shard_count ||
this.shard_id < 0 ||
this.shard_count <= 0
) {
// TODO: why do we even care about this right now? // TODO: why do we even care about this right now?
console.log( console.log(`[Gateway] Invalid sharding from ${user.id}: ${identify.shard}`);
`[Gateway] Invalid sharding from ${user.id}: ${identify.shard}`,
);
return this.close(CLOSECODES.Invalid_shard); return this.close(CLOSECODES.Invalid_shard);
} }
} }
@ -170,13 +153,7 @@ export async function onIdentify(this: WebSocket, data: Payload) {
timePromise(() => timePromise(() =>
ReadState.find({ ReadState.find({
where: { user_id: this.user_id }, where: { user_id: this.user_id },
select: [ select: ["id", "channel_id", "last_message_id", "last_pin_timestamp", "mention_count"],
"id",
"channel_id",
"last_message_id",
"last_pin_timestamp",
"mention_count",
],
}), }),
), ),
@ -185,29 +162,28 @@ export async function onIdentify(this: WebSocket, data: Payload) {
where: { id: this.user_id }, where: { id: this.user_id },
select: { select: {
// We only want some member props // We only want some member props
...Object.fromEntries( ...Object.fromEntries(MemberPrivateProjection.map((x) => [x, true])),
MemberPrivateProjection.map((x) => [x, true]),
),
settings: true, // guild settings settings: true, // guild settings
roles: { id: true }, // the full role is fetched from the `guild` relation roles: { id: true }, // the full role is fetched from the `guild` relation
guild: { id: true },
// TODO: we don't really need every property of // TODO: we don't really need every property of
// guild channels, emoji, roles, stickers // guild channels, emoji, roles, stickers
// but we do want almost everything from guild. // but we do want almost everything from guild.
// How do you do that without just enumerating the guild props? // How do you do that without just enumerating the guild props?
guild: Object.fromEntries( // guild: Object.fromEntries(
getDatabase()! // getDatabase()!
.getMetadata(Guild) // .getMetadata(Guild)
.columns.map((x) => [x.propertyName, true]), // .columns.map((x) => [x.propertyName, true]),
), // ),
}, },
relations: [ relations: [
"guild", // "guild",
"guild.channels", // "guild.channels",
"guild.emojis", // "guild.emojis",
"guild.roles", // "guild.roles",
"guild.stickers", // "guild.stickers",
"guild.voice_states", // "guild.voice_states",
"roles", "roles",
// For these entities, `user` is always just the logged in user we fetched above // For these entities, `user` is always just the logged in user we fetched above
@ -219,11 +195,7 @@ export async function onIdentify(this: WebSocket, data: Payload) {
timePromise(() => timePromise(() =>
Recipient.find({ Recipient.find({
where: { user_id: this.user_id, closed: false }, where: { user_id: this.user_id, closed: false },
relations: [ relations: ["channel", "channel.recipients", "channel.recipients.user"],
"channel",
"channel.recipients",
"channel.recipients.user",
],
select: { select: {
channel: { channel: {
id: true, id: true,
@ -241,9 +213,7 @@ export async function onIdentify(this: WebSocket, data: Payload) {
// at least one column. // at least one column.
id: true, id: true,
// We only want public user data for each dm channel // We only want public user data for each dm channel
user: Object.fromEntries( user: Object.fromEntries(PublicUserProjection.map((x) => [x, true])),
PublicUserProjection.map((x) => [x, true]),
),
}, },
}, },
}, },
@ -251,6 +221,27 @@ export async function onIdentify(this: WebSocket, data: Payload) {
), ),
]); ]);
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]),
),
relations: ["channels", "emojis", "roles", "stickers", "voice_states"],
}),
),
),
);
members.forEach((m) => {
const g = memberGuilds.find((mg) => mg.id === m.guild_id);
if (g) m.guild = g;
});
const totalQueryTime = taskSw.getElapsedAndReset(); const totalQueryTime = taskSw.getElapsedAndReset();
// We forgot to migrate user settings from the JSON column of `users` // We forgot to migrate user settings from the JSON column of `users`
@ -311,9 +302,7 @@ export async function onIdentify(this: WebSocket, data: Payload) {
}) })
*/ */
.map((channel) => { .map((channel) => {
channel.position = member.guild.channel_ordering.indexOf( channel.position = member.guild.channel_ordering.indexOf(channel.id);
channel.id,
);
return channel; return channel;
}) })
.sort((a, b) => a.position - b.position); .sort((a, b) => a.position - b.position);
@ -333,21 +322,18 @@ export async function onIdentify(this: WebSocket, data: Payload) {
const generateGuildsListTime = taskSw.getElapsedAndReset(); const generateGuildsListTime = taskSw.getElapsedAndReset();
// Generate user_guild_settings // Generate user_guild_settings
const user_guild_settings_entries: ReadyUserGuildSettingsEntries[] = const user_guild_settings_entries: ReadyUserGuildSettingsEntries[] = members.map((x) => ({
members.map((x) => ({ ...DefaultUserGuildSettings,
...DefaultUserGuildSettings, ...x.settings,
...x.settings, guild_id: x.guild_id,
guild_id: x.guild_id, channel_overrides: Object.entries(x.settings.channel_overrides ?? {}).map((y) => ({
channel_overrides: Object.entries( ...y[1],
x.settings.channel_overrides ?? {}, channel_id: y[0],
).map((y) => ({ })),
...y[1], }));
channel_id: y[0],
})),
}));
const generateUserGuildSettingsTime = taskSw.getElapsedAndReset(); const generateUserGuildSettingsTime = taskSw.getElapsedAndReset();
// Popultaed with users from private channels, relationships. // Populated with users from private channels, relationships.
// Uses a set to dedupe for us. // Uses a set to dedupe for us.
const users: Set<PublicUser> = new Set(); const users: Set<PublicUser> = new Set();
@ -360,16 +346,11 @@ export async function onIdentify(this: WebSocket, data: Payload) {
const channel = r.channel as DMChannel; const channel = r.channel as DMChannel;
// Remove ourself from the list of other users in dm channel // Remove ourself from the list of other users in dm channel
channel.recipients = channel.recipients.filter( channel.recipients = channel.recipients.filter((recipient) => recipient.user.id !== this.user_id);
(recipient) => recipient.user.id !== this.user_id,
);
const channelUsers = channel.recipients?.map((recipient) => const channelUsers = channel.recipients?.map((recipient) => recipient.user.toPublicUser());
recipient.user.toPublicUser(),
);
if (channelUsers && channelUsers.length > 0) if (channelUsers && channelUsers.length > 0) channelUsers.forEach((user) => users.add(user));
channelUsers.forEach((user) => users.add(user));
return { return {
id: channel.id, id: channel.id,
@ -403,10 +384,7 @@ export async function onIdentify(this: WebSocket, data: Payload) {
})); }));
const findAndGenerateSessionReplaceTime = taskSw.getElapsedAndReset(); const findAndGenerateSessionReplaceTime = taskSw.getElapsedAndReset();
const [ const [{ elapsed: emitSessionsReplaceTime }, { elapsed: emitPresenceUpdateTime }] = await Promise.all([
{ elapsed: emitSessionsReplaceTime },
{ elapsed: emitPresenceUpdateTime },
] = await Promise.all([
timePromise(() => timePromise(() =>
emitEvent({ emitEvent({
event: "SESSIONS_REPLACE", event: "SESSIONS_REPLACE",
@ -437,14 +415,10 @@ export async function onIdentify(this: WebSocket, data: Payload) {
const d: ReadyEventData = { const d: ReadyEventData = {
v: 9, v: 9,
application: application application: application ? { id: application.id, flags: application.flags } : undefined,
? { id: application.id, flags: application.flags }
: undefined,
user: user.toPrivateUser(), user: user.toPrivateUser(),
user_settings: user.settings, user_settings: user.settings,
guilds: this.capabilities.has(Capabilities.FLAGS.CLIENT_STATE_V2) guilds: this.capabilities.has(Capabilities.FLAGS.CLIENT_STATE_V2) ? guilds.map((x) => new ReadyGuildDTO(x).toJSON()) : guilds,
? guilds.map((x) => new ReadyGuildDTO(x).toJSON())
: guilds,
relationships: user.relationships.map((x) => x.toPublicRelationship()), relationships: user.relationships.map((x) => x.toPublicRelationship()),
read_state: { read_state: {
entries: read_states, entries: read_states,
@ -463,16 +437,10 @@ export async function onIdentify(this: WebSocket, data: Payload) {
merged_members: merged_members, merged_members: merged_members,
sessions: allSessions, sessions: allSessions,
resume_gateway_url: resume_gateway_url: Config.get().gateway.endpointClient || Config.get().gateway.endpointPublic || "ws://127.0.0.1:3001",
Config.get().gateway.endpointClient ||
Config.get().gateway.endpointPublic ||
"ws://127.0.0.1:3001",
// lol hack whatever // lol hack whatever
required_action: required_action: Config.get().login.requireVerification && !user.verified ? "REQUIRE_VERIFIED_EMAIL" : undefined,
Config.get().login.requireVerification && !user.verified
? "REQUIRE_VERIFIED_EMAIL"
: undefined,
consents: { consents: {
personalization: { personalization: {
@ -534,6 +502,7 @@ export async function onIdentify(this: WebSocket, data: Payload) {
read_statesQueryTime, read_statesQueryTime,
membersQueryTime, membersQueryTime,
recipientsQueryTime, recipientsQueryTime,
queryGuildsTime,
})) { })) {
if (subvalue) { if (subvalue) {
val.calls.push(subkey, { val.calls.push(subkey, {
@ -564,9 +533,7 @@ export async function onIdentify(this: WebSocket, data: Payload) {
await Promise.all( await Promise.all(
pending_guilds.map((x) => { pending_guilds.map((x) => {
//Even with the GUILD_MEMBERS intent, the bot always receives just itself as the guild members //Even with the GUILD_MEMBERS intent, the bot always receives just itself as the guild members
const botMemberObject = members.find( const botMemberObject = members.find((member) => member.guild_id === x.id);
(member) => member.guild_id === x.id,
);
return Send(this, { return Send(this, {
op: OPCODES.Dispatch, op: OPCODES.Dispatch,
@ -583,21 +550,15 @@ export async function onIdentify(this: WebSocket, data: Payload) {
] ]
: [], : [],
}, },
})?.catch((e) => })?.catch((e) => console.error(`[Gateway] error when sending bot guilds`, e));
console.error(`[Gateway] error when sending bot guilds`, e),
);
}), }),
); );
const pendingGuildsTime = Date.now(); const pendingGuildsTime = Date.now();
const readySupplementalGuilds = ( const readySupplementalGuilds = (guilds.filter((guild) => !guild.unavailable) as Guild[]).map((guild) => {
guilds.filter((guild) => !guild.unavailable) as Guild[]
).map((guild) => {
return { return {
voice_states: guild.voice_states.map((state) => voice_states: guild.voice_states.map((state) => state.toPublicVoiceState()),
state.toPublicVoiceState(),
),
id: guild.id, id: guild.id,
embedded_activities: [], embedded_activities: [],
}; };
@ -631,8 +592,5 @@ export async function onIdentify(this: WebSocket, data: Payload) {
const setupListenerTime = Date.now(); const setupListenerTime = Date.now();
console.log( console.log(`[Gateway] IDENTIFY ${this.user_id} in ${totalSw.elapsed().totalMilliseconds}ms`, JSON.stringify(d._trace, null, 2));
`[Gateway] IDENTIFY ${this.user_id} in ${totalSw.elapsed().totalMilliseconds}ms`,
JSON.stringify(d._trace, null, 2),
);
} }