From 741f10d9b1715ea33e704c33615e6d88ad44cbb0 Mon Sep 17 00:00:00 2001
From: RogueException <rogueexception@gmail.com>
Date: Wed, 22 Jun 2016 20:06:58 -0300
Subject: [PATCH] Improved 429 handling

---
 src/Discord.Net/Net/Queue/RequestQueueBucket.cs | 12 ++++-----
 src/Discord.Net/Net/RateLimitException.cs       |  5 ++--
 src/Discord.Net/Net/Rest/DefaultRestClient.cs   | 35 +++++++++++++++----------
 3 files changed, 30 insertions(+), 22 deletions(-)

diff --git a/src/Discord.Net/Net/Queue/RequestQueueBucket.cs b/src/Discord.Net/Net/Queue/RequestQueueBucket.cs
index a61d71476..3c914315d 100644
--- a/src/Discord.Net/Net/Queue/RequestQueueBucket.cs
+++ b/src/Discord.Net/Net/Queue/RequestQueueBucket.cs
@@ -47,7 +47,7 @@ namespace Discord.Net.Queue
                     RequestQueueBucket bucket;
                     bool success = FindBucket(ex.BucketId, out bucket);
 
-                    await _queue.RaiseRateLimitTriggered(ex.BucketId, success ? bucket.Definition : (Bucket)null, ex.RetryAfterMilliseconds).ConfigureAwait(false);
+                    await _queue.RaiseRateLimitTriggered(ex.BucketId, success ? bucket.Definition : null, ex.RetryAfterMilliseconds).ConfigureAwait(false);
 
                     bucket.Pause(ex.RetryAfterMilliseconds);
                 }
@@ -67,6 +67,7 @@ namespace Discord.Net.Queue
                     //Get our 429 state
                     Task notifier;
                     int resumeTime;
+                    
                     lock (_pauseLock)
                     {
                         notifier = _resumeNotifier.Task;
@@ -133,14 +134,14 @@ namespace Discord.Net.Queue
                 {
                     _resumeNotifier = new TaskCompletionSource<byte>();
                     _pauseEndTick = unchecked(Environment.TickCount + milliseconds);
-                    QueueResumeAsync(milliseconds);
+                    QueueResumeAsync(_resumeNotifier, milliseconds);
                 }
             }
         }
-        private async Task QueueResumeAsync(int millis)
+        private async Task QueueResumeAsync(TaskCompletionSource<byte> resumeNotifier, int millis)
         {
             await Task.Delay(millis).ConfigureAwait(false);
-            _resumeNotifier.SetResult(0);
+            resumeNotifier.SetResult(0);
         }
 
         private async Task EnterAsync(int? endTick)
@@ -151,8 +152,7 @@ namespace Discord.Net.Queue
                 if (millis <= 0 || !await _semaphore.WaitAsync(millis).ConfigureAwait(false))
                     throw new TimeoutException();
             }
-            else
-                await _semaphore.WaitAsync().ConfigureAwait(false);
+            await _semaphore.WaitAsync().ConfigureAwait(false);
         }
         private async Task QueueExitAsync()
         {
diff --git a/src/Discord.Net/Net/RateLimitException.cs b/src/Discord.Net/Net/RateLimitException.cs
index 4fa4900e9..ff594155a 100644
--- a/src/Discord.Net/Net/RateLimitException.cs
+++ b/src/Discord.Net/Net/RateLimitException.cs
@@ -7,9 +7,10 @@ namespace Discord.Net
         public string BucketId { get; }
         public int RetryAfterMilliseconds { get; }
 
-        public HttpRateLimitException(int retryAfterMilliseconds)
-            : base((HttpStatusCode)429)
+        public HttpRateLimitException(string bucketId, int retryAfterMilliseconds, string reason)
+            : base((HttpStatusCode)429, reason)
         {
+            BucketId = bucketId;
             RetryAfterMilliseconds = retryAfterMilliseconds;
         }
     }
diff --git a/src/Discord.Net/Net/Rest/DefaultRestClient.cs b/src/Discord.Net/Net/Rest/DefaultRestClient.cs
index 43b405966..dd937d6ec 100644
--- a/src/Discord.Net/Net/Rest/DefaultRestClient.cs
+++ b/src/Discord.Net/Net/Rest/DefaultRestClient.cs
@@ -124,25 +124,32 @@ namespace Discord.Net.Rest
                 int statusCode = (int)response.StatusCode;
                 if (statusCode < 200 || statusCode >= 300) //2xx = Success
                 {
-                    if (statusCode == 429)
-                    {
-                        //TODO: Include bucket info
-                        int retryAfterMillis = int.Parse(response.Headers.GetValues("retry-after").First());
-                        throw new HttpRateLimitException(retryAfterMillis);
-                    }
-
                     string reason = null;
-                    try
+                    JToken content = null;
+                    if (response.Content.Headers.GetValues("content-type").FirstOrDefault() == "application/json")
                     {
-                        using (var stream = await response.Content.ReadAsStreamAsync().ConfigureAwait(false))
-                        using (var reader = new StreamReader(stream))
-                        using (var json = new JsonTextReader(reader))
+                        try
                         {
-                            reason = (_errorDeserializer.Deserialize(json) as JToken).Value<string>("message");
+                            using (var stream = await response.Content.ReadAsStreamAsync().ConfigureAwait(false))
+                            using (var reader = new StreamReader(stream))
+                            using (var json = new JsonTextReader(reader))
+                            {
+                                content = _errorDeserializer.Deserialize<JToken>(json);
+                                reason = content.Value<string>("message");
+                            }
                         }
+                        catch { } //Might have been HTML Should we check for content-type?
+                    }
+
+                    if (statusCode == 429 && content != null)
+                    {
+                        //TODO: Include bucket info
+                        string bucketId = content.Value<string>("bucket");
+                        int retryAfterMillis = content.Value<int>("retry_after");
+                        throw new HttpRateLimitException(bucketId, retryAfterMillis, reason);
                     }
-                    catch { } //Might have been HTML
-                    throw new HttpException(response.StatusCode, reason);
+                    else
+                        throw new HttpException(response.StatusCode, reason);
                 }
 
                 if (headerOnly)