PluralKit/Myriad/Rest/Ratelimit/Bucket.cs

173 lines
6.2 KiB
C#
Raw Normal View History

2020-12-22 12:15:26 +00:00
using System;
using System.Threading;
using Serilog;
namespace Myriad.Rest.Ratelimit
{
public class Bucket
{
private static readonly TimeSpan Epsilon = TimeSpan.FromMilliseconds(10);
private static readonly TimeSpan FallbackDelay = TimeSpan.FromMilliseconds(200);
private static readonly TimeSpan StaleTimeout = TimeSpan.FromSeconds(5);
private readonly ILogger _logger;
private readonly SemaphoreSlim _semaphore = new(1, 1);
2021-01-30 00:07:43 +00:00
private DateTimeOffset? _nextReset;
2020-12-22 12:15:26 +00:00
private bool _resetTimeValid;
2021-01-30 00:07:43 +00:00
private bool _hasReceivedHeaders;
2020-12-22 12:15:26 +00:00
public Bucket(ILogger logger, string key, ulong major, int limit)
{
_logger = logger.ForContext<Bucket>();
Key = key;
Major = major;
Limit = limit;
Remaining = limit;
_resetTimeValid = false;
}
public string Key { get; }
public ulong Major { get; }
public int Remaining { get; private set; }
public int Limit { get; private set; }
public DateTimeOffset LastUsed { get; private set; } = DateTimeOffset.UtcNow;
public bool TryAcquire()
{
LastUsed = DateTimeOffset.Now;
try
{
_semaphore.Wait();
if (Remaining > 0)
{
_logger.Debug(
"{BucketKey}/{BucketMajor}: Bucket has [{BucketRemaining}/{BucketLimit} left], allowing through",
Key, Major, Remaining, Limit);
Remaining--;
2021-01-30 00:07:43 +00:00
2020-12-22 12:15:26 +00:00
return true;
}
_logger.Debug("{BucketKey}/{BucketMajor}: Bucket has [{BucketRemaining}/{BucketLimit}] left, denying",
Key, Major, Remaining, Limit);
return false;
}
finally
{
_semaphore.Release();
}
}
public void HandleResponse(RatelimitHeaders headers)
{
try
{
_semaphore.Wait();
2021-01-31 15:02:34 +00:00
_logger.Verbose("{BucketKey}/{BucketMajor}: Received rate limit headers: {@RateLimitHeaders}",
Key, Major, headers);
2020-12-22 12:15:26 +00:00
if (headers.ResetAfter != null)
{
var headerNextReset = DateTimeOffset.UtcNow + headers.ResetAfter.Value; // todo: server time
2021-01-31 15:02:34 +00:00
if (_nextReset == null || headerNextReset > _nextReset)
2020-12-22 12:15:26 +00:00
{
2021-01-30 00:07:43 +00:00
_logger.Debug("{BucketKey}/{BucketMajor}: Received reset time {NextReset} from server (after: {NextResetAfter}, remaining: {Remaining}, local remaining: {LocalRemaining})",
Key, Major, headerNextReset, headers.ResetAfter.Value, headers.Remaining, Remaining);
2020-12-22 12:15:26 +00:00
_nextReset = headerNextReset;
_resetTimeValid = true;
}
}
2021-01-30 00:07:43 +00:00
if (headers.Limit != null)
2020-12-22 12:15:26 +00:00
Limit = headers.Limit.Value;
2021-01-30 00:07:43 +00:00
if (headers.Remaining != null && !_hasReceivedHeaders)
{
2021-01-30 00:07:43 +00:00
var oldRemaining = Remaining;
Remaining = Math.Min(headers.Remaining.Value, Remaining);
_logger.Debug("{BucketKey}/{BucketMajor}: Received first remaining of {HeaderRemaining}, previous local remaining is {LocalRemaining}, new local remaining is {Remaining}",
Key, Major, headers.Remaining.Value, oldRemaining, Remaining);
_hasReceivedHeaders = true;
}
2020-12-22 12:15:26 +00:00
}
finally
{
_semaphore.Release();
}
}
public void Tick(DateTimeOffset now)
{
try
{
_semaphore.Wait();
2021-01-30 00:07:43 +00:00
// If we don't have any reset data, "snap" it to now
// This happens before first request and at this point the reset is invalid anyway, so it's fine
// but it ensures the stale timeout doesn't trigger early by using `default` value
if (_nextReset == null)
_nextReset = now;
2020-12-22 12:15:26 +00:00
// If we're past the reset time *and* we haven't reset already, do that
2020-12-24 13:52:44 +00:00
var timeSinceReset = now - _nextReset;
2020-12-22 12:15:26 +00:00
var shouldReset = _resetTimeValid && timeSinceReset > TimeSpan.Zero;
if (shouldReset)
{
_logger.Debug("{BucketKey}/{BucketMajor}: Bucket timed out, refreshing with {BucketLimit} requests",
Key, Major, Limit);
Remaining = Limit;
_resetTimeValid = false;
return;
}
// We've run out of requests without having any new reset time,
// *and* it's been longer than a set amount - add one request back to the pool and hope that one returns
var isBucketStale = !_resetTimeValid && Remaining <= 0 && timeSinceReset > StaleTimeout;
if (isBucketStale)
{
_logger.Warning(
"{BucketKey}/{BucketMajor}: Bucket is stale ({StaleTimeout} passed with no rate limit info), allowing one request through",
Key, Major, StaleTimeout);
Remaining = 1;
// Reset the (still-invalid) reset time to now, so we don't keep hitting this conditional over and over...
_nextReset = now;
}
}
finally
{
_semaphore.Release();
}
}
public TimeSpan GetResetDelay(DateTimeOffset now)
{
// If we don't have a valid reset time, return the fallback delay always
// (so it'll keep spinning until we hopefully have one...)
if (!_resetTimeValid)
return FallbackDelay;
2021-01-30 00:07:43 +00:00
var delay = (_nextReset ?? now) - now;
2020-12-22 12:15:26 +00:00
// If we have a really small (or negative) value, return a fallback delay too
if (delay < Epsilon)
return FallbackDelay;
return delay;
}
}
}