dcord.api.ratelimit source code

1 /**
2   Utilities for tracking and managing rate limits.
3 */
4 module dcord.api.ratelimit;
5 
6 import std.conv,
7        std.math,
8        std.random,
9        core.time,
10        core.sync.mutex;
11 
12 import vibe.core.core;
13 import vibe.core.sync : LocalManualEvent, createManualEvent;
14 
15 import dcord.api.routes,
16        dcord.util.time;
17 
18 /// Return a random backoff duration (by default between 0.5 and 3 seconds)
19 Duration randomBackoff(int low=500, int high=3000) {
20   int milliseconds = uniform(low, high);
21   return milliseconds.msecs;
22 }
23 
24 /**
25   Stores the rate limit state for a given bucket.
26 */
27 struct RateLimitState {
28   int remaining;
29 
30   /// Time at which this rate limit resets
31   long resetTime;
32 
33   /// Returns true if this request is valid
34   bool willRateLimit() {
35     if (this.remaining - 1 < 0) {
36       if (getUnixTime() <= this.resetTime) {
37         return true;
38       }
39     }
40 
41     return false;
42   }
43 
44   /// Return the time that needs to be waited before another request can be made
45   Duration waitTime() {
46     return (this.resetTime - getUnixTime()).seconds + 500.msecs;
47   }
48 }
49 
50 /**
51   RateLimiter provides an interface for rate limiting HTTP Requests.
52 */
53 class RateLimiter {
54   LocalManualEvent[Bucket]     cooldowns;
55   RateLimitState[Bucket]  states;
56 
57   /// Cooldown a bucket for a given duration. Blocks ALL requests from completing.
58   void cooldown(Bucket bucket, Duration duration) {
59     if (bucket in this.cooldowns) {
60       this.cooldowns[bucket].wait();
61     } else {
62       this.cooldowns[bucket] = createManualEvent();
63       sleep(duration);
64       this.cooldowns[bucket].emit();
65       this.cooldowns.remove(bucket);
66     }
67   }
68 
69   /**
70     Check whether a request can be made for a bucket. If the bucket is on cooldown,
71     wait until the cooldown resets before returning.
72   */
73   bool check(Bucket bucket, Duration timeout) {
74     // If we're currently waiting for a cooldown, join the waiting club
75     if (bucket in this.cooldowns) {
76       if (this.cooldowns[bucket].wait(timeout, 0) != 0) {
77         return false;
78       }
79     }
80 
81     // If we don't have the bucket cached, return
82     if (bucket !in this.states) return true;
83 
84     // If this request will rate limit, wait until it won't anymore
85     if (this.states[bucket].willRateLimit()) {
86       this.cooldown(bucket, this.states[bucket].waitTime());
87     }
88 
89     return true;
90   }
91 
92   /// Update a given bucket with headers returned from a request.
93   void update(Bucket bucket, string remaining, string reset, string retryAfter) {
94     long resetSeconds = (reset.to!long);
95 
96     // If we have a retryAfter header, it may be more accurate
97     if (retryAfter != "") {
98       FloatingPointControl fpctrl;
99       fpctrl.rounding = FloatingPointControl.roundUp;
100       long retryAfterSeconds = rndtol(retryAfter.to!long / 1000.0);
101 
102       long nextRequestAt = getUnixTime() + retryAfterSeconds;
103       if (nextRequestAt > resetSeconds) {
104         resetSeconds = nextRequestAt;
105       }
106     }
107 
108     // Create a new RateLimitState if one doesn't exist
109     if (bucket !in this.states) {
110       this.states[bucket] = RateLimitState();
111     }
112 
113     // Save our remaining requests and reset seconds
114     this.states[bucket].remaining = remaining.to!int;
115     this.states[bucket].resetTime = resetSeconds;
116   }
117 }
118