Initial commit, basic proxying working
This commit is contained in:
Ske
152
Myriad/Rest/Ratelimit/Bucket.cs
Normal file
152
Myriad/Rest/Ratelimit/Bucket.cs
Normal file
@@ -0,0 +1,152 @@
|
||||
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);
|
||||
|
||||
private DateTimeOffset _nextReset;
|
||||
private bool _resetTimeValid;
|
||||
|
||||
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--;
|
||||
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();
|
||||
|
||||
if (headers.ResetAfter != null)
|
||||
{
|
||||
var headerNextReset = DateTimeOffset.UtcNow + headers.ResetAfter.Value; // todo: server time
|
||||
if (headerNextReset > _nextReset)
|
||||
{
|
||||
_logger.Debug("{BucketKey}/{BucketMajor}: Received reset time {NextReset} from server",
|
||||
Key, Major, _nextReset);
|
||||
|
||||
_nextReset = headerNextReset;
|
||||
_resetTimeValid = true;
|
||||
}
|
||||
}
|
||||
|
||||
if (headers.Limit != null)
|
||||
Limit = headers.Limit.Value;
|
||||
}
|
||||
finally
|
||||
{
|
||||
_semaphore.Release();
|
||||
}
|
||||
}
|
||||
|
||||
public void Tick(DateTimeOffset now)
|
||||
{
|
||||
try
|
||||
{
|
||||
_semaphore.Wait();
|
||||
|
||||
// If we're past the reset time *and* we haven't reset already, do that
|
||||
var timeSinceReset = _nextReset - now;
|
||||
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;
|
||||
|
||||
var delay = _nextReset - now;
|
||||
|
||||
// If we have a really small (or negative) value, return a fallback delay too
|
||||
if (delay < Epsilon)
|
||||
return FallbackDelay;
|
||||
|
||||
return delay;
|
||||
}
|
||||
}
|
||||
}
|
||||
79
Myriad/Rest/Ratelimit/BucketManager.cs
Normal file
79
Myriad/Rest/Ratelimit/BucketManager.cs
Normal file
@@ -0,0 +1,79 @@
|
||||
using System;
|
||||
using System.Collections.Concurrent;
|
||||
using System.Threading;
|
||||
using System.Threading.Tasks;
|
||||
|
||||
using Serilog;
|
||||
|
||||
namespace Myriad.Rest.Ratelimit
|
||||
{
|
||||
public class BucketManager: IDisposable
|
||||
{
|
||||
private static readonly TimeSpan StaleBucketTimeout = TimeSpan.FromMinutes(5);
|
||||
private static readonly TimeSpan PruneWorkerInterval = TimeSpan.FromMinutes(1);
|
||||
private readonly ConcurrentDictionary<(string key, ulong major), Bucket> _buckets = new();
|
||||
|
||||
private readonly ConcurrentDictionary<string, string> _endpointKeyMap = new();
|
||||
private readonly ConcurrentDictionary<string, int> _knownKeyLimits = new();
|
||||
|
||||
private readonly ILogger _logger;
|
||||
|
||||
private readonly Task _worker;
|
||||
private readonly CancellationTokenSource _workerCts = new();
|
||||
|
||||
public BucketManager(ILogger logger)
|
||||
{
|
||||
_logger = logger.ForContext<BucketManager>();
|
||||
_worker = PruneWorker(_workerCts.Token);
|
||||
}
|
||||
|
||||
public void Dispose()
|
||||
{
|
||||
_workerCts.Dispose();
|
||||
_worker.Dispose();
|
||||
}
|
||||
|
||||
public Bucket? GetBucket(string endpoint, ulong major)
|
||||
{
|
||||
if (!_endpointKeyMap.TryGetValue(endpoint, out var key))
|
||||
return null;
|
||||
|
||||
if (_buckets.TryGetValue((key, major), out var bucket))
|
||||
return bucket;
|
||||
|
||||
if (!_knownKeyLimits.TryGetValue(key, out var knownLimit))
|
||||
return null;
|
||||
|
||||
return _buckets.GetOrAdd((key, major),
|
||||
k => new Bucket(_logger, k.Item1, k.Item2, knownLimit));
|
||||
}
|
||||
|
||||
public void UpdateEndpointInfo(string endpoint, string key, int? limit)
|
||||
{
|
||||
_endpointKeyMap[endpoint] = key;
|
||||
|
||||
if (limit != null)
|
||||
_knownKeyLimits[key] = limit.Value;
|
||||
}
|
||||
|
||||
private async Task PruneWorker(CancellationToken ct)
|
||||
{
|
||||
while (!ct.IsCancellationRequested)
|
||||
{
|
||||
await Task.Delay(PruneWorkerInterval, ct);
|
||||
PruneStaleBuckets(DateTimeOffset.UtcNow);
|
||||
}
|
||||
}
|
||||
|
||||
private void PruneStaleBuckets(DateTimeOffset now)
|
||||
{
|
||||
foreach (var (key, bucket) in _buckets)
|
||||
if (now - bucket.LastUsed > StaleBucketTimeout)
|
||||
{
|
||||
_logger.Debug("Pruning unused bucket {Bucket} (last used at {BucketLastUsed})", bucket,
|
||||
bucket.LastUsed);
|
||||
_buckets.TryRemove(key, out _);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
46
Myriad/Rest/Ratelimit/DiscordRateLimitPolicy.cs
Normal file
46
Myriad/Rest/Ratelimit/DiscordRateLimitPolicy.cs
Normal file
@@ -0,0 +1,46 @@
|
||||
using System;
|
||||
using System.Net.Http;
|
||||
using System.Threading;
|
||||
using System.Threading.Tasks;
|
||||
|
||||
using Polly;
|
||||
|
||||
namespace Myriad.Rest.Ratelimit
|
||||
{
|
||||
public class DiscordRateLimitPolicy: AsyncPolicy<HttpResponseMessage>
|
||||
{
|
||||
public const string EndpointContextKey = "Endpoint";
|
||||
public const string MajorContextKey = "Major";
|
||||
|
||||
private readonly Ratelimiter _ratelimiter;
|
||||
|
||||
public DiscordRateLimitPolicy(Ratelimiter ratelimiter, PolicyBuilder<HttpResponseMessage>? policyBuilder = null)
|
||||
: base(policyBuilder)
|
||||
{
|
||||
_ratelimiter = ratelimiter;
|
||||
}
|
||||
|
||||
protected override async Task<HttpResponseMessage> ImplementationAsync(
|
||||
Func<Context, CancellationToken, Task<HttpResponseMessage>> action, Context context, CancellationToken ct,
|
||||
bool continueOnCapturedContext)
|
||||
{
|
||||
if (!context.TryGetValue(EndpointContextKey, out var endpointObj) || !(endpointObj is string endpoint))
|
||||
throw new ArgumentException("Must provide endpoint in Polly context");
|
||||
|
||||
if (!context.TryGetValue(MajorContextKey, out var majorObj) || !(majorObj is ulong major))
|
||||
throw new ArgumentException("Must provide major in Polly context");
|
||||
|
||||
// Check rate limit, throw if we're not allowed...
|
||||
_ratelimiter.AllowRequestOrThrow(endpoint, major, DateTimeOffset.Now);
|
||||
|
||||
// We're OK, push it through
|
||||
var response = await action(context, ct).ConfigureAwait(continueOnCapturedContext);
|
||||
|
||||
// Update rate limit state with headers
|
||||
var headers = new RatelimitHeaders(response);
|
||||
_ratelimiter.HandleResponse(headers, endpoint, major);
|
||||
|
||||
return response;
|
||||
}
|
||||
}
|
||||
}
|
||||
46
Myriad/Rest/Ratelimit/RatelimitHeaders.cs
Normal file
46
Myriad/Rest/Ratelimit/RatelimitHeaders.cs
Normal file
@@ -0,0 +1,46 @@
|
||||
using System;
|
||||
using System.Linq;
|
||||
using System.Net.Http;
|
||||
|
||||
namespace Myriad.Rest.Ratelimit
|
||||
{
|
||||
public record RatelimitHeaders
|
||||
{
|
||||
public RatelimitHeaders() { }
|
||||
|
||||
public RatelimitHeaders(HttpResponseMessage response)
|
||||
{
|
||||
ServerDate = response.Headers.Date;
|
||||
|
||||
if (response.Headers.TryGetValues("X-RateLimit-Limit", out var limit))
|
||||
Limit = int.Parse(limit!.First());
|
||||
|
||||
if (response.Headers.TryGetValues("X-RateLimit-Remaining", out var remaining))
|
||||
Remaining = int.Parse(remaining!.First());
|
||||
|
||||
if (response.Headers.TryGetValues("X-RateLimit-Reset", out var reset))
|
||||
Reset = DateTimeOffset.FromUnixTimeMilliseconds((long) (double.Parse(reset!.First()) * 1000));
|
||||
|
||||
if (response.Headers.TryGetValues("X-RateLimit-Reset-After", out var resetAfter))
|
||||
ResetAfter = TimeSpan.FromSeconds(double.Parse(resetAfter!.First()));
|
||||
|
||||
if (response.Headers.TryGetValues("X-RateLimit-Bucket", out var bucket))
|
||||
Bucket = bucket.First();
|
||||
|
||||
if (response.Headers.TryGetValues("X-RateLimit-Global", out var global))
|
||||
Global = bool.Parse(global!.First());
|
||||
}
|
||||
|
||||
public bool Global { get; init; }
|
||||
public int? Limit { get; init; }
|
||||
public int? Remaining { get; init; }
|
||||
public DateTimeOffset? Reset { get; init; }
|
||||
public TimeSpan? ResetAfter { get; init; }
|
||||
public string? Bucket { get; init; }
|
||||
|
||||
public DateTimeOffset? ServerDate { get; init; }
|
||||
|
||||
public bool HasRatelimitInfo =>
|
||||
Limit != null && Remaining != null && Reset != null && ResetAfter != null && Bucket != null;
|
||||
}
|
||||
}
|
||||
86
Myriad/Rest/Ratelimit/Ratelimiter.cs
Normal file
86
Myriad/Rest/Ratelimit/Ratelimiter.cs
Normal file
@@ -0,0 +1,86 @@
|
||||
using System;
|
||||
|
||||
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;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user