feat: add Redis identify ratelimiter

2022-01-13 12:26:25 -05:00
parent 8aa2b0b7aa
commit 2aead40449
15 changed files with 449 additions and 40 deletions
--- a/Myriad/Gateway/Cluster.cs
+++ b/Myriad/Gateway/Cluster.cs
@@ -5,6 +5,8 @@ using Myriad.Types;

 using Serilog;

+using StackExchange.Redis;
+
 namespace Myriad.Gateway;

 public class Cluster
@@ -25,14 +27,14 @@ public class Cluster
    public IReadOnlyDictionary<int, Shard> Shards => _shards;
    public event Action<Shard>? ShardCreated;

-    public async Task Start(GatewayInfo.Bot info)
+    public async Task Start(GatewayInfo.Bot info, ConnectionMultiplexer? conn = null)
    {
-        await Start(info.Url, 0, info.Shards - 1, info.Shards, info.SessionStartLimit.MaxConcurrency);
+        await Start(info.Url, 0, info.Shards - 1, info.Shards, info.SessionStartLimit.MaxConcurrency, conn);
    }

-    public async Task Start(string url, int shardMin, int shardMax, int shardTotal, int recommendedConcurrency)
+    public async Task Start(string url, int shardMin, int shardMax, int shardTotal, int recommendedConcurrency, ConnectionMultiplexer? conn = null)
    {
-        _ratelimiter = GetRateLimiter(recommendedConcurrency);
+        _ratelimiter = GetRateLimiter(recommendedConcurrency, conn);

        var shardCount = shardMax - shardMin + 1;
        _logger.Information("Starting {ShardCount} of {ShardTotal} shards (#{ShardMin}-#{ShardMax}) at {Url}",
@@ -73,12 +75,21 @@ public class Cluster
        return Math.Min(_gatewaySettings.MaxShardConcurrency.Value, recommendedConcurrency);
    }

-    private IGatewayRatelimiter GetRateLimiter(int recommendedConcurrency)
+    private IGatewayRatelimiter GetRateLimiter(int recommendedConcurrency, ConnectionMultiplexer? conn = null)
    {
+        var concurrency = GetActualShardConcurrency(recommendedConcurrency);
+
+        if (_gatewaySettings.UseRedisRatelimiter)
+        {
+            if (conn != null)
+                return new RedisRatelimiter(_logger, conn, concurrency);
+            else
+                _logger.Warning("Tried to get Redis ratelimiter but connection is null! Continuing with local ratelimiter.");
+        }
+
        if (_gatewaySettings.GatewayQueueUrl != null)
            return new TwilightGatewayRatelimiter(_logger, _gatewaySettings.GatewayQueueUrl);

-        var concurrency = GetActualShardConcurrency(recommendedConcurrency);
        return new LocalGatewayRatelimiter(_logger, concurrency);
    }
 }
--- a/Myriad/Gateway/GatewaySettings.cs
+++ b/Myriad/Gateway/GatewaySettings.cs
@@ -4,6 +4,7 @@ public record GatewaySettings
 {
    public string Token { get; init; }
    public GatewayIntent Intents { get; init; }
+    public bool UseRedisRatelimiter { get; init; } = false;
    public int? MaxShardConcurrency { get; init; }
    public string? GatewayQueueUrl { get; init; }
 }
--- a/Myriad/Gateway/Limit/RedisRatelimiter.cs
+++ b/Myriad/Gateway/Limit/RedisRatelimiter.cs
@@ -0,0 +1,46 @@
+using Serilog;
+
+using StackExchange.Redis;
+
+namespace Myriad.Gateway.Limit;
+
+public class RedisRatelimiter: IGatewayRatelimiter
+{
+    private readonly ILogger _logger;
+    private readonly ConnectionMultiplexer _redis;
+
+    private int _concurrency { get; init; }
+
+    // todo: these might need to be tweaked a little
+    private static TimeSpan expiry = TimeSpan.FromSeconds(5);
+    private static TimeSpan retryInterval = TimeSpan.FromSeconds(1);
+
+    public RedisRatelimiter(ILogger logger, ConnectionMultiplexer redis, int concurrency)
+    {
+        _logger = logger.ForContext<TwilightGatewayRatelimiter>();
+        _redis = redis;
+        _concurrency = concurrency;
+    }
+
+    public async Task Identify(int shard)
+    {
+        _logger.Information("Shard {ShardId}: requesting identify from Redis", shard);
+        var key = "pluralkit:identify:" + (shard % _concurrency).ToString();
+        await AcquireLock(key);
+    }
+
+    public async Task AcquireLock(string key)
+    {
+        var conn = _redis.GetDatabase();
+
+        async Task<bool> TryAcquire()
+        {
+            _logger.Verbose("Trying to acquire lock on key {key} from Redis...", key);
+            await Task.Delay(retryInterval);
+            return await conn!.StringSetAsync(key, 0, expiry, When.NotExists);
+        }
+
+        var acquired = false;
+        while (!acquired) acquired = await TryAcquire();
+    }
+}
--- a/Myriad/Myriad.csproj
+++ b/Myriad/Myriad.csproj
@@ -24,6 +24,7 @@
        <PackageReference Include="Polly" Version="7.2.1"/>
        <PackageReference Include="Polly.Contrib.WaitAndRetry" Version="1.1.1"/>
        <PackageReference Include="Serilog" Version="2.10.0"/>
+        <PackageReference Include="StackExchange.Redis" Version="2.2.88" />
        <PackageReference Include="System.Linq.Async" Version="5.0.0"/>
    </ItemGroup>

--- a/Myriad/packages.lock.json
+++ b/Myriad/packages.lock.json
@@ -1,31 +1,140 @@
 {
-    "version": 1,
-    "dependencies": {
-        "net6.0": {
-            "Polly": {
-                "type": "Direct",
-                "requested": "[7.2.1, )",
-                "resolved": "7.2.1",
-                "contentHash": "Od8SnPlpQr+PuAS0YzY3jgtzaDNknlIuAaldN2VEIyTvm/wCg22C5nUkUV1BEG8rIsub5RFMoV/NEQ0tM/+7Uw=="
-            },
-            "Polly.Contrib.WaitAndRetry": {
-                "type": "Direct",
-                "requested": "[1.1.1, )",
-                "resolved": "1.1.1",
-                "contentHash": "1MUQLiSo4KDkQe6nzQRhIU05lm9jlexX5BVsbuw0SL82ynZ+GzAHQxJVDPVBboxV37Po3SG077aX8DuSy8TkaA=="
-            },
-            "Serilog": {
-                "type": "Direct",
-                "requested": "[2.10.0, )",
-                "resolved": "2.10.0",
-                "contentHash": "+QX0hmf37a0/OZLxM3wL7V6/ADvC1XihXN4Kq/p6d8lCPfgkRdiuhbWlMaFjR9Av0dy5F0+MBeDmDdRZN/YwQA=="
-            },
-            "System.Linq.Async": {
-                "type": "Direct",
-                "requested": "[5.0.0, )",
-                "resolved": "5.0.0",
-                "contentHash": "cPtIuuH8TIjVHSi2ewwReWGW1PfChPE0LxPIDlfwVcLuTM9GANFTXiMB7k3aC4sk3f0cQU25LNKzx+jZMxijqw=="
-            }
+  "version": 1,
+  "dependencies": {
+    "net6.0": {
+      "Polly": {
+        "type": "Direct",
+        "requested": "[7.2.1, )",
+        "resolved": "7.2.1",
+        "contentHash": "Od8SnPlpQr+PuAS0YzY3jgtzaDNknlIuAaldN2VEIyTvm/wCg22C5nUkUV1BEG8rIsub5RFMoV/NEQ0tM/+7Uw=="
+      },
+      "Polly.Contrib.WaitAndRetry": {
+        "type": "Direct",
+        "requested": "[1.1.1, )",
+        "resolved": "1.1.1",
+        "contentHash": "1MUQLiSo4KDkQe6nzQRhIU05lm9jlexX5BVsbuw0SL82ynZ+GzAHQxJVDPVBboxV37Po3SG077aX8DuSy8TkaA=="
+      },
+      "Serilog": {
+        "type": "Direct",
+        "requested": "[2.10.0, )",
+        "resolved": "2.10.0",
+        "contentHash": "+QX0hmf37a0/OZLxM3wL7V6/ADvC1XihXN4Kq/p6d8lCPfgkRdiuhbWlMaFjR9Av0dy5F0+MBeDmDdRZN/YwQA=="
+      },
+      "StackExchange.Redis": {
+        "type": "Direct",
+        "requested": "[2.2.88, )",
+        "resolved": "2.2.88",
+        "contentHash": "JJi1jcO3/ZiamBhlsC/TR8aZmYf+nqpGzMi0HRRCy5wJkUPmMnRp0kBA6V84uhU8b531FHSdTDaFCAyCUJomjA==",
+        "dependencies": {
+          "Pipelines.Sockets.Unofficial": "2.2.0",
+          "System.Diagnostics.PerformanceCounter": "5.0.0"
        }
+      },
+      "System.Linq.Async": {
+        "type": "Direct",
+        "requested": "[5.0.0, )",
+        "resolved": "5.0.0",
+        "contentHash": "cPtIuuH8TIjVHSi2ewwReWGW1PfChPE0LxPIDlfwVcLuTM9GANFTXiMB7k3aC4sk3f0cQU25LNKzx+jZMxijqw=="
+      },
+      "Microsoft.NETCore.Platforms": {
+        "type": "Transitive",
+        "resolved": "5.0.0",
+        "contentHash": "VyPlqzH2wavqquTcYpkIIAQ6WdenuKoFN0BdYBbCWsclXacSOHNQn66Gt4z5NBqEYW0FAPm5rlvki9ZiCij5xQ=="
+      },
+      "Microsoft.Win32.Registry": {
+        "type": "Transitive",
+        "resolved": "5.0.0",
+        "contentHash": "dDoKi0PnDz31yAyETfRntsLArTlVAVzUzCIvvEDsDsucrl33Dl8pIJG06ePTJTI3tGpeyHS9Cq7Foc/s4EeKcg==",
+        "dependencies": {
+          "System.Security.AccessControl": "5.0.0",
+          "System.Security.Principal.Windows": "5.0.0"
+        }
+      },
+      "Microsoft.Win32.SystemEvents": {
+        "type": "Transitive",
+        "resolved": "5.0.0",
+        "contentHash": "Bh6blKG8VAKvXiLe2L+sEsn62nc1Ij34MrNxepD2OCrS5cpCwQa9MeLyhVQPQ/R4Wlzwuy6wMK8hLb11QPDRsQ==",
+        "dependencies": {
+          "Microsoft.NETCore.Platforms": "5.0.0"
+        }
+      },
+      "Pipelines.Sockets.Unofficial": {
+        "type": "Transitive",
+        "resolved": "2.2.0",
+        "contentHash": "7hzHplEIVOGBl5zOQZGX/DiJDHjq+RVRVrYgDiqXb6RriqWAdacXxp+XO9WSrATCEXyNOUOQg9aqQArsjase/A==",
+        "dependencies": {
+          "System.IO.Pipelines": "5.0.0"
+        }
+      },
+      "System.Configuration.ConfigurationManager": {
+        "type": "Transitive",
+        "resolved": "5.0.0",
+        "contentHash": "aM7cbfEfVNlEEOj3DsZP+2g9NRwbkyiAv2isQEzw7pnkDg9ekCU2m1cdJLM02Uq691OaCS91tooaxcEn8d0q5w==",
+        "dependencies": {
+          "System.Security.Cryptography.ProtectedData": "5.0.0",
+          "System.Security.Permissions": "5.0.0"
+        }
+      },
+      "System.Diagnostics.PerformanceCounter": {
+        "type": "Transitive",
+        "resolved": "5.0.0",
+        "contentHash": "kcQWWtGVC3MWMNXdMDWfrmIlFZZ2OdoeT6pSNVRtk9+Sa7jwdPiMlNwb0ZQcS7NRlT92pCfmjRtkSWUW3RAKwg==",
+        "dependencies": {
+          "Microsoft.NETCore.Platforms": "5.0.0",
+          "Microsoft.Win32.Registry": "5.0.0",
+          "System.Configuration.ConfigurationManager": "5.0.0",
+          "System.Security.Principal.Windows": "5.0.0"
+        }
+      },
+      "System.Drawing.Common": {
+        "type": "Transitive",
+        "resolved": "5.0.0",
+        "contentHash": "SztFwAnpfKC8+sEKXAFxCBWhKQaEd97EiOL7oZJZP56zbqnLpmxACWA8aGseaUExciuEAUuR9dY8f7HkTRAdnw==",
+        "dependencies": {
+          "Microsoft.Win32.SystemEvents": "5.0.0"
+        }
+      },
+      "System.IO.Pipelines": {
+        "type": "Transitive",
+        "resolved": "5.0.0",
+        "contentHash": "irMYm3vhVgRsYvHTU5b2gsT2CwT/SMM6LZFzuJjpIvT5Z4CshxNsaoBC1X/LltwuR3Opp8d6jOS/60WwOb7Q2Q=="
+      },
+      "System.Security.AccessControl": {
+        "type": "Transitive",
+        "resolved": "5.0.0",
+        "contentHash": "dagJ1mHZO3Ani8GH0PHpPEe/oYO+rVdbQjvjJkBRNQkX4t0r1iaeGn8+/ybkSLEan3/slM0t59SVdHzuHf2jmw==",
+        "dependencies": {
+          "Microsoft.NETCore.Platforms": "5.0.0",
+          "System.Security.Principal.Windows": "5.0.0"
+        }
+      },
+      "System.Security.Cryptography.ProtectedData": {
+        "type": "Transitive",
+        "resolved": "5.0.0",
+        "contentHash": "HGxMSAFAPLNoxBvSfW08vHde0F9uh7BjASwu6JF9JnXuEPhCY3YUqURn0+bQV/4UWeaqymmrHWV+Aw9riQCtCA=="
+      },
+      "System.Security.Permissions": {
+        "type": "Transitive",
+        "resolved": "5.0.0",
+        "contentHash": "uE8juAhEkp7KDBCdjDIE3H9R1HJuEHqeqX8nLX9gmYKWwsqk3T5qZlPx8qle5DPKimC/Fy3AFTdV7HamgCh9qQ==",
+        "dependencies": {
+          "System.Security.AccessControl": "5.0.0",
+          "System.Windows.Extensions": "5.0.0"
+        }
+      },
+      "System.Security.Principal.Windows": {
+        "type": "Transitive",
+        "resolved": "5.0.0",
+        "contentHash": "t0MGLukB5WAVU9bO3MGzvlGnyJPgUlcwerXn1kzBRjwLKixT96XV0Uza41W49gVd8zEMFu9vQEFlv0IOrytICA=="
+      },
+      "System.Windows.Extensions": {
+        "type": "Transitive",
+        "resolved": "5.0.0",
+        "contentHash": "c1ho9WU9ZxMZawML+ssPKZfdnrg/OjR3pe0m9v8230z3acqphwvPJqzAkH54xRYm5ntZHGG1EPP3sux9H3qSPg==",
+        "dependencies": {
+          "System.Drawing.Common": "5.0.0"
+        }
+      }
    }
+  }
 }