Skip to content

Commit e243060

Browse files
MinnDevelopmentMinnDevelopment
authored
Cleaned up BotRateLimiter for better handling of partial/missing headers (discord-jda#670)
1 parent 1a62381 commit e243060

File tree

3 files changed

+123
-65
lines changed

3 files changed

+123
-65
lines changed

‎src/main/java/net/dv8tion/jda/core/requests/RateLimiter.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,8 @@ public abstract class RateLimiter
3434
protected final Requester requester;
3535
protected final ScheduledThreadPoolExecutor pool;
3636
protected volatile boolean isShutdown = false;
37-
protected volatile ConcurrentHashMap<String, IBucket> buckets = new ConcurrentHashMap<>();
38-
protected volatile ConcurrentLinkedQueue<IBucket> submittedBuckets = new ConcurrentLinkedQueue<>();
37+
protected final ConcurrentHashMap<String, IBucket> buckets = new ConcurrentHashMap<>();
38+
protected final ConcurrentLinkedQueue<IBucket> submittedBuckets = new ConcurrentLinkedQueue<>();
3939

4040
protected RateLimiter(Requester requester, int poolSize)
4141
{

‎src/main/java/net/dv8tion/jda/core/requests/Route.java

Lines changed: 41 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -26,14 +26,15 @@
2626

2727
import static net.dv8tion.jda.core.requests.Method.*;
2828

29+
@SuppressWarnings("unused")
2930
public class Route
3031
{
3132
public static class Misc
3233
{
33-
public static final Route GET_VOICE_REGIONS = new Route(GET, "voice/regions");
34-
public static final Route GATEWAY = new Route(GET, "gateway");
35-
public static final Route GATEWAY_BOT = new Route(GET, "gateway/bot");
36-
public static final Route TRACK = new Route(POST, "track");
34+
public static final Route TRACK = new Route(POST, true, "track");
35+
public static final Route GET_VOICE_REGIONS = new Route(GET, true, "voice/regions");
36+
public static final Route GATEWAY = new Route(GET, true, "gateway");
37+
public static final Route GATEWAY_BOT = new Route(GET, true, "gateway/bot");
3738
}
3839

3940
public static class Applications
@@ -57,17 +58,15 @@ public static class Applications
5758
public static final Route GET_AUTHORIZED_APPLICATION = new Route(GET, "oauth2/tokens/{auth_id}");
5859
public static final Route DELETE_AUTHORIZED_APPLICATION = new Route(DELETE, "oauth2/tokens/{auth_id}");
5960
}
60-
61+
6162
public static class Self
6263
{
63-
public static final Route GET_SELF = new Route(GET, "users/@me");
64-
public static final Route MODIFY_SELF = new Route(PATCH, "users/@me");
65-
public static final Route GET_GUILDS = new Route(GET, "users/@me/guilds");
66-
public static final Route LEAVE_GUILD = new Route(DELETE, "users/@me/guilds/{guild_id}");
67-
public static final Route GET_PRIVATE_CHANNELS = new Route(GET, "users/@me/channels");
68-
public static final Route CREATE_PRIVATE_CHANNEL = new Route(POST, "users/@me/channels");
69-
70-
public static final Route GATEWAY_BOT = new Route(GET, "gateway/bot");
64+
public static final Route GET_SELF = new Route(GET, true, "users/@me");
65+
public static final Route MODIFY_SELF = new Route(PATCH, "users/@me");
66+
public static final Route GET_GUILDS = new Route(GET, "users/@me/guilds");
67+
public static final Route LEAVE_GUILD = new Route(DELETE, "users/@me/guilds/{guild_id}");
68+
public static final Route GET_PRIVATE_CHANNELS = new Route(GET, "users/@me/channels");
69+
public static final Route CREATE_PRIVATE_CHANNEL = new Route(POST, "users/@me/channels");
7170

7271
// Client only
7372
public static final Route USER_SETTINGS = new Route(GET, "users/@me/settings");
@@ -115,8 +114,8 @@ public static class Guilds
115114
public static final Route GET_GUILD_EMBED = new Route(GET, "guilds/{guild_id}/embed", "guild_id");
116115
public static final Route MODIFY_GUILD_EMBED = new Route(PATCH, "guilds/{guild_id}/embed", "guild_id");
117116
public static final Route GET_GUILD_EMOTES = new Route(GET, "guilds/{guild_id}/emojis", "guild_id");
118-
public static final Route GET_AUDIT_LOGS = new Route(GET, "guilds/{guild_id}/audit-logs", "guild_id");
119-
public static final Route GET_VOICE_REGIONS = new Route(GET, "guilds/{guild_id}/regions", "guild_id");
117+
public static final Route GET_AUDIT_LOGS = new Route(GET, true, "guilds/{guild_id}/audit-logs", "guild_id");
118+
public static final Route GET_VOICE_REGIONS = new Route(GET, true, "guilds/{guild_id}/regions", "guild_id");
120119

121120
public static final Route GET_INTEGRATIONS = new Route(GET, "guilds/{guild_id}/integrations", "guild_id");
122121
public static final Route CREATE_INTEGRATION = new Route(POST, "guilds/{guild_id}/integrations", "guild_id");
@@ -194,7 +193,6 @@ public static class Messages
194193
{
195194
public static final Route SEND_MESSAGE = new Route(POST, "channels/{channel_id}/messages", "channel_id");
196195
public static final Route EDIT_MESSAGE = new Route(PATCH, "channels/{channel_id}/messages/{message_id}", "channel_id");
197-
public static final Route DELETE_MESSAGE = new Route(DELETE, "channels/{channel_id}/messages/{message_id}", "channel_id");
198196
public static final Route GET_PINNED_MESSAGES = new Route(GET, "channels/{channel_id}/pins", "channel_id");
199197
public static final Route ADD_PINNED_MESSAGE = new Route(PUT, "channels/{channel_id}/pins/{message_id}", "channel_id");
200198
public static final Route REMOVE_PINNED_MESSAGE = new Route(DELETE, "channels/{channel_id}/pins/{message_id}", "channel_id");
@@ -206,23 +204,24 @@ public static class Messages
206204
public static final Route REMOVE_ALL_REACTIONS = new Route(DELETE, "channels/{channel_id}/messages/{message_id}/reactions", "channel_id");
207205
public static final Route GET_REACTION_USERS = new Route(GET, "channels/{channel_id}/messages/{message_id}/reactions/{reaction_code}", "channel_id");
208206

209-
public static final Route GET_MESSAGE_HISTORY = new Route(GET, "channels/{channel_id}/messages", "channel_id");
207+
public static final Route DELETE_MESSAGE = new Route(DELETE, true, "channels/{channel_id}/messages/{message_id}", "channel_id");
208+
public static final Route GET_MESSAGE_HISTORY = new Route(GET, true, "channels/{channel_id}/messages", "channel_id");
210209

211210
//Bot only
212-
public static final Route GET_MESSAGE = new Route(GET, "channels/{channel_id}/messages/{message_id}", "channel_id");
213-
public static final Route DELETE_MESSAGES = new Route(POST, "channels/{channel_id}/messages/bulk-delete", "channel_id");
211+
public static final Route GET_MESSAGE = new Route(GET, true, "channels/{channel_id}/messages/{message_id}", "channel_id");
212+
public static final Route DELETE_MESSAGES = new Route(POST, "channels/{channel_id}/messages/bulk-delete", "channel_id");
214213

215214
//Client only
216215
public static final Route ACK_MESSAGE = new Route(POST, "channels/{channel_id}/messages/{message_id}/ack");
217216
}
218217

219218
public static class Invites
220219
{
221-
public static final Route GET_INVITE = new Route(GET, "invites/{code}");
222-
public static final Route DELETE_INVITE = new Route(DELETE, "invites/{code}");
223-
public static final Route GET_GUILD_INVITES = new Route(GET, "guilds/{guild_id}/invites", "guild_id");
224-
public static final Route GET_CHANNEL_INVITES = new Route(GET, "channels/{channel_id}/invites", "channel_id");
225-
public static final Route CREATE_INVITE = new Route(POST, "channels/{channel_id}/invites", "channel_id");
220+
public static final Route GET_INVITE = new Route(GET, true, "invites/{code}");
221+
public static final Route GET_GUILD_INVITES = new Route(GET, true, "guilds/{guild_id}/invites", "guild_id");
222+
public static final Route GET_CHANNEL_INVITES = new Route(GET, true, "channels/{channel_id}/invites", "channel_id");
223+
public static final Route CREATE_INVITE = new Route(POST, "channels/{channel_id}/invites", "channel_id");
224+
public static final Route DELETE_INVITE = new Route(DELETE, "invites/{code}");
226225
}
227226

228227
public static class Custom
@@ -241,15 +240,27 @@ public static class Custom
241240
private final Method method;
242241
private final List<Integer> majorParamIndexes = new ArrayList<>();
243242
private final RateLimit ratelimit;
243+
private final boolean missingHeaders;
244244

245245
private Route(Method method, String route, String... majorParameters)
246246
{
247-
this(method, null, route, majorParameters);
247+
this(method, null, false, route, majorParameters);
248248
}
249249

250250
private Route(Method method, RateLimit rateLimit, String route, String... majorParameters)
251+
{
252+
this(method, rateLimit, false, route, majorParameters);
253+
}
254+
255+
private Route(Method method, boolean missingHeaders, String route, String... majorParameters)
256+
{
257+
this(method, null, missingHeaders, route, majorParameters);
258+
}
259+
260+
private Route(Method method, RateLimit rateLimit, boolean missingHeaders, String route, String... majorParameters)
251261
{
252262
this.method = method;
263+
this.missingHeaders = missingHeaders;
253264
this.ratelimit = rateLimit;
254265
this.route = route;
255266
this.paramCount = Helpers.countMatches(route, '{'); //All parameters start with {
@@ -303,6 +314,11 @@ public String getRoute()
303314
return route;
304315
}
305316

317+
public boolean isMissingHeaders()
318+
{
319+
return missingHeaders;
320+
}
321+
306322
public String getRatelimitRoute()
307323
{
308324
return ratelimitRoute;

‎src/main/java/net/dv8tion/jda/core/requests/ratelimit/BotRateLimiter.java

Lines changed: 80 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,14 @@
3636
import java.util.concurrent.ConcurrentLinkedQueue;
3737
import java.util.concurrent.RejectedExecutionException;
3838
import java.util.concurrent.TimeUnit;
39+
import java.util.function.IntConsumer;
40+
import java.util.function.LongConsumer;
3941

4042
public class BotRateLimiter extends RateLimiter
4143
{
44+
private static final String RESET_HEADER = "X-RateLimit-Reset";
45+
private static final String LIMIT_HEADER = "X-RateLimit-Limit";
46+
private static final String REMAINING_HEADER = "X-RateLimit-Remaining";
4247
protected volatile Long timeOffset = null;
4348

4449
public BotRateLimiter(Requester requester, int poolSize)
@@ -94,21 +99,21 @@ protected Long handleResponse(Route.CompiledRoute route, okhttp3.Response respon
9499
}
95100
}
96101
long retryAfter = Long.parseLong(retry);
97-
if (!Boolean.parseBoolean(global)) //Not global ratelimit
102+
if (Boolean.parseBoolean(global)) //global ratelimit
98103
{
99-
updateBucket(bucket, headers);
104+
//If it is global, lock down the threads.
105+
requester.getJDA().getSessionController().setGlobalRatelimit(getNow() + retryAfter);
100106
}
101107
else
102108
{
103-
//If it is global, lock down the threads.
104-
requester.getJDA().getSessionController().setGlobalRatelimit(getNow() + retryAfter);
109+
updateBucket(bucket, headers, retryAfter);
105110
}
106111

107112
return retryAfter;
108113
}
109114
else
110115
{
111-
updateBucket(bucket, headers);
116+
updateBucket(bucket, headers, -1);
112117
return null;
113118
}
114119
}
@@ -126,7 +131,8 @@ private Bucket getBucket(Route.CompiledRoute route)
126131
bucket = (Bucket) buckets.get(rateLimitRoute);
127132
if (bucket == null)
128133
{
129-
bucket = new Bucket(rateLimitRoute, route.getBaseRoute().getRatelimit());
134+
Route baseRoute = route.getBaseRoute();
135+
bucket = new Bucket(rateLimitRoute, baseRoute.getRatelimit(), baseRoute.isMissingHeaders());
130136
buckets.put(rateLimitRoute, bucket);
131137
}
132138
}
@@ -162,59 +168,85 @@ private void setTimeOffset(Headers headers)
162168
}
163169
}
164170

165-
private void updateBucket(Bucket bucket, Headers headers)
171+
private void updateBucket(Bucket bucket, Headers headers, long retryAfter)
166172
{
167-
try
173+
int headerCount = 0;
174+
if (retryAfter > 0)
168175
{
169-
if (bucket.hasRatelimit()) // Check if there's a hardcoded rate limit
170-
{
171-
bucket.resetTime = getNow() + bucket.getRatelimit().getResetTime();
172-
//routeUsageLimit provided by the ratelimit object already in the bucket.
173-
}
174-
else
175-
{
176-
bucket.resetTime = Long.parseLong(headers.get("X-RateLimit-Reset")) * 1000; //Seconds to milliseconds
177-
bucket.routeUsageLimit = Integer.parseInt(headers.get("X-RateLimit-Limit"));
178-
}
176+
bucket.resetTime = getNow() + retryAfter;
177+
bucket.routeUsageRemaining = 0;
178+
}
179179

180-
//Currently, we check the remaining amount even for hardcoded ratelimits just to further respect Discord
181-
// however, if there should ever be a case where Discord informs that the remaining is less than what
182-
// it actually is and we add a custom ratelimit to handle that, we will need to instead move this to the
183-
// above else statement and add a bucket.routeUsageRemaining-- decrement to the above if body.
184-
//An example of this statement needing to be moved would be if the custom ratelimit reset time interval is
185-
// equal to or greater than 1000ms, and the remaining count provided by discord is less than the ACTUAL
186-
// amount that their systems allow in such a way that isn't a bug.
187-
//The custom ratelimit system is primarily for ratelimits that can't be properly represented by Discord's
188-
// header system due to their headers only supporting accuracy to the second. The custom ratelimit system
189-
// allows for hardcoded ratelimits that allow accuracy to the millisecond which is important for some
190-
// ratelimits like Reactions which is 1/0.25s, but discord reports the ratelimit as 1/1s with headers.
191-
bucket.routeUsageRemaining = Integer.parseInt(headers.get("X-RateLimit-Remaining"));
180+
if (bucket.hasRatelimit()) // Check if there's a hardcoded rate limit
181+
{
182+
bucket.resetTime = getNow() + bucket.getRatelimit().getResetTime();
183+
//routeUsageLimit provided by the ratelimit object already in the bucket.
192184
}
193-
catch (NumberFormatException ex)
185+
else
194186
{
195-
if (!bucket.getRoute().equals("gateway")
196-
&& !bucket.getRoute().equals("users/@me"))
197-
{
198-
Requester.LOG.debug("Encountered issue with headers when updating a bucket\nRoute: {}\nHeaders: {}",
199-
bucket.getRoute(), headers);
200-
}
187+
headerCount += parseLong(headers.get(RESET_HEADER), bucket, (time, b) -> b.resetTime = time * 1000); //Seconds to milliseconds
188+
headerCount += parseInt(headers.get(LIMIT_HEADER), bucket, (limit, b) -> b.routeUsageLimit = limit);
189+
}
190+
191+
//Currently, we check the remaining amount even for hardcoded ratelimits just to further respect Discord
192+
// however, if there should ever be a case where Discord informs that the remaining is less than what
193+
// it actually is and we add a custom ratelimit to handle that, we will need to instead move this to the
194+
// above else statement and add a bucket.routeUsageRemaining-- decrement to the above if body.
195+
//An example of this statement needing to be moved would be if the custom ratelimit reset time interval is
196+
// equal to or greater than 1000ms, and the remaining count provided by discord is less than the ACTUAL
197+
// amount that their systems allow in such a way that isn't a bug.
198+
//The custom ratelimit system is primarily for ratelimits that can't be properly represented by Discord's
199+
// header system due to their headers only supporting accuracy to the second. The custom ratelimit system
200+
// allows for hardcoded ratelimits that allow accuracy to the millisecond which is important for some
201+
// ratelimits like Reactions which is 1/0.25s, but discord reports the ratelimit as 1/1s with headers.
202+
headerCount += parseInt(headers.get(REMAINING_HEADER), bucket, (remaining, b) -> b.routeUsageRemaining = remaining);
203+
if (!bucket.missingHeaders && headerCount < 3)
204+
{
205+
Requester.LOG.debug("Encountered issue with headers when updating a bucket\n" +
206+
"Route: {}\nHeaders: {}",
207+
bucket.getRoute(), headers);
208+
}
209+
}
201210

211+
private int parseInt(String input, Bucket bucket, IntObjectConsumer<? super Bucket> consumer)
212+
{
213+
try
214+
{
215+
int parsed = Integer.parseInt(input);
216+
consumer.accept(parsed, bucket);
217+
return 1;
202218
}
219+
catch (NumberFormatException ignored) {}
220+
return 0;
221+
}
222+
223+
private int parseLong(String input, Bucket bucket, LongObjectConsumer<? super Bucket> consumer)
224+
{
225+
try
226+
{
227+
long parsed = Long.parseLong(input);
228+
consumer.accept(parsed, bucket);
229+
return 1;
230+
}
231+
catch (NumberFormatException ignored) {}
232+
return 0;
203233
}
204234

205235
private class Bucket implements IBucket, Runnable
206236
{
207237
final String route;
238+
final boolean missingHeaders;
208239
final RateLimit rateLimit;
209240
final ConcurrentLinkedQueue<Request> requests = new ConcurrentLinkedQueue<>();
210241
volatile long resetTime = 0;
211242
volatile int routeUsageRemaining = 1; //These are default values to only allow 1 request until we have properly
212243
volatile int routeUsageLimit = 1; // ratelimit information.
213244

214-
public Bucket(String route, RateLimit rateLimit)
245+
public Bucket(String route, RateLimit rateLimit, boolean missingHeaders)
215246
{
216247
this.route = route;
217248
this.rateLimit = rateLimit;
249+
this.missingHeaders = missingHeaders;
218250
if (rateLimit != null)
219251
{
220252
this.routeUsageRemaining = rateLimit.getUsageLimit();
@@ -370,4 +402,14 @@ public Queue<Request> getRequests()
370402
return requests;
371403
}
372404
}
405+
406+
private interface LongObjectConsumer<T>
407+
{
408+
void accept(long n, T t);
409+
}
410+
411+
private interface IntObjectConsumer<T>
412+
{
413+
void accept(int n, T t);
414+
}
373415
}

0 commit comments

Comments
 (0)