using System.Collections.Concurrent; using System.Globalization; using System.Net.Http.Headers; using Serilog; namespace Foxcord.Rest.Rate; // Most of this code is taken from discordgo: // https://github.com/bwmarrin/discordgo/blob/master/ratelimit.go public class RateLimiter(ILogger logger) { private readonly ILogger _logger = logger.ForContext(); private readonly ConcurrentDictionary _buckets = new(); private readonly ConcurrentDictionary _customRateLimits = new([ new KeyValuePair("//reactions//", new CustomRateLimit { Requests = 1, Reset = TimeSpan.FromMilliseconds(200) }) ]); internal long Global; internal Bucket GetBucket(string key) { key = BucketKeyUtils.Parse(key); var bucket = _buckets.GetOrAdd(key, _ => new Bucket { Key = key, Remaining = 1, RateLimiter = this, Logger = _logger }); if (_customRateLimits.Any(r => key.EndsWith(r.Key))) bucket.CustomRateLimit = _customRateLimits.First(r => key.EndsWith(r.Key)).Value; return bucket; } internal TimeSpan GetWaitTime(Bucket b, int minRemaining) { if (b.Remaining < minRemaining && b.Reset > DateTimeOffset.UtcNow) return b.Reset - DateTimeOffset.UtcNow; var sleepTo = DateTimeOffset.FromUnixTimeMilliseconds(Global); if (sleepTo > DateTimeOffset.UtcNow) return sleepTo - DateTimeOffset.UtcNow; return TimeSpan.Zero; } internal async Task LockBucket(string bucketId, CancellationToken ct = default) => await LockBucket(GetBucket(bucketId), ct); internal async Task LockBucket(Bucket b, CancellationToken ct = default) { _logger.Verbose("Locking bucket {Bucket}", b.Key); await b.Semaphore.WaitAsync(ct); var waitTime = GetWaitTime(b, 1); if (waitTime > TimeSpan.Zero) await Task.Delay(waitTime, ct); b.Remaining--; _logger.Verbose("Letting request for bucket {Bucket} through", b.Key); return b; } } internal class CustomRateLimit { internal int Requests; internal TimeSpan Reset; } internal class Bucket { internal readonly SemaphoreSlim Semaphore = new(1); internal required string Key; internal required ILogger Logger { private get; init; } internal int Remaining; internal DateTimeOffset Reset; private DateTimeOffset _lastReset; internal CustomRateLimit? CustomRateLimit; internal required RateLimiter RateLimiter; // discordgo mentions that this is required to prevent 429s, I trust that private static readonly TimeSpan ExtraResetTime = TimeSpan.FromMilliseconds(250); internal void Release(HttpHeaders headers) { try { if (CustomRateLimit != null) { if (DateTimeOffset.UtcNow - _lastReset >= CustomRateLimit.Reset) { Remaining = CustomRateLimit.Requests - 1; _lastReset = DateTimeOffset.UtcNow; } if (Remaining < 1) { Reset = DateTimeOffset.UtcNow + CustomRateLimit.Reset; } return; } var remaining = TryGetHeaderValue(headers, "X-RateLimit-Remaining"); var reset = TryGetHeaderValue(headers, "X-RateLimit-Reset"); var global = TryGetHeaderValue(headers, "X-RateLimit-Global"); var resetAfter = TryGetHeaderValue(headers, "X-RateLimit-Reset-After"); if (resetAfter != null) { if (!double.TryParse(resetAfter, out var parsedResetAfter)) throw new InvalidRateLimitHeaderException("X-RateLimit-Reset-After was not a valid double"); var resetAt = DateTimeOffset.UtcNow + TimeSpan.FromSeconds(parsedResetAfter); if (global != null) RateLimiter.Global = resetAt.ToUnixTimeMilliseconds(); else Reset = resetAt; } else if (reset != null) { var dateHeader = TryGetHeaderValue(headers, "Date"); if (dateHeader == null) throw new InvalidRateLimitHeaderException("Date header was not set"); if (!DateTimeOffset.TryParseExact(dateHeader, "r", CultureInfo.InvariantCulture, DateTimeStyles.AssumeUniversal, out var parsedDate)) throw new InvalidRateLimitHeaderException("Date was not a valid date"); if (!long.TryParse(reset, out var parsedReset)) throw new InvalidRateLimitHeaderException("X-RateLimit-Reset was not a valid long"); var delta = DateTimeOffset.FromUnixTimeMilliseconds(parsedReset) - parsedDate + ExtraResetTime; Reset = DateTimeOffset.UtcNow + delta; } if (remaining == null) return; if (!int.TryParse(remaining, out var parsedRemaining)) throw new InvalidRateLimitHeaderException("X-RateLimit-Remaining was not a valid integer"); Remaining = parsedRemaining; Logger.Verbose("New remaining for bucket {Bucket} is {Remaining}", Key, Remaining); } finally { Logger.Verbose("Releasing bucket {Bucket}", Key); Semaphore.Release(); } } private static string? TryGetHeaderValue(HttpHeaders headers, string key) => headers.TryGetValues(key, out var values) ? values.FirstOrDefault() : null; } public class InvalidRateLimitHeaderException(string message) : Exception(message);