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(
10    clippy::module_name_repetitions,
11    clippy::must_use_candidate,
12    clippy::unnecessary_wraps
13)]
14
15pub mod headers;
16pub mod in_memory;
17pub mod request;
18pub mod ticket;
19
20pub use self::{
21    headers::RatelimitHeaders,
22    in_memory::InMemoryRatelimiter,
23    request::{Method, Path},
24};
25
26use self::ticket::{TicketReceiver, TicketSender};
27use std::{
28    error::Error,
29    fmt::Debug,
30    future::Future,
31    pin::Pin,
32    time::{Duration, Instant},
33};
34
35/// A bucket containing ratelimiting information for a [`Path`].
36pub struct Bucket {
37    /// Total number of tickets allotted in a cycle.
38    limit: u64,
39    /// Number of tickets remaining.
40    remaining: u64,
41    /// Duration after [`Self::started_at`] time the bucket will refresh.
42    reset_after: Duration,
43    /// When the bucket's ratelimit refresh countdown started.
44    started_at: Option<Instant>,
45}
46
47impl Bucket {
48    /// Create a representation of a ratelimiter bucket.
49    ///
50    /// Buckets are returned by ratelimiters via [`Ratelimiter::bucket`] method.
51    /// Its primary use is for informational purposes, including information
52    /// such as the [number of remaining tickets][`Self::limit`] or determining
53    /// how much time remains
54    /// [until the bucket interval resets][`Self::time_remaining`].
55    #[must_use]
56    pub const fn new(
57        limit: u64,
58        remaining: u64,
59        reset_after: Duration,
60        started_at: Option<Instant>,
61    ) -> Self {
62        Self {
63            limit,
64            remaining,
65            reset_after,
66            started_at,
67        }
68    }
69
70    /// Total number of tickets allotted in a cycle.
71    #[must_use]
72    pub const fn limit(&self) -> u64 {
73        self.limit
74    }
75
76    /// Number of tickets remaining.
77    #[must_use]
78    pub const fn remaining(&self) -> u64 {
79        self.remaining
80    }
81
82    /// Duration after the [`Self::started_at`] time the bucket will
83    /// refresh.
84    #[must_use]
85    pub const fn reset_after(&self) -> Duration {
86        self.reset_after
87    }
88
89    /// When the bucket's ratelimit refresh countdown started.
90    #[must_use]
91    pub const fn started_at(&self) -> Option<Instant> {
92        self.started_at
93    }
94
95    /// How long until the bucket will refresh.
96    ///
97    /// May return `None` if the refresh timer has not been started yet or
98    /// the bucket has already refreshed.
99    #[must_use]
100    pub fn time_remaining(&self) -> Option<Duration> {
101        let reset_at = self.started_at? + self.reset_after;
102
103        reset_at.checked_duration_since(Instant::now())
104    }
105}
106
107/// A generic error type that implements [`Error`].
108pub type GenericError = Box<dyn Error + Send + Sync>;
109
110/// Future returned by [`Ratelimiter::bucket`].
111pub type GetBucketFuture =
112    Pin<Box<dyn Future<Output = Result<Option<Bucket>, GenericError>> + Send + 'static>>;
113
114/// Future returned by [`Ratelimiter::is_globally_locked`].
115pub type IsGloballyLockedFuture =
116    Pin<Box<dyn Future<Output = Result<bool, GenericError>> + Send + 'static>>;
117
118/// Future returned by [`Ratelimiter::has`].
119pub type HasBucketFuture =
120    Pin<Box<dyn Future<Output = Result<bool, GenericError>> + Send + 'static>>;
121
122/// Future returned by [`Ratelimiter::ticket`].
123pub type GetTicketFuture =
124    Pin<Box<dyn Future<Output = Result<TicketReceiver, GenericError>> + Send + 'static>>;
125
126/// Future returned by [`Ratelimiter::wait_for_ticket`].
127pub type WaitForTicketFuture =
128    Pin<Box<dyn Future<Output = Result<TicketSender, GenericError>> + Send + 'static>>;
129
130/// An implementation of a ratelimiter for the Discord REST API.
131///
132/// A default implementation can be found in [`InMemoryRatelimiter`].
133///
134/// All operations are asynchronous to allow for custom implementations to
135/// use different storage backends, for example databases.
136///
137/// Ratelimiters should keep track of two kids of ratelimits:
138/// * The global ratelimit status
139/// * [`Path`]-specific ratelimits
140///
141/// To do this, clients utilizing a ratelimiter will send back response
142/// ratelimit headers via a [`TicketSender`].
143///
144/// The ratelimiter itself will hand a [`TicketReceiver`] to the caller
145/// when a ticket is being requested.
146pub trait Ratelimiter: Debug + Send + Sync {
147    /// Retrieve the basic information of the bucket for a given path.
148    fn bucket(&self, path: &Path) -> GetBucketFuture;
149
150    /// Whether the ratelimiter is currently globally locked.
151    fn is_globally_locked(&self) -> IsGloballyLockedFuture;
152
153    /// Determine if the ratelimiter has a bucket for the given path.
154    fn has(&self, path: &Path) -> HasBucketFuture;
155
156    /// Retrieve a ticket to know when to send a request.
157    /// The provided future will be ready when a ticket in the bucket is
158    /// available. Tickets are ready in order of retrieval.
159    fn ticket(&self, path: Path) -> GetTicketFuture;
160
161    /// Retrieve a ticket to send a request.
162    /// Other than [`Self::ticket`], this method will return
163    /// a [`TicketSender`].
164    ///
165    /// This is identical to calling [`Self::ticket`] and then
166    /// awaiting the [`TicketReceiver`].
167    fn wait_for_ticket(&self, path: Path) -> WaitForTicketFuture {
168        let get_ticket = self.ticket(path);
169        Box::pin(async move {
170            match get_ticket.await {
171                Ok(rx) => rx.await.map_err(From::from),
172                Err(e) => Err(e),
173            }
174        })
175    }
176}