PluralKit/Myriad/Rest/Ratelimit/Ratelimiter.cs

using Myriad.Rest.Exceptions;

using Serilog;

namespace Myriad.Rest.Ratelimit;

public class Ratelimiter: IDisposable
{
    private readonly BucketManager _buckets;
    private readonly ILogger _logger;

    private DateTimeOffset? _globalRateLimitExpiry;

    public Ratelimiter(ILogger logger)
    {
        _logger = logger.ForContext<Ratelimiter>();
        _buckets = new BucketManager(logger);
    }

    public void Dispose()
    {
        _buckets.Dispose();
    }

    public void AllowRequestOrThrow(string endpoint, ulong major, DateTimeOffset now)
    {
        if (IsGloballyRateLimited(now))
        {
            _logger.Warning("Globally rate limited until {GlobalRateLimitExpiry}, cancelling request",
                _globalRateLimitExpiry);
            throw new GloballyRatelimitedException();
        }

        var bucket = _buckets.GetBucket(endpoint, major);
        if (bucket == null)
        {
            // No rate limit for this endpoint (yet), allow through
            _logger.Debug("No rate limit data for endpoint {Endpoint}, allowing through", endpoint);
            return;
        }

        bucket.Tick(now);

        if (bucket.TryAcquire())
            // We're allowed to send it! :)
            return;

        // We can't send this request right now; retrying...
        var waitTime = bucket.GetResetDelay(now);

        // add a small buffer for Timing:tm:
        waitTime += TimeSpan.FromMilliseconds(50);

        // (this is caught by a WaitAndRetry Polly handler, if configured)
        throw new RatelimitBucketExhaustedException(bucket, waitTime);
    }

    public void HandleResponse(RatelimitHeaders headers, string endpoint, ulong major)
    {
        if (!headers.HasRatelimitInfo)
            return;

        // TODO: properly calculate server time?
        if (headers.Global)
        {
            _logger.Warning(
                "Global rate limit hit, resetting at {GlobalRateLimitExpiry} (in {GlobalRateLimitResetAfter}!",
                _globalRateLimitExpiry, headers.ResetAfter);
            _globalRateLimitExpiry = headers.Reset;
        }
        else
        {
            // Update buckets first, then get it again, to properly "transfer" this info over to the new value
            _buckets.UpdateEndpointInfo(endpoint, headers.Bucket!, headers.Limit);

            var bucket = _buckets.GetBucket(endpoint, major);
            bucket?.HandleResponse(headers);
        }
    }

    private bool IsGloballyRateLimited(DateTimeOffset now) =>
        _globalRateLimitExpiry > now;
}