From 13e3289c26b1b0bf5946e88f8bb74ae31ceb304c Mon Sep 17 00:00:00 2001
From: Ske <voltasalt@gmail.com>
Date: Thu, 18 Mar 2021 09:47:58 +0100
Subject: [PATCH] Add config option for max shard concurrency

---
 Myriad/Gateway/Cluster.cs         | 12 ++++++++++--
 Myriad/Gateway/GatewaySettings.cs |  1 +
 PluralKit.Bot/BotConfig.cs        |  2 ++
 PluralKit.Bot/Modules.cs          | 23 ++++++++++++++---------
 4 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/Myriad/Gateway/Cluster.cs b/Myriad/Gateway/Cluster.cs
index cbb0bd51..bc5805fa 100644
--- a/Myriad/Gateway/Cluster.cs
+++ b/Myriad/Gateway/Cluster.cs
@@ -2,7 +2,6 @@
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
-using System.Threading;
 using System.Threading.Tasks;
 
 using Myriad.Types;
@@ -69,9 +68,10 @@ namespace Myriad.Gateway
 
             await StartShards(concurrency);
         }
-
         private async Task StartShards(int concurrency)
         {
+            concurrency = GetActualShardConcurrency(concurrency);
+            
             var lastTime = DateTimeOffset.UtcNow;
             var identifyCalls = 0;
             
@@ -112,5 +112,13 @@ namespace Myriad.Gateway
             if (EventReceived != null)
                 await EventReceived(shard, evt);
         }
+        
+        private int GetActualShardConcurrency(int recommendedConcurrency)
+        {
+            if (_gatewaySettings.MaxShardConcurrency == null)
+                return recommendedConcurrency;
+            
+            return Math.Min(_gatewaySettings.MaxShardConcurrency.Value, recommendedConcurrency);
+        }
     }
 }
\ No newline at end of file
diff --git a/Myriad/Gateway/GatewaySettings.cs b/Myriad/Gateway/GatewaySettings.cs
index fdaf13ea..1fdb12fd 100644
--- a/Myriad/Gateway/GatewaySettings.cs
+++ b/Myriad/Gateway/GatewaySettings.cs
@@ -4,5 +4,6 @@
     {
         public string Token { get; init; }
         public GatewayIntent Intents { get; init; }
+        public int? MaxShardConcurrency { get; init; }
     }
 }
\ No newline at end of file
diff --git a/PluralKit.Bot/BotConfig.cs b/PluralKit.Bot/BotConfig.cs
index 153b965c..095b77cd 100644
--- a/PluralKit.Bot/BotConfig.cs
+++ b/PluralKit.Bot/BotConfig.cs
@@ -11,5 +11,7 @@ namespace PluralKit.Bot
         // and fall back to the separate default array at the use site :)
         // This does bind [] as null (therefore default) instead of an empty array, but I can live w/ that. 
         public string[] Prefixes { get; set; }
+        
+        public int? MaxShardConcurrency { get; set; }
     }
 }
\ No newline at end of file
diff --git a/PluralKit.Bot/Modules.cs b/PluralKit.Bot/Modules.cs
index f3265e2c..cef29570 100644
--- a/PluralKit.Bot/Modules.cs
+++ b/PluralKit.Bot/Modules.cs
@@ -21,16 +21,21 @@ namespace PluralKit.Bot
         protected override void Load(ContainerBuilder builder)
         {
             // Clients
-            builder.Register(c => new GatewaySettings
+            builder.Register(c =>
             {
-                Token = c.Resolve<BotConfig>().Token,
-                Intents = GatewayIntent.Guilds |
-                          GatewayIntent.DirectMessages |
-                          GatewayIntent.DirectMessageReactions |
-                          GatewayIntent.GuildEmojis |
-                          GatewayIntent.GuildMessages |
-                          GatewayIntent.GuildWebhooks |
-                          GatewayIntent.GuildMessageReactions
+                var botConfig = c.Resolve<BotConfig>();
+                return new GatewaySettings
+                {
+                    Token = botConfig.Token,
+                    MaxShardConcurrency = botConfig.MaxShardConcurrency,
+                    Intents = GatewayIntent.Guilds |
+                              GatewayIntent.DirectMessages |
+                              GatewayIntent.DirectMessageReactions |
+                              GatewayIntent.GuildEmojis |
+                              GatewayIntent.GuildMessages |
+                              GatewayIntent.GuildWebhooks |
+                              GatewayIntent.GuildMessageReactions
+                };
             }).AsSelf().SingleInstance();
             builder.RegisterType<Cluster>().AsSelf().SingleInstance();
             builder.Register(c => new Myriad.Rest.DiscordApiClient(c.Resolve<BotConfig>().Token, c.Resolve<ILogger>()))