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(getEpochTime() <= this.resetTime) return true; 
37 
38     return false;
39   }
40 
41   /// Return the time that needs to be waited before another request can be made
42   Duration waitTime() {
43     return (this.resetTime - getEpochTime()).seconds + 500.msecs;
44   }
45 }
46 
47 /**
48   RateLimiter provides an interface for rate limiting HTTP Requests.
49 */
50 class RateLimiter {
51   LocalManualEvent[Bucket] cooldowns;
52   RateLimitState[Bucket] states;
53 
54   /// Cooldown a bucket for a given duration. Blocks ALL requests from completing.
55   void cooldown(Bucket bucket, Duration duration) {
56     if(bucket in this.cooldowns) {
57       this.cooldowns[bucket].wait();
58     } else {
59       this.cooldowns[bucket] = createManualEvent();
60       sleep(duration);
61       this.cooldowns[bucket].emit();
62       this.cooldowns.remove(bucket);
63     }
64   }
65 
66   /**
67     Check whether a request can be made for a bucket. If the bucket is on cooldown,
68     wait until the cooldown resets before returning.
69   */
70   bool check(Bucket bucket, Duration timeout) {
71     // If we're currently waiting for a cooldown, join the waiting club
72     if (bucket in this.cooldowns) {
73       if (this.cooldowns[bucket].wait(timeout, 0) != 0) {
74         return false;
75       }
76     }
77 
78     // If we don't have the bucket cached, return
79     if (bucket !in this.states) return true;
80 
81     // If this request will rate limit, wait until it won't anymore
82     if(this.states[bucket].willRateLimit()) 
83       this.cooldown(bucket, this.states[bucket].waitTime());
84 
85     return true;
86   }
87 
88   /// Update a given bucket with headers returned from a request.
89   void update(Bucket bucket, string remaining, string reset, string retryAfter) {
90     long resetSeconds = (reset.to!long);
91 
92     // If we have a retryAfter header, it may be more accurate
93     if(retryAfter != "") {
94       FloatingPointControl fpctrl;
95       fpctrl.rounding = FloatingPointControl.roundUp;
96       long retryAfterSeconds = rndtol(retryAfter.to!long / 1000.0);
97       long nextRequestAt = getEpochTime() + retryAfterSeconds;
98       if(nextRequestAt > resetSeconds) resetSeconds = nextRequestAt;
99     }
100 
101     // Create a new RateLimitState if one doesn't exist
102     if(bucket !in this.states) 
103       this.states[bucket] = RateLimitState();
104     
105     // Save our remaining requests and reset seconds
106     this.states[bucket].remaining = remaining.to!int;
107     this.states[bucket].resetTime = resetSeconds;
108   }
109 }
110