foxcord/Foxcord/Rest/Rate/RateLimiter.cs

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<RateLimiter>();
    private readonly ConcurrentDictionary<string, Bucket> _buckets = new();

    private readonly ConcurrentDictionary<string, CustomRateLimit> _customRateLimits = new([
        new KeyValuePair<string, CustomRateLimit>("//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<Bucket> LockBucket(string bucketId, CancellationToken ct = default) =>
        await LockBucket(GetBucket(bucketId), ct);

    internal async Task<Bucket> 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);