From 68b3120aa87871d77a77d5522b65ad7c08ba0333 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Florian=20Spie=C3=9F?= <business@minnced.club>
Date: Fri, 18 Aug 2017 01:06:50 +0200
Subject: [PATCH] Added proper backoff for IDENTIFY rate limit handling

---
 .../jda/core/requests/WebSocketClient.java    | 21 ++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/src/main/java/net/dv8tion/jda/core/requests/WebSocketClient.java b/src/main/java/net/dv8tion/jda/core/requests/WebSocketClient.java
index ac3b22dff3..3176803b65 100644
--- a/src/main/java/net/dv8tion/jda/core/requests/WebSocketClient.java
+++ b/src/main/java/net/dv8tion/jda/core/requests/WebSocketClient.java
@@ -54,6 +54,7 @@ public class WebSocketClient extends WebSocketAdapter implements WebSocketListen
 {
     public static final SimpleLog LOG = SimpleLog.getLog("JDASocket");
     public static final int DISCORD_GATEWAY_VERSION = 6;
+    public static final int IDENTIFY_DELAY = 5;
 
     protected final JDAImpl api;
     protected final JDA.ShardInfo shardInfo;
@@ -90,6 +91,7 @@ public class WebSocketClient extends WebSocketAdapter implements WebSocketListen
 
     protected boolean firstInit = true;
     protected boolean processingReady = true;
+    protected boolean handleIdentifyRateLimit = false;
 
     public WebSocketClient(JDAImpl api)
     {
@@ -475,13 +477,24 @@ else if (closeCode != null)
 
     protected void reconnect()
     {
-        LOG.warn("Got disconnected from WebSocket (Internet?!)... Attempting to reconnect in " + reconnectTimeoutS + "s");
+        if (!handleIdentifyRateLimit)
+            LOG.warn("Got disconnected from WebSocket (Internet?!)... Attempting to reconnect in " + reconnectTimeoutS + "s");
         while(shouldReconnect)
         {
             try
             {
                 api.setStatus(JDA.Status.WAITING_TO_RECONNECT);
-                Thread.sleep(reconnectTimeoutS * 1000);
+                if (handleIdentifyRateLimit)
+                {
+                    handleIdentifyRateLimit = false;
+                    LOG.fatal("Encountered IDENTIFY (OP " + WebSocketCode.IDENTIFY + ") Rate Limit! " +
+                        "Waiting " + IDENTIFY_DELAY + " seconds before trying again!");
+                    Thread.sleep(IDENTIFY_DELAY * 1000);
+                }
+                else
+                {
+                    Thread.sleep(reconnectTimeoutS * 1000);
+                }
                 api.setStatus(JDA.Status.ATTEMPTING_TO_RECONNECT);
             }
             catch(InterruptedException ignored) {}
@@ -531,7 +544,7 @@ public void onTextMessage(WebSocket websocket, String message)
                 int closeCode = isResume ? 4000 : 1000;
                 if (isResume)
                     LOG.debug("Session can be recovered... Closing and sending new RESUME request");
-                else // this can also mean we got rate limited in IDENTIFY
+                else if (!handleIdentifyRateLimit) // this can also mean we got rate limited in IDENTIFY (no need to invalidate then)
                     invalidate();
 
                 close(closeCode);
@@ -621,6 +634,7 @@ protected void sendIdentify()
                     .put(shardInfo.getShardTotal()));
         }
         send(identify.toString(), true);
+        handleIdentifyRateLimit = true;
         sentAuthInfo = true;
     }
 
@@ -791,6 +805,7 @@ protected void handleEvent(JSONObject raw)
                 case "READY":
                     //LOG.debug(String.format("%s -> %s", type, content.toString())); already logged on trace level
                     processingReady = true;
+                    handleIdentifyRateLimit = false;
                     sessionId = content.getString("session_id");
                     if (!content.isNull("_trace"))
                         updateTraces(content.getJSONArray("_trace"), "READY", WebSocketCode.DISPATCH);