twilight_http_ratelimiting/
lib.rs

1#![doc = include_str!("../README.md")]
2#![warn(
3    clippy::missing_const_for_fn,
4    clippy::missing_docs_in_private_items,
5    clippy::pedantic,
6    missing_docs,
7    unsafe_code
8)]
9#![allow(clippy::module_name_repetitions, clippy::must_use_candidate)]
10
11mod actor;
12
13use std::{
14    future::Future,
15    hash::{Hash as _, Hasher},
16    num::NonZero,
17    pin::Pin,
18    task::{Context, Poll},
19    time::{Duration, Instant},
20};
21use tokio::sync::{mpsc, oneshot};
22
23/// Duration from the first globally limited request until the remaining count
24/// resets to the global limit count.
25pub const GLOBAL_LIMIT_PERIOD: Duration = Duration::from_secs(1);
26
27/// User actionable description that the actor panicked.
28const ACTOR_PANIC_MESSAGE: &str =
29    "actor task panicked: report its panic message to the maintainers";
30
31/// HTTP request [method].
32///
33/// [method]: https://developer.mozilla.org/en-US/docs/Web/HTTP/Methods
34#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
35#[non_exhaustive]
36pub enum Method {
37    /// Delete a resource.
38    Delete,
39    /// Retrieve a resource.
40    Get,
41    /// Update a resource.
42    Patch,
43    /// Create a resource.
44    Post,
45    /// Replace a resource.
46    Put,
47}
48
49impl Method {
50    /// Name of the method.
51    pub const fn name(self) -> &'static str {
52        match self {
53            Method::Delete => "DELETE",
54            Method::Get => "GET",
55            Method::Patch => "PATCH",
56            Method::Post => "POST",
57            Method::Put => "PUT",
58        }
59    }
60}
61
62/// Rate limited endpoint.
63///
64/// The rate limiter dynamically supports new or unknown API paths, but is consequently unable to
65/// catch invalid arguments. Invalidly structured endpoints may be permitted at an improper time.
66///
67/// # Example
68///
69/// ```no_run
70/// # let rt = tokio::runtime::Builder::new_current_thread()
71/// #     .enable_time()
72/// #     .build()
73/// #     .unwrap();
74/// # rt.block_on(async {
75/// # let rate_limiter = twilight_http_ratelimiting::RateLimiter::default();
76/// use twilight_http_ratelimiting::{Endpoint, Method};
77///
78/// let url = "https://discord.com/api/v10/guilds/745809834183753828/audit-logs?limit=10";
79/// let endpoint = Endpoint {
80///     method: Method::Get,
81///     path: String::from("guilds/745809834183753828/audit-logs"),
82/// };
83/// let permit = rate_limiter.acquire(endpoint).await;
84/// let headers = unimplemented!("GET {url}");
85/// permit.complete(headers);
86/// # });
87/// ```
88#[derive(Clone, Debug, Eq, Hash, PartialEq)]
89pub struct Endpoint {
90    /// Method of the endpoint.
91    pub method: Method,
92    /// API path of the endpoint.
93    ///
94    /// Should not start with a slash (`/`) or include query parameters (`?key=value`).
95    pub path: String,
96}
97
98impl Endpoint {
99    /// Whether the endpoint is properly structured.
100    pub(crate) fn is_valid(&self) -> bool {
101        !self.path.as_bytes().starts_with(b"/") && !self.path.as_bytes().contains(&b'?')
102    }
103
104    /// Whether the endpoint is an interaction.
105    pub(crate) fn is_interaction(&self) -> bool {
106        self.path.as_bytes().starts_with(b"webhooks")
107            || self.path.as_bytes().starts_with(b"interactions")
108    }
109
110    /// Feeds the top-level resources of this endpoint into the given [`Hasher`].
111    ///
112    /// Top-level resources represent the bucket namespace in which they are unique.
113    ///
114    /// Top-level resources are currently:
115    /// - `channels/<channel_id>`
116    /// - `guilds/<guild_id>`
117    /// - `webhooks/<webhook_id>`
118    /// - `webhooks/<webhook_id>/<webhook_token>`
119    pub(crate) fn hash_resources(&self, state: &mut impl Hasher) {
120        let mut segments = self.path.as_bytes().split(|&s| s == b'/');
121        match segments.next().unwrap_or_default() {
122            b"channels" => {
123                if let Some(s) = segments.next() {
124                    "channels".hash(state);
125                    s.hash(state);
126                }
127            }
128            b"guilds" => {
129                if let Some(s) = segments.next() {
130                    "guilds".hash(state);
131                    s.hash(state);
132                }
133            }
134            b"webhooks" => {
135                if let Some(s) = segments.next() {
136                    "webhooks".hash(state);
137                    s.hash(state);
138                }
139                if let Some(s) = segments.next() {
140                    s.hash(state);
141                }
142            }
143            _ => {}
144        }
145    }
146}
147
148/// Parsed user response rate limit headers.
149///
150/// A `limit` of zero marks the [`Bucket`] as exhausted until `reset_at` elapses.
151///
152/// # Global limits
153///
154/// Please open an issue if the [`RateLimiter`] exceeded the global limit.
155///
156/// # Shared limits
157///
158/// You may preemptively exhaust the bucket until `Reset-After` by completing
159/// the [`Permit`] with [`RateLimitHeaders::shared`], but are not required to
160/// since these limits do not count towards the invalid request limit.
161#[derive(Clone, Debug, Eq, Hash, PartialEq)]
162pub struct RateLimitHeaders {
163    /// Bucket identifier.
164    pub bucket: Vec<u8>,
165    /// Total number of requests until the bucket becomes exhausted.
166    pub limit: u16,
167    /// Number of remaining requests until the bucket becomes exhausted.
168    pub remaining: u16,
169    /// Time at which the bucket resets.
170    pub reset_at: Instant,
171}
172
173impl RateLimitHeaders {
174    /// Lowercased name for the bucket header.
175    pub const BUCKET: &'static str = "x-ratelimit-bucket";
176
177    /// Lowercased name for the limit header.
178    pub const LIMIT: &'static str = "x-ratelimit-limit";
179
180    /// Lowercased name for the remaining header.
181    pub const REMAINING: &'static str = "x-ratelimit-remaining";
182
183    /// Lowercased name for the reset-after header.
184    pub const RESET_AFTER: &'static str = "x-ratelimit-reset-after";
185
186    /// Lowercased name for the scope header.
187    pub const SCOPE: &'static str = "x-ratelimit-scope";
188
189    /// Emulates a shared resource limit as a user limit by setting `limit` and
190    /// `remaining` to zero.
191    pub fn shared(bucket: Vec<u8>, retry_after: u16) -> Self {
192        Self {
193            bucket,
194            limit: 0,
195            remaining: 0,
196            reset_at: Instant::now() + Duration::from_secs(retry_after.into()),
197        }
198    }
199}
200
201/// Permit to send a Discord HTTP API request to the acquired endpoint.
202#[derive(Debug)]
203#[must_use = "dropping the permit immediately cancels itself"]
204pub struct Permit(oneshot::Sender<Option<RateLimitHeaders>>);
205
206impl Permit {
207    /// Update the [`RateLimiter`] based on the response headers.
208    ///
209    /// Non-completed permits are regarded as cancelled, so only call this
210    /// on receiving a response.
211    #[allow(clippy::missing_panics_doc)]
212    pub fn complete(self, headers: Option<RateLimitHeaders>) {
213        assert!(self.0.send(headers).is_ok(), "{ACTOR_PANIC_MESSAGE}");
214    }
215}
216
217/// Future that completes when a permit is ready.
218#[derive(Debug)]
219#[must_use = "futures do nothing unless you `.await` or poll them"]
220pub struct PermitFuture(oneshot::Receiver<oneshot::Sender<Option<RateLimitHeaders>>>);
221
222impl Future for PermitFuture {
223    type Output = Permit;
224
225    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
226        #[allow(clippy::match_wild_err_arm)]
227        Pin::new(&mut self.0).poll(cx).map(|r| match r {
228            Ok(sender) => Permit(sender),
229            Err(_) => panic!("{ACTOR_PANIC_MESSAGE}"),
230        })
231    }
232}
233
234/// Future that completes when a permit is ready or cancelled.
235#[derive(Debug)]
236#[must_use = "futures do nothing unless you `.await` or poll them"]
237pub struct MaybePermitFuture(oneshot::Receiver<oneshot::Sender<Option<RateLimitHeaders>>>);
238
239impl Future for MaybePermitFuture {
240    type Output = Option<Permit>;
241
242    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
243        Pin::new(&mut self.0).poll(cx).map(|r| r.ok().map(Permit))
244    }
245}
246
247/// Rate limit information for one or more paths from previous
248/// [`RateLimitHeaders`].
249#[non_exhaustive]
250#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
251pub struct Bucket {
252    /// Total number of permits until the bucket becomes exhausted.
253    pub limit: u16,
254    /// Number of remaining permits until the bucket becomes exhausted.
255    pub remaining: u16,
256    /// Time at which the bucket resets.
257    pub reset_at: Instant,
258}
259
260/// Actor run closure pre-enqueue for early [`MaybePermitFuture`] cancellation.
261type Predicate = Box<dyn FnOnce(Option<Bucket>) -> bool + Send>;
262
263/// Discord HTTP client API rate limiter.
264///
265/// The [`RateLimiter`] runs an associated actor task to concurrently handle permit
266/// requests and responses.
267///
268/// Cloning a [`RateLimiter`] increments just the amount of senders for the actor.
269/// The actor completes when there are no senders and non-completed permits left.
270#[derive(Clone, Debug)]
271pub struct RateLimiter {
272    /// Actor message sender.
273    tx: mpsc::UnboundedSender<(actor::Message, Option<Predicate>)>,
274}
275
276impl RateLimiter {
277    /// Create a new [`RateLimiter`] with a custom global limit.
278    ///
279    /// Note that the global limit may be practically disabled by passing `u16::MAX`.
280    ///
281    /// # Panics
282    ///
283    /// Panics if the global limit is 0.
284    #[track_caller]
285    pub fn new(global_limit: u16) -> Self {
286        let global_limit = NonZero::new(global_limit).expect("global limit is non-zero");
287        let (tx, rx) = mpsc::unbounded_channel();
288        tokio::spawn(actor::runner(global_limit, rx));
289
290        Self { tx }
291    }
292
293    /// Await a single permit for this endpoint.
294    ///
295    /// Permits are queued per endpoint in the order they were requested.
296    pub fn acquire(&self, endpoint: Endpoint) -> PermitFuture {
297        PermitFuture(self.acquire_inner(endpoint, None))
298    }
299
300    /// Await a single permit for this endpoint, but only if the predicate evaluates
301    /// to `true`.
302    ///
303    /// Permits are queued per endpoint in the order they were requested.
304    ///
305    /// Note that the predicate is asynchronously called in the actor task.
306    ///
307    /// # Example
308    ///
309    /// ```no_run
310    /// # let rt = tokio::runtime::Builder::new_current_thread()
311    /// #     .enable_time()
312    /// #     .build()
313    /// #     .unwrap();
314    /// # rt.block_on(async {
315    /// # let rate_limiter = twilight_http_ratelimiting::RateLimiter::default();
316    /// use twilight_http_ratelimiting::{Endpoint, Method};
317    ///
318    /// let endpoint = Endpoint {
319    ///     method: Method::Get,
320    ///     path: String::from("applications/@me"),
321    /// };
322    /// if let Some(permit) = rate_limiter
323    ///     .acquire_if(endpoint, |b| b.is_none_or(|b| b.remaining > 10))
324    ///     .await
325    /// {
326    ///     let headers = unimplemented!("GET /applications/@me");
327    ///     permit.complete(headers);
328    /// }
329    /// # });
330    /// ```
331    pub fn acquire_if<P>(&self, endpoint: Endpoint, predicate: P) -> MaybePermitFuture
332    where
333        P: FnOnce(Option<Bucket>) -> bool + Send + 'static,
334    {
335        MaybePermitFuture(self.acquire_inner(endpoint, Some(Box::new(predicate))))
336    }
337
338    /// Requests a permit for this endpoint.
339    #[allow(clippy::missing_panics_doc)]
340    fn acquire_inner(
341        &self,
342        endpoint: Endpoint,
343        predicate: Option<Predicate>,
344    ) -> oneshot::Receiver<oneshot::Sender<Option<RateLimitHeaders>>> {
345        let (notifier, rx) = oneshot::channel();
346        let message = actor::Message { endpoint, notifier };
347
348        let res = self.tx.send((message, predicate));
349        assert!(res.is_ok(), "{ACTOR_PANIC_MESSAGE}");
350
351        rx
352    }
353
354    /// Retrieve the [`Bucket`] for this endpoint.
355    ///
356    /// The bucket is internally retrieved via [`acquire_if`][Self::acquire_if].
357    #[allow(clippy::missing_panics_doc)]
358    pub async fn bucket(&self, endpoint: Endpoint) -> Option<Bucket> {
359        let (tx, rx) = oneshot::channel();
360        self.acquire_if(endpoint, |bucket| {
361            // Ignore cancellation error.
362            _ = tx.send(bucket);
363            false
364        })
365        .await;
366
367        #[allow(clippy::match_wild_err_arm)]
368        match rx.await {
369            Ok(bucket) => bucket,
370            Err(_) => panic!("{ACTOR_PANIC_MESSAGE}"),
371        }
372    }
373}
374
375impl Default for RateLimiter {
376    /// Create a new [`RateLimiter`] with Discord's default global limit.
377    ///
378    /// Currently this is `50`.
379    fn default() -> Self {
380        Self::new(50)
381    }
382}
383
384#[cfg(test)]
385mod tests {
386    use super::{
387        Bucket, Endpoint, MaybePermitFuture, Method, Permit, PermitFuture, RateLimitHeaders,
388        RateLimiter,
389    };
390    use static_assertions::assert_impl_all;
391    use std::{
392        fmt::Debug,
393        future::Future,
394        hash::{DefaultHasher, Hash, Hasher as _},
395        time::{Duration, Instant},
396    };
397    use tokio::task;
398
399    assert_impl_all!(Bucket: Clone, Copy, Debug, Eq, Hash, PartialEq, Send, Sync);
400    assert_impl_all!(Endpoint: Clone, Debug, Eq, Hash, PartialEq, Send, Sync);
401    assert_impl_all!(MaybePermitFuture: Debug, Future<Output = Option<Permit>>);
402    assert_impl_all!(Method: Clone, Copy, Debug, Eq, PartialEq);
403    assert_impl_all!(Permit: Debug, Send, Sync);
404    assert_impl_all!(PermitFuture: Debug, Future<Output = Permit>);
405    assert_impl_all!(RateLimitHeaders: Clone, Debug, Eq, Hash, PartialEq, Send, Sync);
406    assert_impl_all!(RateLimiter: Clone, Debug, Default, Send, Sync);
407
408    const ENDPOINT: fn() -> Endpoint = || Endpoint {
409        method: Method::Get,
410        path: String::from("applications/@me"),
411    };
412
413    #[tokio::test]
414    async fn acquire_if() {
415        let rate_limiter = RateLimiter::default();
416
417        assert!(
418            rate_limiter
419                .acquire_if(ENDPOINT(), |_| false)
420                .await
421                .is_none()
422        );
423        assert!(
424            rate_limiter
425                .acquire_if(ENDPOINT(), |_| true)
426                .await
427                .is_some()
428        );
429    }
430
431    #[tokio::test]
432    async fn bucket() {
433        let rate_limiter = RateLimiter::default();
434
435        let limit = 2;
436        let remaining = 1;
437        let reset_at = Instant::now() + Duration::from_secs(1);
438        let headers = RateLimitHeaders {
439            bucket: vec![1, 2, 3],
440            limit,
441            remaining,
442            reset_at,
443        };
444
445        rate_limiter
446            .acquire(ENDPOINT())
447            .await
448            .complete(Some(headers));
449        task::yield_now().await;
450
451        let bucket = rate_limiter.bucket(ENDPOINT()).await.unwrap();
452        assert_eq!(bucket.limit, limit);
453        assert_eq!(bucket.remaining, remaining);
454        assert!(
455            bucket.reset_at.saturating_duration_since(reset_at) < Duration::from_millis(1)
456                && reset_at.saturating_duration_since(bucket.reset_at) < Duration::from_millis(1)
457        );
458    }
459
460    fn with_hasher(f: impl FnOnce(&mut DefaultHasher)) -> u64 {
461        let mut hasher = DefaultHasher::new();
462        f(&mut hasher);
463        hasher.finish()
464    }
465
466    #[test]
467    fn endpoint() {
468        let invalid = Endpoint {
469            method: Method::Get,
470            path: String::from("/guilds/745809834183753828/audit-logs?limit=10"),
471        };
472        let delete_webhook = Endpoint {
473            method: Method::Delete,
474            path: String::from("webhooks/1"),
475        };
476        let interaction_response = Endpoint {
477            method: Method::Post,
478            path: String::from("interactions/1/abc/callback"),
479        };
480
481        assert!(!invalid.is_valid());
482        assert!(delete_webhook.is_valid());
483        assert!(interaction_response.is_valid());
484
485        assert!(delete_webhook.is_interaction());
486        assert!(interaction_response.is_interaction());
487
488        assert_eq!(
489            with_hasher(|state| invalid.hash_resources(state)),
490            with_hasher(|_| {})
491        );
492        assert_eq!(
493            with_hasher(|state| delete_webhook.hash_resources(state)),
494            with_hasher(|state| {
495                "webhooks".hash(state);
496                b"1".hash(state);
497            })
498        );
499        assert_eq!(
500            with_hasher(|state| interaction_response.hash_resources(state)),
501            with_hasher(|_| {})
502        );
503    }
504
505    #[test]
506    fn method_conversions() {
507        assert_eq!("DELETE", Method::Delete.name());
508        assert_eq!("GET", Method::Get.name());
509        assert_eq!("PATCH", Method::Patch.name());
510        assert_eq!("POST", Method::Post.name());
511        assert_eq!("PUT", Method::Put.name());
512    }
513}
twilight_http_ratelimiting/lib.rs

twilight_http_ratelimiting/
lib.rs
