twilight-http-ratelimiting 0.17.1

#![doc = include_str!("../README.md")]
#![warn(
    clippy::missing_const_for_fn,
    clippy::missing_docs_in_private_items,
    clippy::pedantic,
    missing_docs,
    unsafe_code
)]
#![allow(clippy::module_name_repetitions, clippy::must_use_candidate)]

mod actor;

use std::{
    future::Future,
    hash::{Hash as _, Hasher},
    pin::Pin,
    task::{Context, Poll},
    time::{Duration, Instant},
};
use tokio::sync::{mpsc, oneshot};

/// Duration from the first globally limited request until the remaining count
/// resets to the global limit count.
pub const GLOBAL_LIMIT_PERIOD: Duration = Duration::from_secs(1);

/// User actionable description that the actor panicked.
const ACTOR_PANIC_MESSAGE: &str =
    "actor task panicked: report its panic message to the maintainers";

/// HTTP request [method].
///
/// [method]: https://developer.mozilla.org/en-US/docs/Web/HTTP/Methods
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
#[non_exhaustive]
pub enum Method {
    /// Delete a resource.
    Delete,
    /// Retrieve a resource.
    Get,
    /// Update a resource.
    Patch,
    /// Create a resource.
    Post,
    /// Replace a resource.
    Put,
}

impl Method {
    /// Name of the method.
    pub const fn name(self) -> &'static str {
        match self {
            Method::Delete => "DELETE",
            Method::Get => "GET",
            Method::Patch => "PATCH",
            Method::Post => "POST",
            Method::Put => "PUT",
        }
    }
}

/// Rate limited endpoint.
///
/// The rate limiter dynamically supports new or unknown API paths, but is consequently unable to
/// catch invalid arguments. Invalidly structured endpoints may be permitted at an improper time.
///
/// # Example
///
/// ```no_run
/// # let rt = tokio::runtime::Builder::new_current_thread()
/// #     .enable_time()
/// #     .build()
/// #     .unwrap();
/// # rt.block_on(async {
/// # let rate_limiter = twilight_http_ratelimiting::RateLimiter::default();
/// use twilight_http_ratelimiting::{Endpoint, Method};
///
/// let url = "https://discord.com/api/v10/guilds/745809834183753828/audit-logs?limit=10";
/// let endpoint = Endpoint {
///     method: Method::Get,
///     path: String::from("guilds/745809834183753828/audit-logs"),
/// };
/// let permit = rate_limiter.acquire(endpoint).await;
/// let headers = unimplemented!("GET {url}");
/// permit.complete(headers);
/// # });
/// ```
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct Endpoint {
    /// Method of the endpoint.
    pub method: Method,
    /// API path of the endpoint.
    ///
    /// Should not start with a slash (`/`) or include query parameters (`?key=value`).
    pub path: String,
}

impl Endpoint {
    /// Whether the endpoint is properly structured.
    pub(crate) fn is_valid(&self) -> bool {
        !self.path.as_bytes().starts_with(b"/") && !self.path.as_bytes().contains(&b'?')
    }

    /// Whether the endpoint is an interaction.
    pub(crate) fn is_interaction(&self) -> bool {
        self.path.as_bytes().starts_with(b"webhooks")
            || self.path.as_bytes().starts_with(b"interactions")
    }

    /// Feeds the top-level resources of this endpoint into the given [`Hasher`].
    ///
    /// Top-level resources represent the bucket namespace in which they are unique.
    ///
    /// Top-level resources are currently:
    /// - `channels/<channel_id>`
    /// - `guilds/<guild_id>`
    /// - `webhooks/<webhook_id>`
    /// - `webhooks/<webhook_id>/<webhook_token>`
    pub(crate) fn hash_resources(&self, state: &mut impl Hasher) {
        let mut segments = self.path.as_bytes().split(|&s| s == b'/');
        match segments.next().unwrap_or_default() {
            b"channels" => {
                if let Some(s) = segments.next() {
                    "channels".hash(state);
                    s.hash(state);
                }
            }
            b"guilds" => {
                if let Some(s) = segments.next() {
                    "guilds".hash(state);
                    s.hash(state);
                }
            }
            b"webhooks" => {
                if let Some(s) = segments.next() {
                    "webhooks".hash(state);
                    s.hash(state);
                }
                if let Some(s) = segments.next() {
                    s.hash(state);
                }
            }
            _ => {}
        }
    }
}

/// Parsed user response rate limit headers.
///
/// A `limit` of zero marks the [`Bucket`] as exhausted until `reset_at` elapses.
///
/// # Global limits
///
/// Please open an issue if the [`RateLimiter`] exceeded the global limit.
///
/// # Shared limits
///
/// You may preemptively exhaust the bucket until `Reset-After` by completing
/// the [`Permit`] with [`RateLimitHeaders::shared`], but are not required to
/// since these limits do not count towards the invalid request limit.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct RateLimitHeaders {
    /// Bucket identifier.
    pub bucket: Vec<u8>,
    /// Total number of requests until the bucket becomes exhausted.
    pub limit: u16,
    /// Number of remaining requests until the bucket becomes exhausted.
    pub remaining: u16,
    /// Time at which the bucket resets.
    pub reset_at: Instant,
}

impl RateLimitHeaders {
    /// Lowercased name for the bucket header.
    pub const BUCKET: &'static str = "x-ratelimit-bucket";

    /// Lowercased name for the limit header.
    pub const LIMIT: &'static str = "x-ratelimit-limit";

    /// Lowercased name for the remaining header.
    pub const REMAINING: &'static str = "x-ratelimit-remaining";

    /// Lowercased name for the reset-after header.
    pub const RESET_AFTER: &'static str = "x-ratelimit-reset-after";

    /// Lowercased name for the scope header.
    pub const SCOPE: &'static str = "x-ratelimit-scope";

    /// Emulates a shared resource limit as a user limit by setting `limit` and
    /// `remaining` to zero.
    pub fn shared(bucket: Vec<u8>, retry_after: u16) -> Self {
        Self {
            bucket,
            limit: 0,
            remaining: 0,
            reset_at: Instant::now() + Duration::from_secs(retry_after.into()),
        }
    }
}

/// Permit to send a Discord HTTP API request to the acquired endpoint.
#[derive(Debug)]
#[must_use = "dropping the permit immediately cancels itself"]
pub struct Permit(oneshot::Sender<Option<RateLimitHeaders>>);

impl Permit {
    /// Update the [`RateLimiter`] based on the response headers.
    ///
    /// Non-completed permits are regarded as cancelled, so only call this
    /// on receiving a response.
    #[allow(clippy::missing_panics_doc)]
    pub fn complete(self, headers: Option<RateLimitHeaders>) {
        assert!(self.0.send(headers).is_ok(), "{ACTOR_PANIC_MESSAGE}");
    }
}

/// Future that completes when a permit is ready.
#[derive(Debug)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct PermitFuture(oneshot::Receiver<oneshot::Sender<Option<RateLimitHeaders>>>);

impl Future for PermitFuture {
    type Output = Permit;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        #[allow(clippy::match_wild_err_arm)]
        Pin::new(&mut self.0).poll(cx).map(|r| match r {
            Ok(sender) => Permit(sender),
            Err(_) => panic!("{ACTOR_PANIC_MESSAGE}"),
        })
    }
}

/// Future that completes when a permit is ready or cancelled.
#[derive(Debug)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct MaybePermitFuture(oneshot::Receiver<oneshot::Sender<Option<RateLimitHeaders>>>);

impl Future for MaybePermitFuture {
    type Output = Option<Permit>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        Pin::new(&mut self.0).poll(cx).map(|r| r.ok().map(Permit))
    }
}

/// Rate limit information for one or more paths from previous
/// [`RateLimitHeaders`].
#[non_exhaustive]
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Bucket {
    /// Total number of permits until the bucket becomes exhausted.
    pub limit: u16,
    /// Number of remaining permits until the bucket becomes exhausted.
    pub remaining: u16,
    /// Time at which the bucket resets.
    pub reset_at: Instant,
}

/// Actor run closure pre-enqueue for early [`MaybePermitFuture`] cancellation.
type Predicate = Box<dyn FnOnce(Option<Bucket>) -> bool + Send>;

/// Discord HTTP client API rate limiter.
///
/// The [`RateLimiter`] runs an associated actor task to concurrently handle permit
/// requests and responses.
///
/// Cloning a [`RateLimiter`] increments just the amount of senders for the actor.
/// The actor completes when there are no senders and non-completed permits left.
#[derive(Clone, Debug)]
pub struct RateLimiter {
    /// Actor message sender.
    tx: mpsc::UnboundedSender<(actor::Message, Option<Predicate>)>,
}

impl RateLimiter {
    /// Create a new [`RateLimiter`] with a custom global limit.
    pub fn new(global_limit: u16) -> Self {
        let (tx, rx) = mpsc::unbounded_channel();
        tokio::spawn(actor::runner(global_limit, rx));

        Self { tx }
    }

    /// Await a single permit for this endpoint.
    ///
    /// Permits are queued per endpoint in the order they were requested.
    #[allow(clippy::missing_panics_doc)]
    pub fn acquire(&self, endpoint: Endpoint) -> PermitFuture {
        let (notifier, rx) = oneshot::channel();
        let message = actor::Message { endpoint, notifier };
        assert!(
            self.tx.send((message, None)).is_ok(),
            "{ACTOR_PANIC_MESSAGE}"
        );

        PermitFuture(rx)
    }

    /// Await a single permit for this endpoint, but only if the predicate evaluates
    /// to `true`.
    ///
    /// Permits are queued per endpoint in the order they were requested.
    ///
    /// Note that the predicate is asynchronously called in the actor task.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # let rt = tokio::runtime::Builder::new_current_thread()
    /// #     .enable_time()
    /// #     .build()
    /// #     .unwrap();
    /// # rt.block_on(async {
    /// # let rate_limiter = twilight_http_ratelimiting::RateLimiter::default();
    /// use twilight_http_ratelimiting::{Endpoint, Method};
    ///
    /// let endpoint = Endpoint {
    ///     method: Method::Get,
    ///     path: String::from("applications/@me"),
    /// };
    /// if let Some(permit) = rate_limiter
    ///     .acquire_if(endpoint, |b| b.is_none_or(|b| b.remaining > 10))
    ///     .await
    /// {
    ///     let headers = unimplemented!("GET /applications/@me");
    ///     permit.complete(headers);
    /// }
    /// # });
    /// ```
    #[allow(clippy::missing_panics_doc)]
    pub fn acquire_if<P>(&self, endpoint: Endpoint, predicate: P) -> MaybePermitFuture
    where
        P: FnOnce(Option<Bucket>) -> bool + Send + 'static,
    {
        let (notifier, rx) = oneshot::channel();
        let message = actor::Message { endpoint, notifier };
        assert!(
            self.tx.send((message, Some(Box::new(predicate)))).is_ok(),
            "{ACTOR_PANIC_MESSAGE}"
        );

        MaybePermitFuture(rx)
    }

    /// Retrieve the [`Bucket`] for this endpoint.
    ///
    /// The bucket is internally retrieved via [`acquire_if`][Self::acquire_if].
    #[allow(clippy::missing_panics_doc)]
    pub async fn bucket(&self, endpoint: Endpoint) -> Option<Bucket> {
        let (tx, rx) = oneshot::channel();
        self.acquire_if(endpoint, |bucket| {
            // Ignore cancellation error.
            _ = tx.send(bucket);
            false
        })
        .await;

        #[allow(clippy::match_wild_err_arm)]
        match rx.await {
            Ok(bucket) => bucket,
            Err(_) => panic!("{ACTOR_PANIC_MESSAGE}"),
        }
    }
}

impl Default for RateLimiter {
    /// Create a new [`RateLimiter`] with Discord's default global limit.
    ///
    /// Currently this is `50`.
    fn default() -> Self {
        Self::new(50)
    }
}

#[cfg(test)]
mod tests {
    use super::{
        Bucket, Endpoint, MaybePermitFuture, Method, Permit, PermitFuture, RateLimitHeaders,
        RateLimiter,
    };
    use static_assertions::assert_impl_all;
    use std::{
        fmt::Debug,
        future::Future,
        hash::{DefaultHasher, Hash, Hasher as _},
        time::{Duration, Instant},
    };
    use tokio::task;

    assert_impl_all!(Bucket: Clone, Copy, Debug, Eq, Hash, PartialEq, Send, Sync);
    assert_impl_all!(Endpoint: Clone, Debug, Eq, Hash, PartialEq, Send, Sync);
    assert_impl_all!(MaybePermitFuture: Debug, Future<Output = Option<Permit>>);
    assert_impl_all!(Method: Clone, Copy, Debug, Eq, PartialEq);
    assert_impl_all!(Permit: Debug, Send, Sync);
    assert_impl_all!(PermitFuture: Debug, Future<Output = Permit>);
    assert_impl_all!(RateLimitHeaders: Clone, Debug, Eq, Hash, PartialEq, Send, Sync);
    assert_impl_all!(RateLimiter: Clone, Debug, Default, Send, Sync);

    const ENDPOINT: fn() -> Endpoint = || Endpoint {
        method: Method::Get,
        path: String::from("applications/@me"),
    };

    #[tokio::test]
    async fn acquire_if() {
        let rate_limiter = RateLimiter::default();

        assert!(
            rate_limiter
                .acquire_if(ENDPOINT(), |_| false)
                .await
                .is_none()
        );
        assert!(
            rate_limiter
                .acquire_if(ENDPOINT(), |_| true)
                .await
                .is_some()
        );
    }

    #[tokio::test]
    async fn bucket() {
        let rate_limiter = RateLimiter::default();

        let limit = 2;
        let remaining = 1;
        let reset_at = Instant::now() + Duration::from_secs(1);
        let headers = RateLimitHeaders {
            bucket: vec![1, 2, 3],
            limit,
            remaining,
            reset_at,
        };

        rate_limiter
            .acquire(ENDPOINT())
            .await
            .complete(Some(headers));
        task::yield_now().await;

        let bucket = rate_limiter.bucket(ENDPOINT()).await.unwrap();
        assert_eq!(bucket.limit, limit);
        assert_eq!(bucket.remaining, remaining);
        assert!(
            bucket.reset_at.saturating_duration_since(reset_at) < Duration::from_millis(1)
                && reset_at.saturating_duration_since(bucket.reset_at) < Duration::from_millis(1)
        );
    }

    fn with_hasher(f: impl FnOnce(&mut DefaultHasher)) -> u64 {
        let mut hasher = DefaultHasher::new();
        f(&mut hasher);
        hasher.finish()
    }

    #[test]
    fn endpoint() {
        let invalid = Endpoint {
            method: Method::Get,
            path: String::from("/guilds/745809834183753828/audit-logs?limit=10"),
        };
        let delete_webhook = Endpoint {
            method: Method::Delete,
            path: String::from("webhooks/1"),
        };
        let interaction_response = Endpoint {
            method: Method::Post,
            path: String::from("interactions/1/abc/callback"),
        };

        assert!(!invalid.is_valid());
        assert!(delete_webhook.is_valid());
        assert!(interaction_response.is_valid());

        assert!(delete_webhook.is_interaction());
        assert!(interaction_response.is_interaction());

        assert_eq!(
            with_hasher(|state| invalid.hash_resources(state)),
            with_hasher(|_| {})
        );
        assert_eq!(
            with_hasher(|state| delete_webhook.hash_resources(state)),
            with_hasher(|state| {
                "webhooks".hash(state);
                b"1".hash(state);
            })
        );
        assert_eq!(
            with_hasher(|state| interaction_response.hash_resources(state)),
            with_hasher(|_| {})
        );
    }

    #[test]
    fn method_conversions() {
        assert_eq!("DELETE", Method::Delete.name());
        assert_eq!("GET", Method::Get.name());
        assert_eq!("PATCH", Method::Patch.name());
        assert_eq!("POST", Method::Post.name());
        assert_eq!("PUT", Method::Put.name());
    }
}