Skip to content
This repository has been archived by the owner on Jul 15, 2024. It is now read-only.

Commit

Permalink
feat: cache optimization and guildCreate throttling (#372)
Browse files Browse the repository at this point in the history
  • Loading branch information
Hazmi35 authored Feb 22, 2024
1 parent f97a73c commit 80a7f37
Show file tree
Hide file tree
Showing 7 changed files with 51 additions and 30 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ export class ChannelCreateListener extends Listener {

public async run(payload: { data: GatewayChannelCreateDispatch; shardId: number; }): Promise<void> {
if (stateChannels) {
const channel = await this.container.client.drizzle.insert(channels).values({
await this.container.client.drizzle.insert(channels).values({
id: payload.data.d.id,
guildId: "guild_id" in payload.data.d ? payload.data.d.guild_id : null,
name: payload.data.d.name,
Expand All @@ -42,15 +42,19 @@ export class ChannelCreateListener extends Listener {
.then(c => c[0]);

if ("permission_overwrites" in payload.data.d && payload.data.d.permission_overwrites !== undefined && payload.data.d.permission_overwrites.length > 0) {
await this.container.client.drizzle.insert(channelsOverwrite).values(payload.data.d.permission_overwrites.map(overwrite => ({
userOrRole: overwrite.id,
channelId: channel.id,
type: overwrite.type,
allow: overwrite.allow,
deny: overwrite.deny
}))).onConflictDoNothing({
target: [channelsOverwrite.userOrRole, channelsOverwrite.channelId]
});
for (const overwrite of payload.data.d.permission_overwrites) {
// @ts-expect-error Intended to avoid .map
overwrite.channelId = payload.data.d.id;

// @ts-expect-error Intended to avoid .map
overwrite.userOrRole = overwrite.id;
}

await this.container.client.drizzle.insert(channelsOverwrite)
.values(payload.data.d.permission_overwrites)
.onConflictDoNothing({
target: [channelsOverwrite.userOrRole, channelsOverwrite.channelId]
});
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ export class ChannelUpdateListener extends Listener {
}

public async run(payload: { data: GatewayChannelUpdateDispatch; shardId: number; }): Promise<void> {
const channel = await this.container.client.drizzle.insert(channels).values({
await this.container.client.drizzle.insert(channels).values({
id: payload.data.d.id,
guildId: "guild_id" in payload.data.d ? payload.data.d.guild_id : null,
name: payload.data.d.name,
Expand All @@ -41,18 +41,23 @@ export class ChannelUpdateListener extends Listener {
.returning({ id: channels.id })
.then(c => c[0]);

// TODO [2024-03-01]: Avoid delete all, intelligently delete only the ones that are not in the new payload
await this.container.client.drizzle.delete(channelsOverwrite).where(eq(channelsOverwrite.channelId, payload.data.d.id));

if ("permission_overwrites" in payload.data.d && payload.data.d.permission_overwrites !== undefined && payload.data.d.permission_overwrites.length > 0) {
await this.container.client.drizzle.insert(channelsOverwrite).values(payload.data.d.permission_overwrites.map(overwrite => ({
userOrRole: overwrite.id,
channelId: channel.id,
type: overwrite.type,
allow: overwrite.allow,
deny: overwrite.deny
}))).onConflictDoNothing({
target: [channelsOverwrite.userOrRole, channelsOverwrite.channelId]
});
for (const overwrite of payload.data.d.permission_overwrites) {
// @ts-expect-error Intended to avoid .map
overwrite.channelId = payload.data.d.id;

// @ts-expect-error Intended to avoid .map
overwrite.userOrRole = overwrite.id;
}

await this.container.client.drizzle.insert(channelsOverwrite)
.values(payload.data.d.permission_overwrites)
.onConflictDoNothing({
target: [channelsOverwrite.userOrRole, channelsOverwrite.channelId]
});
}

await this.container.client.amqp.publish(RabbitMQ.GATEWAY_QUEUE_SEND, RoutingKey(clientId, payload.shardId), Buffer.from(JSON.stringify(payload.data)));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ export class GuildMemberUpdateListener extends Listener {
}
});

// TODO [2024-03-01]: Avoid delete all, intelligently delete only the ones that are not in the new payload
await this.container.client.drizzle.delete(memberRoles).where(and(eq(memberRoles.memberId, payload.data.d.user.id), eq(memberRoles.guildId, payload.data.d.guild_id)));

if (payload.data.d.roles.length > 0) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ export class GuildMembersChunkListener extends Listener {
}

public async run(payload: { data: GatewayGuildMembersChunkDispatch; shardId: number; }): Promise<void> {
// TODO [2024-03-01]: Avoid .map for large arrays (unless if it's just string[])
const chunks = chunk(payload.data.d.members, 1_000);

for (const memberChunk of chunks) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,6 @@ import { Listener } from "../../../Stores/Listener.js";
import { clientId, guildCreateGcEvery, stateChannels, stateRoles } from "../../../config.js";

export class GuildCreateListener extends Listener {
public count = 0;
public gcEvery = guildCreateGcEvery;
public constructor(context: ListenerContext) {
super(context, {
event: GatewayDispatchEvents.GuildCreate
Expand Down Expand Up @@ -304,11 +302,10 @@ export class GuildCreateListener extends Listener {
Buffer.from(JSON.stringify(payload.data))
);

this.count++;

if (global.gc && this.count % this.gcEvery === 0) {
this.logger.info(`Running garbage collection in ${payload.shardId}, ${this.count} Guilds flushed to the database so far`);
if (global.gc && this.container.client.guildsCreateThrottle % guildCreateGcEvery === 0) {
this.logger.info(`Running garbage collection in ${payload.shardId}, ${this.container.client.guildsCreateThrottle} Guilds flushed to the database so far`);
global.gc();
this.container.client.guildsCreateThrottle = 0;
}
}
}
18 changes: 15 additions & 3 deletions services/kanao-cache/src/Structures/KanaoCache.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,13 @@ import * as schema from "@nezuchan/kanao-schema";
import { createAmqpChannel } from "@nezuchan/utilities";
import { StoreRegistry, container } from "@sapphire/pieces";
import type { Channel } from "amqplib";
import type { GatewayDispatchPayload } from "discord-api-types/v10.js";
import { GatewayDispatchEvents } from "discord-api-types/v10.js";
import { drizzle } from "drizzle-orm/node-postgres";
import pg from "pg";
import { ListenerStore } from "../Stores/ListenerStore.js";
import { createLogger } from "../Utilities/Logger.js";
import { clientId, storeLogs, lokiHost, databaseUrl, amqp, databaseConnectionLimit } from "../config.js";
import { clientId, storeLogs, lokiHost, databaseUrl, amqp, databaseConnectionLimit, guildCreateGcEvery, disableGuildCreateGcThrottle } from "../config.js";

export class KanaoCache extends EventEmitter {
public amqp = createAmqpChannel(amqp, {
Expand All @@ -22,6 +24,8 @@ export class KanaoCache extends EventEmitter {

public stores = new StoreRegistry();

public guildsCreateThrottle = 0;

public async connect(): Promise<void> {
container.client = this;
await this.pgClient.connect();
Expand All @@ -43,8 +47,16 @@ export class KanaoCache extends EventEmitter {

await channel.consume(queue, message => {
if (message && message.properties.replyTo === clientId) {
channel.ack(message);
this.emit("dispatch", JSON.parse(message.content.toString()));
const payload = JSON.parse(message.content.toString()) as { shardId: number; data: { data: GatewayDispatchPayload; }; };

if (payload.data.data.t === GatewayDispatchEvents.GuildCreate && !disableGuildCreateGcThrottle) {
if (this.guildsCreateThrottle++ % guildCreateGcEvery === 0) return;
channel.ack(message);
this.emit("dispatch", payload);
} else {
channel.ack(message);
this.emit("dispatch", payload);
}
}
});
}
Expand Down
3 changes: 2 additions & 1 deletion services/kanao-cache/src/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,4 +19,5 @@ export const stateRoles = process.env.STATE_ROLE === "true";
export const stateChannels = process.env.STATE_CHANNEL === "true";
export const stateMessages = process.env.STATE_MESSAGE === "true";

export const guildCreateGcEvery = Number(process.env.GUILD_CREATE_GC_EVERY ?? 50);
export const guildCreateGcEvery = Number(process.env.GUILD_CREATE_GC_EVERY ?? 150);
export const disableGuildCreateGcThrottle = process.env.DISABLE_GUILD_CREATE_GC_THROTTLE === "true";

0 comments on commit 80a7f37

Please sign in to comment.