From ab2d6e2ef06b843341a8eb18387e5672823c41fe Mon Sep 17 00:00:00 2001
From: Essem <smswessem@gmail.com>
Date: Tue, 4 Oct 2022 12:10:06 -0500
Subject: [PATCH] Implement pm2 shard management

---
 .github/workflows/build.yml |  2 +-
 app.js                      |  1 +
 ecosystem.config.cjs        |  7 ----
 utils/pm2/ext.js            | 81 ++++++++++++++++++++++++++++++++++++-
 4 files changed, 82 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index c36757d..78a7b88 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -27,7 +27,7 @@ jobs:
       run: pnpm install --config.strict-peer-dependencies=false && pnpm run build
 
   win64:
-    runs-on: windows-latest
+    runs-on: windows-2019
     steps:
     - name: Checkout
       uses: actions/checkout@v1
diff --git a/app.js b/app.js
index 3ecb9f0..019b1f2 100644
--- a/app.js
+++ b/app.js
@@ -133,6 +133,7 @@ esmBot ${esmBotVersion} (${process.env.GIT_REV})
     gateway: {
       concurrency: "auto",
       maxShards: "auto",
+      shardIDs: process.env.SHARDS ? JSON.parse(process.env.SHARDS)[process.env.pm_id - 1] : null,
       presence: {
         status: "idle",
         activities: [{
diff --git a/ecosystem.config.cjs b/ecosystem.config.cjs
index 7a6b307..bc19548 100644
--- a/ecosystem.config.cjs
+++ b/ecosystem.config.cjs
@@ -6,12 +6,5 @@ module.exports = {
     exp_backoff_restart_delay: 1000,
     watch: false,
     exec_mode: "fork"
-  }, {
-    name: "esmBot",
-    script: "app.js",
-    autorestart: true,
-    exp_backoff_restart_delay: 1000,
-    watch: false,
-    exec_mode: "cluster"
   }]
 };
\ No newline at end of file
diff --git a/utils/pm2/ext.js b/utils/pm2/ext.js
index 17e3f22..e004304 100644
--- a/utils/pm2/ext.js
+++ b/utils/pm2/ext.js
@@ -10,7 +10,11 @@ import { createServer } from "http";
 import { config } from "dotenv";
 config({ path: resolve(dirname(fileURLToPath(import.meta.url)), "../../.env") });
 
+// oceanic client used for getting shard counts
+import { Client } from "oceanic.js";
+
 import database from "../database.js";
+import { cpus } from "os";
 
 const dbl = process.env.NODE_ENV === "production" && process.env.DBL ? new Api(process.env.DBL) : null;
 
@@ -163,4 +167,79 @@ if (dbl) setInterval(dblPost, 1800000);
 
 setTimeout(updateStats, 10000);
 
-logger.info("Started esmBot management process.");
\ No newline at end of file
+logger.info("Started esmBot management process.");
+
+// from eris-fleet
+function calcShards(shards, procs) {
+  if (procs < 2) return [shards];
+
+  const length = shards.length;
+  const r = [];
+  let i = 0;
+  let size;
+
+  if (length % procs === 0) {
+    size = Math.floor(length / procs);
+    while (i < length) {
+      r.push(shards.slice(i, (i += size)));
+    }
+  } else {
+    while (i < length) {
+      size = Math.ceil((length - i) / procs--);
+      r.push(shards.slice(i, (i += size)));
+    }
+  }
+
+  return r;
+}
+
+(async function init() {
+  logger.main("Getting gateway connection data...");
+  const client = new Client({
+    auth: `Bot ${process.env.TOKEN}`,
+    gateway: {
+      concurrency: "auto",
+      maxShards: "auto",
+      presence: {
+        status: "idle",
+        activities: [{
+          type: 0,
+          name: "Starting esmBot..."
+        }]
+      },
+      intents: []
+    }
+  });
+
+  const connectionData = await client.rest.getBotGateway();
+  const cpuAmount = cpus().length;
+  const procAmount = Math.min(connectionData.shards, cpuAmount);
+  logger.main(`Obtained data, connecting with ${connectionData.shards} shard(s) across ${procAmount} process(es)...`);
+
+  const lastShard = connectionData.shards - 1;
+  const shardArray = [];
+  for (let i = 0; i <= lastShard; i++) {
+    shardArray.push(i);
+  }
+  const shardArrays = calcShards(shardArray, procAmount);
+
+  pm2.start({
+    name: "esmBot",
+    script: "app.js",
+    autorestart: true,
+    exp_backoff_restart_delay: 1000,
+    watch: false,
+    exec_mode: "cluster",
+    instances: shardArrays.length,
+    env: {
+      "SHARDS": JSON.stringify(shardArrays)
+    }
+  }, (err) => {
+    if (err) {
+      logger.error(`Failed to start esmBot: ${err}`);
+      process.exit(0);
+    } else {
+      logger.info("Started esmBot processes.");
+    }
+  });
+})();
\ No newline at end of file