From d61f61fead3cb432c539635eb8db3780ccc24400 Mon Sep 17 00:00:00 2001
From: Ske <voltasalt@gmail.com>
Date: Wed, 9 Jun 2021 14:49:12 +0200
Subject: [PATCH] Add basic support for multi-node clustering

---
 Myriad/Gateway/Cluster.cs  | 18 ++++++++++--------
 PluralKit.Bot/BotConfig.cs |  9 +++++++++
 PluralKit.Bot/Init.cs      | 36 ++++++++++++++++++++++++++++++++++--
 3 files changed, 53 insertions(+), 10 deletions(-)

diff --git a/Myriad/Gateway/Cluster.cs b/Myriad/Gateway/Cluster.cs
index 3f6a1134..07695a5e 100644
--- a/Myriad/Gateway/Cluster.cs
+++ b/Myriad/Gateway/Cluster.cs
@@ -32,17 +32,19 @@ namespace Myriad.Gateway
         
         public async Task Start(GatewayInfo.Bot info)
         {
-            var concurrency = GetActualShardConcurrency(info.SessionStartLimit.MaxConcurrency);
-            _ratelimiter = new(_logger, concurrency);
-            
-            await Start(info.Url, info.Shards);
+            await Start(info.Url, 0, info.Shards - 1, info.Shards, info.SessionStartLimit.MaxConcurrency);
         }
 
-        public async Task Start(string url, int shardCount)
+        public async Task Start(string url, int shardMin, int shardMax, int shardTotal, int concurrency)
         {
-            _logger.Information("Starting {ShardCount} shards at {Url}", shardCount, url);
-            for (var i = 0; i < shardCount; i++)
-                CreateAndAddShard(url, new ShardInfo(i, shardCount));
+            concurrency = GetActualShardConcurrency(concurrency);
+            _ratelimiter = new(_logger, concurrency);
+
+            var shardCount = shardMax - shardMin + 1;
+            _logger.Information("Starting {ShardCount} of {ShardTotal} shards (#{ShardMin}-#{ShardMax}) at {Url}",
+                shardCount, shardTotal, shardMin, shardMax, url);
+            for (var i = shardMin; i <= shardMax; i++)
+                CreateAndAddShard(url, new ShardInfo(i, shardTotal));
 
             await StartShards();
         }
diff --git a/PluralKit.Bot/BotConfig.cs b/PluralKit.Bot/BotConfig.cs
index 378e20d1..8215afbb 100644
--- a/PluralKit.Bot/BotConfig.cs
+++ b/PluralKit.Bot/BotConfig.cs
@@ -15,5 +15,14 @@ namespace PluralKit.Bot
         public int? MaxShardConcurrency { get; set; }
         
         public ulong? AdminRole { get; set; }
+        
+        public ClusterSettings? Cluster { get; set; }
+
+        public record ClusterSettings
+        {
+            public string NodeName { get; set; }
+            public int TotalShards { get; set; }
+            public int TotalNodes { get; set; }
+        }
     }
 }
\ No newline at end of file
diff --git a/PluralKit.Bot/Init.cs b/PluralKit.Bot/Init.cs
index f4b7c6f5..7ccd9fe0 100644
--- a/PluralKit.Bot/Init.cs
+++ b/PluralKit.Bot/Init.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -48,8 +49,7 @@ namespace PluralKit.Bot
 
                 // Start the Discord shards themselves (handlers already set up)
                 logger.Information("Connecting to Discord");
-                var info = await services.Resolve<DiscordApiClient>().GetGatewayBot();
-                await services.Resolve<Cluster>().Start(info);
+                await StartCluster(services);
                 logger.Information("Connected! All is good (probably).");
 
                 // Lastly, we just... wait. Everything else is handled in the DiscordClient event loop
@@ -128,5 +128,37 @@ namespace PluralKit.Bot
             builder.RegisterModule<BotModule>();
             return builder.Build();
         }
+
+        private static async Task StartCluster(IComponentContext services)
+        {
+            var info = await services.Resolve<DiscordApiClient>().GetGatewayBot();
+
+            var cluster = services.Resolve<Cluster>();
+            var config = services.Resolve<BotConfig>();
+            
+            if (config.Cluster != null)
+            {
+                // For multi-instance deployments, calculate the "span" of shards this node is responsible for
+                var totalNodes = config.Cluster.TotalNodes;
+                var totalShards = config.Cluster.TotalShards;
+                var nodeIndex = ExtractNodeIndex(config.Cluster.NodeName);
+                
+                // Should evenly distribute shards even with an uneven amount of nodes
+                var shardMin = (int) Math.Round(totalShards * (float) nodeIndex / totalNodes);
+                var shardMax = (int) Math.Round(totalShards * (float) (nodeIndex + 1) / totalNodes) - 1;
+                
+                await cluster.Start(info.Url, shardMin, shardMax, totalShards, info.SessionStartLimit.MaxConcurrency);
+            }
+            else
+            {
+                await cluster.Start(info);
+            }
+        }
+
+        private static int ExtractNodeIndex(string nodeName)
+        {
+            // Node name eg. "pluralkit-3", want to extract the 3. blame k8s :p
+            return int.Parse(nodeName.Split("-").Last());
+        }
     }
 }
\ No newline at end of file