From fa70df8f9830628667ec3ddac062b5b6ae243181 Mon Sep 17 00:00:00 2001 From: Ske Date: Fri, 27 Mar 2020 00:01:42 +0100 Subject: [PATCH] Add proper webhook rate limit support --- PluralKit.Bot/Modules.cs | 1 + .../Services/WebhookExecutorService.cs | 25 ++++- .../Services/WebhookRateLimitService.cs | 104 ++++++++++++++++++ PluralKit.Bot/Utils/MiscUtils.cs | 1 + 4 files changed, 125 insertions(+), 6 deletions(-) create mode 100644 PluralKit.Bot/Services/WebhookRateLimitService.cs diff --git a/PluralKit.Bot/Modules.cs b/PluralKit.Bot/Modules.cs index aa371e38..e4e0c67b 100644 --- a/PluralKit.Bot/Modules.cs +++ b/PluralKit.Bot/Modules.cs @@ -60,6 +60,7 @@ namespace PluralKit.Bot builder.RegisterType().AsSelf().SingleInstance(); builder.RegisterType().AsSelf().SingleInstance(); builder.RegisterType().AsSelf().SingleInstance(); + builder.RegisterType().AsSelf().SingleInstance(); builder.RegisterType().AsSelf().SingleInstance(); builder.RegisterType().AsSelf().SingleInstance(); builder.RegisterType().AsSelf().SingleInstance(); diff --git a/PluralKit.Bot/Services/WebhookExecutorService.cs b/PluralKit.Bot/Services/WebhookExecutorService.cs index fb963d49..104d379b 100644 --- a/PluralKit.Bot/Services/WebhookExecutorService.cs +++ b/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(); } @@ -56,7 +63,6 @@ namespace PluralKit.Bot private async Task ExecuteWebhookInner(IWebhook webhook, string name, string avatarUrl, string content, IReadOnlyCollection 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 - throw new WebhookExecutionErrorOnDiscordsEnd(); - + // Rate limits should be respected, we bail early (already updated the limit info so we hopefully won't hit this again) + 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(); } - private IReadOnlyCollection> ChunkAttachmentsOrThrow( IReadOnlyCollection attachments, int sizeThreshold) { diff --git a/PluralKit.Bot/Services/WebhookRateLimitService.cs b/PluralKit.Bot/Services/WebhookRateLimitService.cs new file mode 100644 index 00000000..0dddd9d3 --- /dev/null +++ b/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 _info = new ConcurrentDictionary(); + + public WebhookRateLimitService(ILogger logger) + { + _logger = logger.ForContext(); + } + + 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; + } + } +} \ No newline at end of file diff --git a/PluralKit.Bot/Utils/MiscUtils.cs b/PluralKit.Bot/Utils/MiscUtils.cs index 238c1711..ed4d63dc 100644 --- a/PluralKit.Bot/Utils/MiscUtils.cs +++ b/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*