From 9a55479349e37cb9a112e69fa31a1ee84b8c53c3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Florian=20Spie=C3=9F?= <business@minn.dev>
Date: Sat, 18 Nov 2023 19:54:54 +0100
Subject: [PATCH] Fix orphaned rate-limit buckets (#2585)

---
 .../requests/SequentialRestRateLimiter.java   | 30 +++++++++----------
 1 file changed, 14 insertions(+), 16 deletions(-)

diff --git a/src/main/java/net/dv8tion/jda/api/requests/SequentialRestRateLimiter.java b/src/main/java/net/dv8tion/jda/api/requests/SequentialRestRateLimiter.java
index 877a125f75..bafc368ab8 100644
--- a/src/main/java/net/dv8tion/jda/api/requests/SequentialRestRateLimiter.java
+++ b/src/main/java/net/dv8tion/jda/api/requests/SequentialRestRateLimiter.java
@@ -179,7 +179,7 @@ private void cleanup()
                 bucket.requests.removeIf(Work::isSkipped); // Remove cancelled requests
 
                 // Check if the bucket is empty
-                if (bucket.requests.isEmpty())
+                if (bucket.requests.isEmpty() && !rateLimitQueue.containsKey(bucket))
                 {
                     // remove uninit if requests are empty
                     if (bucket.isUninit())
@@ -231,8 +231,8 @@ private void runBucket(Bucket bucket)
             return;
         // Schedule a new bucket worker if no worker is running
         MiscUtil.locked(lock, () ->
-                rateLimitQueue.computeIfAbsent(bucket,
-                        (k) -> config.getPool().schedule(bucket, bucket.getRateLimit(), TimeUnit.MILLISECONDS)));
+            rateLimitQueue.computeIfAbsent(bucket,
+                k -> config.getPool().schedule(bucket, bucket.getRateLimit(), TimeUnit.MILLISECONDS)));
     }
 
     private long parseLong(String input)
@@ -252,9 +252,9 @@ private long getNow()
         return System.currentTimeMillis();
     }
 
-    private void updateBucket(Route.CompiledRoute route, Response response)
+    private Bucket updateBucket(Route.CompiledRoute route, Response response)
     {
-        MiscUtil.locked(lock, () ->
+        return MiscUtil.locked(lock, () ->
         {
             try
             {
@@ -302,7 +302,7 @@ else if (cloudflare)
                         boolean firstHit = hitRatelimit.add(baseRoute) && retryAfter < 60000;
                         // Update the bucket to the new information
                         bucket.remaining = 0;
-                        bucket.reset = getNow() + retryAfter;
+                        bucket.reset = now + retryAfter;
                         // don't log warning if we hit the rate limit for the first time, likely due to initialization of the bucket
                         // unless its a long retry-after delay (more than a minute)
                         if (firstHit)
@@ -310,6 +310,8 @@ else if (cloudflare)
                         else
                             log.warn("Encountered 429 on route {} with bucket {} Retry-After: {} ms Scope: {}", baseRoute, bucket.bucketId, retryAfter, scope);
                     }
+
+                    log.trace("Updated bucket {} to retry after {}", bucket.bucketId, bucket.reset - now);
                     return bucket;
                 }
 
@@ -367,7 +369,8 @@ public void enqueue(Work request)
 
         public void retry(Work request)
         {
-            requests.addFirst(request);
+            if (!moveRequest(request))
+                requests.addFirst(request);
         }
 
         public long getReset()
@@ -423,7 +426,7 @@ public Queue<Work> getRequests()
             return requests;
         }
 
-        protected Boolean moveRequest(Work request)
+        protected boolean moveRequest(Work request)
         {
             return MiscUtil.locked(lock, () ->
             {
@@ -433,9 +436,8 @@ protected Boolean moveRequest(Work request)
                 {
                     bucket.enqueue(request);
                     runBucket(bucket);
-                    return true;
                 }
-                return false;
+                return bucket != this;
             });
         }
 
@@ -480,12 +482,8 @@ public void run()
                 if (request.isSkipped())
                     continue;
 
-                // Check if a bucket has been discovered and initialized for this route
-                if (isUninit())
-                {
-                    boolean shouldSkip = moveRequest(request);
-                    if (shouldSkip) continue;
-                }
+                if (isUninit() && moveRequest(request))
+                    continue;
 
                 if (execute(request)) break;
             }