Add proper webhook rate limit support

PluralKit.Bot/Modules.cs

@@ -60,6 +60,7 @@ namespace PluralKit.Bot
             builder.RegisterType<ProxyService>().AsSelf().SingleInstance();
             builder.RegisterType<LogChannelService>().AsSelf().SingleInstance();
             builder.RegisterType<DataFileService>().AsSelf().SingleInstance();
+            builder.RegisterType<WebhookRateLimitService>().AsSelf().SingleInstance();
             builder.RegisterType<WebhookExecutorService>().AsSelf().SingleInstance();
             builder.RegisterType<WebhookCacheService>().AsSelf().SingleInstance();
             builder.RegisterType<ShardInfoService>().AsSelf().SingleInstance();

PluralKit.Bot/Services/WebhookExecutorService.cs

@@ -22,18 +22,25 @@ namespace PluralKit.Bot
     public class WebhookExecutionErrorOnDiscordsEnd: Exception {
     }
     
+    public class WebhookRateLimited: WebhookExecutionErrorOnDiscordsEnd {
+        // Exceptions for control flow? don't mind if I do
+        // TODO: rewrite both of these as a normal exceptional return value (0?) in case of error to be discarded by caller 
+    }
+    
     public class WebhookExecutorService
     {
         private WebhookCacheService _webhookCache;
+        private WebhookRateLimitService _rateLimit;
         private ILogger _logger;
         private IMetrics _metrics;
         private HttpClient _client;
 
-        public WebhookExecutorService(IMetrics metrics, WebhookCacheService webhookCache, ILogger logger, HttpClient client)
+        public WebhookExecutorService(IMetrics metrics, WebhookCacheService webhookCache, ILogger logger, HttpClient client, WebhookRateLimitService rateLimit)
         {
             _metrics = metrics;
             _webhookCache = webhookCache;
             _client = client;
+            _rateLimit = rateLimit;
             _logger = logger.ForContext<WebhookExecutorService>();
         }
 
@@ -56,7 +63,6 @@ namespace PluralKit.Bot
         private async Task<ulong> ExecuteWebhookInner(IWebhook webhook, string name, string avatarUrl, string content,
             IReadOnlyCollection<IAttachment> attachments, bool hasRetried = false)
         {
-
             using var mfd = new MultipartFormDataContent
             {
                 {new StringContent(content.Truncate(2000)), "content"},
@@ -70,15 +76,23 @@ namespace PluralKit.Bot
                 _logger.Information("Invoking webhook with {AttachmentCount} attachments totalling {AttachmentSize} MiB in {AttachmentChunks} chunks", attachments.Count, attachments.Select(a => a.Size).Sum() / 1024 / 1024, attachmentChunks.Count);
                 await AddAttachmentsToMultipart(mfd, attachmentChunks.First());
             }
+            
+            mfd.Headers.Add("X-RateLimit-Precision", "millisecond"); // Need this for better rate limit support
+            
+            // Adding this check as close to the actual send call as possible to prevent potential race conditions (unlikely, but y'know)
+            if (!_rateLimit.TryExecuteWebhook(webhook))
+                throw new WebhookRateLimited();
 
             var timerCtx = _metrics.Measure.Timer.Time(BotMetrics.WebhookResponseTime);
             using var response = await _client.PostAsync($"{DiscordConfig.APIUrl}webhooks/{webhook.Id}/{webhook.Token}?wait=true", mfd);
             timerCtx.Dispose();
+            
+            _rateLimit.UpdateRateLimitInfo(webhook, response);
 
             if (response.StatusCode == HttpStatusCode.TooManyRequests)
-                // Rate limits should be respected, we bail early
+                // Rate limits should be respected, we bail early (already updated the limit info so we hopefully won't hit this again)
-                throw new WebhookExecutionErrorOnDiscordsEnd();
+                throw new WebhookRateLimited();
-
+            
             var responseString = await response.Content.ReadAsStringAsync();
 
             JObject responseJson;
@@ -136,7 +150,6 @@ namespace PluralKit.Bot
             // TODO: can we do this without a round-trip to a string?
             return responseJson["id"].Value<ulong>();
         }
-
         private IReadOnlyCollection<IReadOnlyCollection<IAttachment>> ChunkAttachmentsOrThrow(
             IReadOnlyCollection<IAttachment> attachments, int sizeThreshold)
         {

PluralKit.Bot/Services/WebhookRateLimitService.cs

@@ -0,0 +1,104 @@
+using System;
+using System.Collections.Concurrent;
+using System.Globalization;
+using System.Linq;
+using System.Net;
+using System.Net.Http;
+
+using Discord;
+
+using NodaTime;
+
+using Serilog;
+
+namespace PluralKit.Bot
+{
+    // Simplified rate limit handler for webhooks only, disregards most bucket functionality since scope is limited and just denies requests if too fast.
+    public class WebhookRateLimitService
+    {
+        private ILogger _logger;
+        private ConcurrentDictionary<ulong, WebhookRateLimitInfo> _info = new ConcurrentDictionary<ulong, WebhookRateLimitInfo>();
+
+        public WebhookRateLimitService(ILogger logger)
+        {
+            _logger = logger.ForContext<WebhookRateLimitService>();
+        }
+
+        public bool TryExecuteWebhook(IWebhook webhook)
+        {
+            // If we have nothing saved, just allow it (we'll save something once the response returns)
+            if (!_info.TryGetValue(webhook.Id, out var info)) return true;
+            
+            // If we're past the reset time, allow the request and update the bucket limit
+            if (SystemClock.Instance.GetCurrentInstant() > info.resetTime)
+            {
+                if (!info.hasResetTimeExpired)
+                    info.remaining = info.maxLimit;
+                info.hasResetTimeExpired = true;
+                return true;
+            }
+
+            // If we don't have any more requests left, deny the request
+            if (info.remaining == 0)
+            {
+                _logger.Debug("Rate limit bucket for {Webhook} out of requests, denying request", webhook.Id);
+                return false;
+            }
+            
+            // Otherwise, decrement the request count and allow the request
+            info.remaining--;
+            return true;
+        }
+
+        public void UpdateRateLimitInfo(IWebhook webhook, HttpResponseMessage response)
+        {
+            var info = _info.GetOrAdd(webhook.Id, _ => new WebhookRateLimitInfo());
+
+            if (int.TryParse(GetHeader(response, "X-RateLimit-Limit"), out var limit))
+                info.maxLimit = limit;
+
+            // Max "safe" is way above UNIX timestamp values, and we get fractional seconds, hence the double
+            // but need culture/format specifiers to get around Some Locales (cough, my local PC) having different settings for decimal point...
+            // We also use Reset-After to avoid issues with clock desync between us and Discord's server, this way it's all relative (plus latency errors work in our favor)    
+            if (double.TryParse(GetHeader(response, "X-RateLimit-Reset-After"), NumberStyles.Float, CultureInfo.InvariantCulture, out var resetTimestampDelta))
+            {
+                var resetTime = SystemClock.Instance.GetCurrentInstant() + Duration.FromSeconds(resetTimestampDelta);
+                if (resetTime > info.resetTime)
+                {
+                    // Set to the *latest* reset value we have (for safety), since we rely on relative times this can jitter a bit
+                    info.resetTime = resetTime;
+                    info.hasResetTimeExpired = false;
+                }
+            }
+            
+            if (int.TryParse(GetHeader(response, "X-RateLimit-Remaining"), out var remainingRequests))
+                // Overwrite a negative "we don't know" value with whatever we just got
+                // Otherwise, *lower* remaining requests takes precedence
+                if (info.remaining < 0 || remainingRequests < info.remaining)
+                    info.remaining = remainingRequests;
+            
+            _logger.Debug("Updated rate limit information for {Webhook}, bucket has {RequestsRemaining} requests remaining, reset in {ResetTime}", webhook.Id, info.remaining, info.resetTime - SystemClock.Instance.GetCurrentInstant());
+            
+            if (response.StatusCode == HttpStatusCode.TooManyRequests)
+            {
+                // 429, we're *definitely* out of requests
+                info.remaining = 0;
+                _logger.Warning("Got 429 Too Many Requests when invoking webhook {Webhook}, next bucket reset in {ResetTime}", webhook.Id, info.resetTime - SystemClock.Instance.GetCurrentInstant());
+            }
+        }
+
+        private string GetHeader(HttpResponseMessage response, string key)
+        {
+            var firstPair = response.Headers.FirstOrDefault(pair => pair.Key.Equals(key, StringComparison.InvariantCultureIgnoreCase));
+            return firstPair.Value?.FirstOrDefault(); // If key is missing, default value is null
+        }
+
+        private class WebhookRateLimitInfo
+        {
+            public Instant resetTime;
+            public bool hasResetTimeExpired;
+            public int remaining = -1;
+            public int maxLimit = 0;
+        } 
+    }
+}

PluralKit.Bot/Utils/MiscUtils.cs

@@ -23,6 +23,7 @@ namespace PluralKit.Bot
             if (e is HttpException he && ((int) he.HttpCode) >= 500) return false;
             
             // Webhook server errors are also *not our problem*
+            // (this includes rate limit errors, WebhookRateLimited is a subclass)
             if (e is WebhookExecutionErrorOnDiscordsEnd) return false;
             
             // Socket errors are *not our problem*