feat: setup postgres replicas (#421)

* chore: setup postgres replicas * chore: properly setup rpc querying * fix: withReplicas slaves --------- Co-authored-by: hzmi <[email protected]>
NezuChan · May 1, 2024 · 3e59893 · 3e59893
1 parent 02f68b0
commit 3e59893
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 9 deletions.
diff --git 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<void> {
         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
@@ -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";