diff --git a/src/Discord.Net.Rest/Net/Queue/RequestQueue.cs b/src/Discord.Net.Rest/Net/Queue/RequestQueue.cs index 4baf76433..a363615b1 100644 --- a/src/Discord.Net.Rest/Net/Queue/RequestQueue.cs +++ b/src/Discord.Net.Rest/Net/Queue/RequestQueue.cs @@ -14,7 +14,8 @@ namespace Discord.Net.Queue { public event Func RateLimitTriggered; - private readonly ConcurrentDictionary _buckets; + private readonly ConcurrentDictionary _bucketsByHash; + private readonly ConcurrentDictionary _bucketsById; private readonly SemaphoreSlim _tokenLock; private readonly CancellationTokenSource _cancelTokenSource; //Dispose token private CancellationTokenSource _clearToken; @@ -34,7 +35,8 @@ namespace Discord.Net.Queue _requestCancelToken = CancellationToken.None; _parentToken = CancellationToken.None; - _buckets = new ConcurrentDictionary(); + _bucketsByHash = new ConcurrentDictionary(); + _bucketsById = new ConcurrentDictionary(); _cleanupTask = RunCleanup(); } @@ -112,12 +114,24 @@ namespace Discord.Net.Queue private RequestBucket GetOrCreateBucket(string id, RestRequest request) { - return _buckets.GetOrAdd(id, x => new RequestBucket(this, request, x)); + object obj = _bucketsById.GetOrAdd(id, x => new RequestBucket(this, request, x)); + if (obj is string hash) + return _bucketsByHash.GetOrAdd(hash, x => new RequestBucket(this, request, x)); + return (RequestBucket)obj; } internal async Task RaiseRateLimitTriggered(string bucketId, RateLimitInfo? info) { await RateLimitTriggered(bucketId, info).ConfigureAwait(false); } + internal void UpdateBucketHash(string id, string discordHash) + { + if (_bucketsById.TryGetValue(id, out object obj) && obj is RequestBucket bucket) + { + string hash = discordHash + id.Split(new char[] { ' ' }, 2)[1]; //remove http method, using hash now + _bucketsByHash.GetOrAdd(hash, bucket); + _bucketsById.TryUpdate(id, hash, bucket); + } + } private async Task RunCleanup() { @@ -126,10 +140,21 @@ namespace Discord.Net.Queue while (!_cancelTokenSource.IsCancellationRequested) { var now = DateTimeOffset.UtcNow; - foreach (var bucket in _buckets.Select(x => x.Value)) + foreach (var bucket in _bucketsById.Where(x => x.Value is RequestBucket).Select(x => (RequestBucket)x.Value)) { if ((now - bucket.LastAttemptAt).TotalMinutes > 1.0) - _buckets.TryRemove(bucket.Id, out _); + _bucketsById.TryRemove(bucket.Id, out _); + } + foreach (var kvp in _bucketsByHash) + { + var kvpHash = kvp.Key; + var kvpBucket = kvp.Value; + if ((now - kvpBucket.LastAttemptAt).TotalMinutes > 1.0) + { + _bucketsByHash.TryRemove(kvpHash, out _); + foreach (var key in _bucketsById.Where(x => x.Value is string hash && hash == kvpHash).Select(x => x.Key)) + _bucketsById.TryRemove(key, out _); + } } await Task.Delay(60000, _cancelTokenSource.Token).ConfigureAwait(false); //Runs each minute } diff --git a/src/Discord.Net.Rest/Net/Queue/RequestQueueBucket.cs b/src/Discord.Net.Rest/Net/Queue/RequestQueueBucket.cs index d9d1e888f..46eee8d73 100644 --- a/src/Discord.Net.Rest/Net/Queue/RequestQueueBucket.cs +++ b/src/Discord.Net.Rest/Net/Queue/RequestQueueBucket.cs @@ -233,6 +233,9 @@ namespace Discord.Net.Queue #endif } + if (info.Bucket != null) + _queue.UpdateBucketHash(request.Options.BucketId, info.Bucket); + var now = DateTimeOffset.UtcNow.ToUnixTimeSeconds(); DateTimeOffset? resetTick = null; diff --git a/src/Discord.Net.Rest/Net/RateLimitInfo.cs b/src/Discord.Net.Rest/Net/RateLimitInfo.cs index 13e9e39a7..74ccf6ccb 100644 --- a/src/Discord.Net.Rest/Net/RateLimitInfo.cs +++ b/src/Discord.Net.Rest/Net/RateLimitInfo.cs @@ -12,6 +12,7 @@ namespace Discord.Net public int? RetryAfter { get; } public DateTimeOffset? Reset { get; } public TimeSpan? ResetAfter { get; } + public string Bucket { get; } public TimeSpan? Lag { get; } internal RateLimitInfo(Dictionary headers) @@ -28,6 +29,7 @@ namespace Discord.Net int.TryParse(temp, NumberStyles.None, CultureInfo.InvariantCulture, out var retryAfter) ? retryAfter : (int?)null; ResetAfter = headers.TryGetValue("X-RateLimit-Reset-After", out temp) && float.TryParse(temp, out var resetAfter) ? TimeSpan.FromMilliseconds((long)(resetAfter * 1000)) : (TimeSpan?)null; + Bucket = headers.TryGetValue("X-RateLimit-Bucket", out temp) ? temp : null; Lag = headers.TryGetValue("Date", out temp) && DateTimeOffset.TryParse(temp, CultureInfo.InvariantCulture, DateTimeStyles.None, out var date) ? DateTimeOffset.UtcNow - date : (TimeSpan?)null; }