diff --git a/services/kanao-cache/src/Structures/KanaoCache.ts b/services/kanao-cache/src/Structures/KanaoCache.ts index 1f9ba762..2a9390cb 100644 --- a/services/kanao-cache/src/Structures/KanaoCache.ts +++ b/services/kanao-cache/src/Structures/KanaoCache.ts @@ -7,10 +7,11 @@ import { StoreRegistry, container } from "@sapphire/pieces"; import type { Channel } from "amqplib"; import { count } from "drizzle-orm"; import { drizzle } from "drizzle-orm/node-postgres"; +import { withReplicas } from "drizzle-orm/pg-core"; import pg from "pg"; import { ListenerStore } from "../Stores/ListenerStore.js"; import { createLogger } from "../Utilities/Logger.js"; -import { clientId, storeLogs, lokiHost, databaseUrl, amqp, databaseConnectionLimit, prefetchCount } from "../config.js"; +import { clientId, storeLogs, lokiHost, databaseUrl, amqp, databaseConnectionLimit, prefetchCount, databaseReadUrls } from "../config.js"; export class KanaoCache extends EventEmitter { public cacheQueue = createAmqpChannel(amqp, { @@ -26,16 +27,29 @@ export class KanaoCache extends EventEmitter { }); public logger = createLogger("kanao-cache", clientId, storeLogs, lokiHost); - public pgClient = new pg.Pool({ connectionString: databaseUrl, max: databaseConnectionLimit }); - public drizzle = drizzle(this.pgClient, { schema }); + public postgresInstance = { + master: new pg.Pool({ connectionString: databaseUrl, max: databaseConnectionLimit }), + slaves: databaseReadUrls.map(x => new pg.Pool({ connectionString: x, max: databaseConnectionLimit })) + }; + + public drizzle = withReplicas( + drizzle(this.postgresInstance.master, { schema }), + [this.postgresInstance.slaves.shift()!, this.postgresInstance.slaves] + ); public stores = new StoreRegistry(); public async connect(): Promise { container.client = this; - await this.pgClient.connect(); - this.pgClient.on("error", e => this.logger.error(e, "Postgres emitted error")); + + await this.postgresInstance.master.connect(); + this.postgresInstance.master.on("error", e => this.logger.error(e, "Postgres emitted error")); + + for (const instance of this.postgresInstance.slaves) { + await instance.connect(); + instance.on("error", e => this.logger.error(e, "Postgres slave emitted error")); + } this.stores.register(new ListenerStore()); @@ -81,10 +95,18 @@ export class KanaoCache extends EventEmitter { }; try { - const result = await this.pgClient.query(query); - await this.queryRpcQueue.sendToQueue(message.properties.replyTo as string, Buffer.from( - JSON.stringify({ route: rpc.key, rows: result.rows, message: `Successfully querying with ${result.rowCount} results !` }) - ), { correlationId: message.properties.correlationId as string }); + if (query.text.includes("INSERT") || query.text.includes("UPDATE") || query.text.includes("DELETE")) { + const result = await this.postgresInstance.master.query(query); + await this.queryRpcQueue.sendToQueue(message.properties.replyTo as string, Buffer.from( + JSON.stringify({ route: rpc.key, rows: result.rows, message: `Successfully querying with ${result.rowCount} results !` }) + ), { correlationId: message.properties.correlationId as string }); + } else { + const server = this.postgresInstance.slaves[Math.floor(Math.random() * this.postgresInstance.slaves.length)]; + const result = await server.query(query); + await this.queryRpcQueue.sendToQueue(message.properties.replyTo as string, Buffer.from( + JSON.stringify({ route: rpc.key, rows: result.rows, message: `Successfully querying with ${result.rowCount} results !` }) + ), { correlationId: message.properties.correlationId as string }); + } } catch (error) { await this.queryRpcQueue.sendToQueue(message.properties.replyTo as string, Buffer.from( JSON.stringify({ route: rpc.key, rows: [], message: (error as Error).message }) diff --git a/services/kanao-cache/src/config.ts b/services/kanao-cache/src/config.ts index dfd030b8..973a5d28 100644 --- a/services/kanao-cache/src/config.ts +++ b/services/kanao-cache/src/config.ts @@ -10,6 +10,7 @@ export const amqp = process.env.AMQP_HOST ?? process.env.AMQP_URL ?? "amqp://loc export const clientId = process.env.CLIENT_ID ?? Buffer.from(discordToken.split(".")[0], "base64").toString(); export const production = process.env.NODE_ENV === "production"; export const databaseUrl = process.env.DATABASE_GATEWAY_URL ?? process.env.DATABASE_URL!; +export const databaseReadUrls = JSON.parse(process.env.DATABASE_READ_URLS ?? "[]") as string[]; export const databaseConnectionLimit = Number(process.env.DATABASE_CONNECTION_LIMIT ?? 15); export const stateMembers = process.env.STATE_MEMBER === "true";