leaky_bucket/

lib.rs

//! [<img alt="github" src="https://img.shields.io/badge/github-udoprog/leaky--bucket-8da0cb?style=for-the-badge&logo=github" height="20">](https://github.com/udoprog/leaky-bucket)
//! [<img alt="crates.io" src="https://img.shields.io/crates/v/leaky-bucket.svg?style=for-the-badge&color=fc8d62&logo=rust" height="20">](https://crates.io/crates/leaky-bucket)
//! [<img alt="docs.rs" src="https://img.shields.io/badge/docs.rs-leaky--bucket-66c2a5?style=for-the-badge&logoColor=white&logo=data:image/svg+xml;base64,PHN2ZyByb2xlPSJpbWciIHhtbG5zPSJodHRwOi8vd3d3LnczLm9yZy8yMDAwL3N2ZyIgdmlld0JveD0iMCAwIDUxMiA1MTIiPjxwYXRoIGZpbGw9IiNmNWY1ZjUiIGQ9Ik00ODguNiAyNTAuMkwzOTIgMjE0VjEwNS41YzAtMTUtOS4zLTI4LjQtMjMuNC0zMy43bC0xMDAtMzcuNWMtOC4xLTMuMS0xNy4xLTMuMS0yNS4zIDBsLTEwMCAzNy41Yy0xNC4xIDUuMy0yMy40IDE4LjctMjMuNCAzMy43VjIxNGwtOTYuNiAzNi4yQzkuMyAyNTUuNSAwIDI2OC45IDAgMjgzLjlWMzk0YzAgMTMuNiA3LjcgMjYuMSAxOS45IDMyLjJsMTAwIDUwYzEwLjEgNS4xIDIyLjEgNS4xIDMyLjIgMGwxMDMuOS01MiAxMDMuOSA1MmMxMC4xIDUuMSAyMi4xIDUuMSAzMi4yIDBsMTAwLTUwYzEyLjItNi4xIDE5LjktMTguNiAxOS45LTMyLjJWMjgzLjljMC0xNS05LjMtMjguNC0yMy40LTMzLjd6TTM1OCAyMTQuOGwtODUgMzEuOXYtNjguMmw4NS0zN3Y3My4zek0xNTQgMTA0LjFsMTAyLTM4LjIgMTAyIDM4LjJ2LjZsLTEwMiA0MS40LTEwMi00MS40di0uNnptODQgMjkxLjFsLTg1IDQyLjV2LTc5LjFsODUtMzguOHY3NS40em0wLTExMmwtMTAyIDQxLjQtMTAyLTQxLjR2LS42bDEwMi0zOC4yIDEwMiAzOC4ydi42em0yNDAgMTEybC04NSA0Mi41di03OS4xbDg1LTM4Ljh2NzUuNHptMC0xMTJsLTEwMiA0MS40LTEwMi00MS40di0uNmwxMDItMzguMiAxMDIgMzguMnYuNnoiPjwvcGF0aD48L3N2Zz4K" height="20">](https://docs.rs/leaky-bucket)
//!
//! A token-based rate limiter based on the [leaky bucket] algorithm.
//!
//! If the bucket overflows and goes over its max configured capacity, the task
//! that tried to acquire the tokens will be suspended until the required number
//! of tokens has been drained from the bucket.
//!
//! Since this crate uses timing facilities from tokio it has to be used within
//! a Tokio runtime with the [`time` feature] enabled.
//!
//! This library has some neat features, which includes:
//!
//! **Not requiring a background task**. This is usually needed by token bucket
//! rate limiters to drive progress. Instead, one of the waiting tasks
//! temporarily assumes the role as coordinator (called the *core*). This
//! reduces the amount of tasks needing to sleep, which can be a source of
//! jitter for imprecise sleeping implementations and tight limiters. See below
//! for more details.
//!
//! **Dropped tasks** release any resources they've reserved. So that
//! constructing and cancellaing asynchronous tasks to not end up taking up wait
//! slots it never uses which would be the case for cell-based rate limiters.
//!
//! <br>
//!
//! ## Usage
//!
//! The core type is [`RateLimiter`], which allows for limiting the throughput
//! of a section using its [`acquire`], [`try_acquire`], and [`acquire_one`]
//! methods.
//!
//! The following is a simple example where we wrap requests through a HTTP
//! `Client`, to ensure that we don't exceed a given limit:
//!
//! ```
//! use leaky_bucket::RateLimiter;
//! # struct Client;
//! # impl Client { async fn request<T>(&self, path: &str) -> Result<T> { todo!() } }
//! # trait DeserializeOwned {}
//! # impl DeserializeOwned for Vec<Post> {}
//! # type Result<T> = core::result::Result<T, ()>;
//!
//! /// A blog client.
//! pub struct BlogClient {
//!     limiter: RateLimiter,
//!     client: Client,
//! }
//!
//! struct Post {
//!     // ..
//! }
//!
//! impl BlogClient {
//!     /// Get all posts from the service.
//!     pub async fn get_posts(&self) -> Result<Vec<Post>> {
//!         self.request("posts").await
//!     }
//!
//!     /// Perform a request against the service, limiting requests to abide by a rate limit.
//!     async fn request<T>(&self, path: &str) -> Result<T>
//!     where
//!         T: DeserializeOwned
//!     {
//!         // Before we start sending a request, we block on acquiring one token.
//!         self.limiter.acquire(1).await;
//!         self.client.request::<T>(path).await
//!     }
//! }
//! ```
//!
//! <br>
//!
//! ## Implementation details
//!
//! Each rate limiter has two acquisition modes. A fast path and a slow path.
//! The fast path is used if the desired number of tokens are readily available,
//! and simply involves decrementing the number of tokens available in the
//! shared pool.
//!
//! If the required number of tokens is not available, the task will be forced
//! to be suspended until the next refill interval. Here one of the acquiring
//! tasks will switch over to work as a *core*. This is known as *core
//! switching*.
//!
//! ```
//! use leaky_bucket::RateLimiter;
//! use tokio::time::Duration;
//!
//! # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() {
//! let limiter = RateLimiter::builder()
//!     .initial(10)
//!     .interval(Duration::from_millis(100))
//!     .build();
//!
//! // This is instantaneous since the rate limiter starts with 10 tokens to
//! // spare.
//! limiter.acquire(10).await;
//!
//! // This however needs to core switch and wait for a while until the desired
//! // number of tokens is available.
//! limiter.acquire(3).await;
//! # }
//! ```
//!
//! The core is responsible for sleeping for the configured interval so that
//! more tokens can be added. After which it ensures that any tasks that are
//! waiting to acquire including itself are appropriately unsuspended.
//!
//! On-demand core switching is what allows this rate limiter implementation to
//! work without a coordinating background thread. But we need to ensure that
//! any asynchronous tasks that uses [`RateLimiter`] must either run an
//! [`acquire`] call to completion, or be *cancelled* by being dropped.
//!
//! If none of these hold, the core might leak and be locked indefinitely
//! preventing any future use of the rate limiter from making progress. This is
//! similar to if you would lock an asynchronous [`Mutex`] but never drop its
//! guard.
//!
//! > You can run this example with:
//! >
//! > ```sh
//! > cargo run --example block_forever
//! > ```
//!
//! ```no_run
//! use std::future::Future;
//! use std::sync::Arc;
//! use std::task::Context;
//!
//! use leaky_bucket::RateLimiter;
//!
//! struct Waker;
//! # impl std::task::Wake for Waker { fn wake(self: Arc<Self>) { } }
//!
//! # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() {
//! let limiter = Arc::new(RateLimiter::builder().build());
//!
//! let waker = Arc::new(Waker).into();
//! let mut cx = Context::from_waker(&waker);
//!
//! let mut a0 = Box::pin(limiter.acquire(1));
//! // Poll once to ensure that the core task is assigned.
//! assert!(a0.as_mut().poll(&mut cx).is_pending());
//! assert!(a0.is_core());
//!
//! // We leak the core task, preventing the rate limiter from making progress
//! // by assigning new core tasks.
//! std::mem::forget(a0);
//!
//! // Awaiting acquire here would block forever.
//! // limiter.acquire(1).await;
//! # }
//! ```
//!
//! <br>
//!
//! ## Fairness
//!
//! By default [`RateLimiter`] uses a *fair* scheduler. This ensures that the
//! core task makes progress even if there are many tasks waiting to acquire
//! tokens. This might cause more core switching, increasing the total work
//! needed. An unfair scheduler is expected to do a bit less work under
//! contention. But without fair scheduling some tasks might end up taking
//! longer to acquire than expected.
//!
//! Unfair rate limiters also have access to a fast path for acquiring tokens,
//! which might further improve throughput.
//!
//! This behavior can be tweaked with the [`Builder::fair`] option.
//!
//! ```
//! use leaky_bucket::RateLimiter;
//!
//! let limiter = RateLimiter::builder()
//!     .fair(false)
//!     .build();
//! ```
//!
//! The `unfair-scheduling` example can showcase this phenomenon.
//!
//! ```sh
//! cargo run --example unfair_scheduling
//! ```
//!
//! ```text
//! # fair
//! Max: 1011ms, Total: 1012ms
//! Timings:
//!  0: 101ms
//!  1: 101ms
//!  2: 101ms
//!  3: 101ms
//!  4: 101ms
//!  ...
//! # unfair
//! Max: 1014ms, Total: 1014ms
//! Timings:
//!  0: 1014ms
//!  1: 101ms
//!  2: 101ms
//!  3: 101ms
//!  4: 101ms
//!  ...
//! ```
//!
//! As can be seen above the first task in the *unfair* scheduler takes longer
//! to run because it prioritises releasing other tasks waiting to acquire over
//! itself.
//!
//! [`acquire_one`]: https://docs.rs/leaky-bucket/1/leaky_bucket/struct.RateLimiter.html#method.acquire_one
//! [`acquire`]: https://docs.rs/leaky-bucket/1/leaky_bucket/struct.RateLimiter.html#method.acquire
//! [`Builder::fair`]: https://docs.rs/leaky-bucket/1/leaky_bucket/struct.Builder.html#method.fair
//! [`Mutex`]: https://docs.rs/tokio/1/tokio/sync/struct.Mutex.html
//! [`RateLimiter`]: https://docs.rs/leaky-bucket/1/leaky_bucket/struct.RateLimiter.html
//! [`time` feature]: https://docs.rs/tokio/1/tokio/#feature-flags
//! [`try_acquire`]: https://docs.rs/leaky-bucket/1/leaky_bucket/struct.RateLimiter.html#method.try_acquire
//! [leaky bucket]: https://en.wikipedia.org/wiki/Leaky_bucket

#![no_std]
#![deny(missing_docs)]

extern crate alloc;

#[macro_use]
extern crate std;

use core::cell::UnsafeCell;
use core::convert::TryFrom as _;
use core::fmt;
use core::future::Future;
use core::mem::{self, ManuallyDrop};
use core::ops::{Deref, DerefMut};
use core::pin::Pin;
use core::ptr;
use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use core::task::{Context, Poll, Waker};

use alloc::sync::Arc;

use parking_lot::{Mutex, MutexGuard};
use pin_project_lite::pin_project;
use tokio::time::{self, Duration, Instant};

#[cfg(feature = "tracing")]
macro_rules! trace {
    ($($arg:tt)*) => {
        tracing::trace!($($arg)*)
    };
}

#[cfg(not(feature = "tracing"))]
macro_rules! trace {
    ($($arg:tt)*) => {};
}

mod linked_list;
use self::linked_list::{LinkedList, Node};

/// Default factor for how to calculate max refill value.
const DEFAULT_REFILL_MAX_FACTOR: usize = 10;

/// Interval to bump the shared mutex guard to allow other parts of the system
/// to make process. Processes which loop should use this number to determine
/// how many times it should loop before calling [Guard::bump].
///
/// If we do not respect this limit we might inadvertently end up starving other
/// tasks from making progress so that they can unblock.
const BUMP_LIMIT: usize = 16;

/// The maximum supported balance.
const MAX_BALANCE: usize = isize::MAX as usize;

/// Marker trait which indicates that a type represents a unique held critical section.
trait IsCritical {}
impl IsCritical for Critical {}
impl IsCritical for Guard<'_> {}

/// Linked task state.
struct Task {
    /// Remaining tokens that need to be satisfied.
    remaining: usize,
    /// If this node has been released or not. We make this an atomic to permit
    /// access to it without synchronization.
    complete: AtomicBool,
    /// The waker associated with the node.
    waker: Option<Waker>,
}

impl Task {
    /// Construct a new task state with the given permits remaining.
    const fn new() -> Self {
        Self {
            remaining: 0,
            complete: AtomicBool::new(false),
            waker: None,
        }
    }

    /// Test if the current node is completed.
    fn is_completed(&self) -> bool {
        self.remaining == 0
    }

    /// Fill the current node from the given pool of tokens and modify it.
    fn fill(&mut self, current: &mut usize) {
        let removed = usize::min(self.remaining, *current);
        self.remaining -= removed;
        *current -= removed;
    }
}

/// A borrowed rate limiter.
struct BorrowedRateLimiter<'a>(&'a RateLimiter);

impl Deref for BorrowedRateLimiter<'_> {
    type Target = RateLimiter;

    #[inline]
    fn deref(&self) -> &RateLimiter {
        self.0
    }
}

struct Critical {
    /// Waiter list.
    waiters: LinkedList<Task>,
    /// The deadline for when more tokens can be be added.
    deadline: Instant,
}

#[repr(transparent)]
struct Guard<'a> {
    critical: MutexGuard<'a, Critical>,
}

impl Guard<'_> {
    #[inline]
    fn bump(this: &mut Guard<'_>) {
        MutexGuard::bump(&mut this.critical)
    }
}

impl Deref for Guard<'_> {
    type Target = Critical;

    #[inline]
    fn deref(&self) -> &Critical {
        &self.critical
    }
}

impl DerefMut for Guard<'_> {
    #[inline]
    fn deref_mut(&mut self) -> &mut Critical {
        &mut self.critical
    }
}

impl Critical {
    #[inline]
    fn push_task_front(&mut self, task: &mut Node<Task>) {
        // SAFETY: We both have mutable access to the node being pushed, and
        // mutable access to the critical section through `self`. So we know we
        // have exclusive tampering rights to the waiter queue.
        unsafe {
            self.waiters.push_front(task.into());
        }
    }

    #[inline]
    fn push_task(&mut self, task: &mut Node<Task>) {
        // SAFETY: We both have mutable access to the node being pushed, and
        // mutable access to the critical section through `self`. So we know we
        // have exclusive tampering rights to the waiter queue.
        unsafe {
            self.waiters.push_back(task.into());
        }
    }

    #[inline]
    fn remove_task(&mut self, task: &mut Node<Task>) {
        // SAFETY: We both have mutable access to the node being pushed, and
        // mutable access to the critical section through `self`. So we know we
        // have exclusive tampering rights to the waiter queue.
        unsafe {
            self.waiters.remove(task.into());
        }
    }

    /// Release the current core. Beyond this point the current task may no
    /// longer interact exclusively with the core.
    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self), level = "trace"))]
    fn release(&mut self, state: &mut State<'_>) {
        trace!("releasing core");
        state.available = true;

        // Find another task that might take over as core. Once it has acquired
        // core status it will have to make sure it is no longer linked into the
        // wait queue.
        unsafe {
            if let Some(node) = self.waiters.front() {
                trace!(node = ?node, "waking next core");

                if let Some(ref waker) = node.as_ref().waker {
                    waker.wake_by_ref();
                }
            }
        }
    }
}

#[derive(Debug)]
struct State<'a> {
    /// Original state.
    state: usize,
    /// If the core is available or not.
    available: bool,
    /// The balance.
    balance: usize,
    /// The rate limiter the state is associated with.
    lim: &'a RateLimiter,
}

impl<'a> State<'a> {
    fn try_fast_path(mut self, permits: usize) -> bool {
        let mut attempts = 0;

        // Fast path where we just try to nab any available permit without
        // locking.
        //
        // We do have to race against anyone else grabbing permits here when
        // storing the state back.
        while self.balance >= permits {
            // Abandon fast path if we've tried too many times.
            if attempts == BUMP_LIMIT {
                break;
            }

            self.balance -= permits;

            if let Err(new_state) = self.try_save() {
                self = new_state;
                attempts += 1;
                continue;
            }

            return true;
        }

        false
    }

    /// Add tokens and release any pending tasks.
    #[cfg_attr(
        feature = "tracing",
        tracing::instrument(skip(self, critical, f), level = "trace")
    )]
    #[inline]
    fn add_tokens<F, O>(&mut self, critical: &mut Guard<'_>, tokens: usize, f: F) -> O
    where
        F: FnOnce(&mut Guard<'_>, &mut State) -> O,
    {
        if tokens > 0 {
            debug_assert!(
                tokens <= MAX_BALANCE,
                "Additional tokens {} must be less than {}",
                tokens,
                MAX_BALANCE
            );

            self.balance = (self.balance + tokens).min(self.lim.max);
            drain_wait_queue(critical, self);
            let output = f(critical, self);
            return output;
        }

        f(critical, self)
    }

    #[inline]
    fn decode(state: usize, lim: &'a RateLimiter) -> Self {
        State {
            state,
            available: state & 1 == 1,
            balance: state >> 1,
            lim,
        }
    }

    #[inline]
    fn encode(&self) -> usize {
        (self.balance << 1) | usize::from(self.available)
    }

    /// Try to save the state, but only succeed if it hasn't been modified.
    #[inline]
    fn try_save(self) -> Result<(), Self> {
        let this = ManuallyDrop::new(self);

        match this.lim.state.compare_exchange(
            this.state,
            this.encode(),
            Ordering::Release,
            Ordering::Relaxed,
        ) {
            Ok(_) => Ok(()),
            Err(state) => Err(State::decode(state, this.lim)),
        }
    }
}

impl Drop for State<'_> {
    #[inline]
    fn drop(&mut self) {
        self.lim.state.store(self.encode(), Ordering::Release);
    }
}

/// A token-bucket rate limiter.
pub struct RateLimiter {
    /// Tokens to add every `per` duration.
    refill: usize,
    /// Interval in milliseconds to add tokens.
    interval: Duration,
    /// Max number of tokens associated with the rate limiter.
    max: usize,
    /// If the rate limiter is fair or not.
    fair: bool,
    /// The state of the rate limiter.
    state: AtomicUsize,
    /// Critical state of the rate limiter.
    critical: Mutex<Critical>,
}

impl RateLimiter {
    /// Construct a new [`Builder`] for a [`RateLimiter`].
    ///
    /// # Examples
    ///
    /// ```
    /// use leaky_bucket::RateLimiter;
    /// use tokio::time::Duration;
    ///
    /// let limiter = RateLimiter::builder()
    ///     .initial(100)
    ///     .refill(100)
    ///     .max(1000)
    ///     .interval(Duration::from_millis(250))
    ///     .fair(false)
    ///     .build();
    /// ```
    pub fn builder() -> Builder {
        Builder::default()
    }

    /// Get the refill amount  of this rate limiter as set through
    /// [`Builder::refill`].
    ///
    /// # Examples
    ///
    /// ```
    /// use leaky_bucket::RateLimiter;
    ///
    /// let limiter = RateLimiter::builder()
    ///     .refill(1024)
    ///     .build();
    ///
    /// assert_eq!(limiter.refill(), 1024);
    /// ```
    pub fn refill(&self) -> usize {
        self.refill
    }

    /// Get the refill interval of this rate limiter as set through
    /// [`Builder::interval`].
    ///
    /// # Examples
    ///
    /// ```
    /// use leaky_bucket::RateLimiter;
    /// use tokio::time::Duration;
    ///
    /// let limiter = RateLimiter::builder()
    ///     .interval(Duration::from_millis(1000))
    ///     .build();
    ///
    /// assert_eq!(limiter.interval(), Duration::from_millis(1000));
    /// ```
    pub fn interval(&self) -> Duration {
        self.interval
    }

    /// Get the max value of this rate limiter as set through [`Builder::max`].
    ///
    /// # Examples
    ///
    /// ```
    /// use leaky_bucket::RateLimiter;
    ///
    /// let limiter = RateLimiter::builder()
    ///     .max(1024)
    ///     .build();
    ///
    /// assert_eq!(limiter.max(), 1024);
    /// ```
    pub fn max(&self) -> usize {
        self.max
    }

    /// Test if the current rate limiter is fair as specified through
    /// [`Builder::fair`].
    ///
    /// # Examples
    ///
    /// ```
    /// use leaky_bucket::RateLimiter;
    ///
    /// let limiter = RateLimiter::builder()
    ///     .fair(true)
    ///     .build();
    ///
    /// assert_eq!(limiter.is_fair(), true);
    /// ```
    pub fn is_fair(&self) -> bool {
        self.fair
    }

    /// Get the current token balance.
    ///
    /// This indicates how many tokens can be requested without blocking.
    ///
    /// # Examples
    ///
    /// ```
    /// use leaky_bucket::RateLimiter;
    ///
    /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() {
    /// let limiter = RateLimiter::builder()
    ///     .initial(100)
    ///     .build();
    ///
    /// assert_eq!(limiter.balance(), 100);
    /// limiter.acquire(10).await;
    /// assert_eq!(limiter.balance(), 90);
    /// # }
    /// ```
    pub fn balance(&self) -> usize {
        self.state.load(Ordering::Acquire) >> 1
    }

    /// Acquire a single permit.
    ///
    /// # Examples
    ///
    /// ```
    /// use leaky_bucket::RateLimiter;
    ///
    /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() {
    /// let limiter = RateLimiter::builder()
    ///     .initial(10)
    ///     .build();
    ///
    /// limiter.acquire_one().await;
    /// # }
    /// ```
    pub fn acquire_one(&self) -> Acquire<'_> {
        self.acquire(1)
    }

    /// Acquire the given number of permits, suspending the current task until
    /// they are available.
    ///
    /// If zero permits are specified, this function never suspends the current
    /// task.
    ///
    /// # Examples
    ///
    /// ```
    /// use leaky_bucket::RateLimiter;
    ///
    /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() {
    /// let limiter = RateLimiter::builder()
    ///     .initial(10)
    ///     .build();
    ///
    /// limiter.acquire(10).await;
    /// # }
    /// ```
    pub fn acquire(&self, permits: usize) -> Acquire<'_> {
        Acquire {
            inner: AcquireFut::new(BorrowedRateLimiter(self), permits),
        }
    }

    /// Try to acquire the given number of permits, returning `true` if the
    /// given number of permits were successfully acquired.
    ///
    /// If the scheduler is fair, and there are pending tasks waiting to acquire
    /// tokens this method will return `false`.
    ///
    /// If zero permits are specified, this method returns `true`.
    ///
    /// # Examples
    ///
    /// ```
    /// use leaky_bucket::RateLimiter;
    /// use tokio::time;
    ///
    /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() {
    /// let limiter = RateLimiter::builder().refill(1).initial(1).build();
    ///
    /// assert!(limiter.try_acquire(1));
    /// assert!(!limiter.try_acquire(1));
    /// assert!(limiter.try_acquire(0));
    ///
    /// time::sleep(limiter.interval() * 2).await;
    ///
    /// assert!(limiter.try_acquire(1));
    /// assert!(limiter.try_acquire(1));
    /// assert!(!limiter.try_acquire(1));
    /// # }
    /// ```
    pub fn try_acquire(&self, permits: usize) -> bool {
        if self.try_fast_path(permits) {
            return true;
        }

        let mut critical = self.lock();

        // Reload the state while we are under the critical lock, this
        // ensures that the `available` flag is up-to-date since it is only
        // ever modified while holding the critical lock.
        let mut state = self.take();

        // The core is *not* available, which also implies that there are tasks
        // ahead which are busy.
        if !state.available {
            return false;
        }

        let now = Instant::now();

        // Here we try to assume core duty temporarily to see if we can
        // release a sufficient number of tokens to allow the current task
        // to proceed.
        if let Some((tokens, deadline)) = self.calculate_drain(critical.deadline, now) {
            state.balance = (state.balance + tokens).min(self.max);
            critical.deadline = deadline;
        }

        if state.balance >= permits {
            state.balance -= permits;
            return true;
        }

        false
    }

    /// Acquire a permit using an owned future.
    ///
    /// If zero permits are specified, this function never suspends the current
    /// task.
    ///
    /// This required the [`RateLimiter`] to be wrapped inside of an
    /// [`std::sync::Arc`] but will in contrast permit the acquire operation to
    /// be owned by another struct making it more suitable for embedding.
    ///
    /// # Examples
    ///
    /// ```
    /// use leaky_bucket::RateLimiter;
    /// use std::sync::Arc;
    ///
    /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() {
    /// let limiter = Arc::new(RateLimiter::builder().initial(10).build());
    ///
    /// limiter.acquire_owned(10).await;
    /// # }
    /// ```
    ///
    /// Example when embedded into another future. This wouldn't be possible
    /// with [`RateLimiter::acquire`] since it would otherwise hold a reference
    /// to the corresponding [`RateLimiter`] instance.
    ///
    /// ```
    /// use std::future::Future;
    /// use std::pin::Pin;
    /// use std::sync::Arc;
    /// use std::task::{Context, Poll};
    ///
    /// use leaky_bucket::{AcquireOwned, RateLimiter};
    /// use pin_project::pin_project;
    ///
    /// #[pin_project]
    /// struct MyFuture {
    ///     limiter: Arc<RateLimiter>,
    ///     #[pin]
    ///     acquire: Option<AcquireOwned>,
    /// }
    ///
    /// impl Future for MyFuture {
    ///     type Output = ();
    ///
    ///     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
    ///         let mut this = self.project();
    ///
    ///         loop {
    ///             if let Some(acquire) = this.acquire.as_mut().as_pin_mut() {
    ///                 futures::ready!(acquire.poll(cx));
    ///                 return Poll::Ready(());
    ///             }
    ///
    ///             this.acquire.set(Some(this.limiter.clone().acquire_owned(100)));
    ///         }
    ///     }
    /// }
    ///
    /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() {
    /// let limiter = Arc::new(RateLimiter::builder().initial(100).build());
    ///
    /// let future = MyFuture { limiter, acquire: None };
    /// future.await;
    /// # }
    /// ```
    pub fn acquire_owned(self: Arc<Self>, permits: usize) -> AcquireOwned {
        AcquireOwned {
            inner: AcquireFut::new(self, permits),
        }
    }

    /// Lock the critical section of the rate limiter and return the associated guard.
    fn lock(&self) -> Guard<'_> {
        Guard {
            critical: self.critical.lock(),
        }
    }

    /// Load the current state.
    fn load(&self) -> State<'_> {
        State::decode(self.state.load(Ordering::Acquire), self)
    }

    /// Take the current state, leaving the core state intact.
    fn take(&self) -> State<'_> {
        State::decode(self.state.swap(0, Ordering::Acquire), self)
    }

    /// Try to use fast path.
    fn try_fast_path(&self, permits: usize) -> bool {
        if permits == 0 {
            return true;
        }

        if self.fair {
            return false;
        }

        self.load().try_fast_path(permits)
    }

    /// Calculate refill amount. Returning a tuple of how much to fill and remaining
    /// duration to sleep until the next refill time if appropriate.
    ///
    /// The maximum number of additional tokens this method will ever return is
    /// limited to [`MAX_BALANCE`] to ensure that addition with an existing
    /// balance will never overflow.
    fn calculate_drain(&self, deadline: Instant, now: Instant) -> Option<(usize, Instant)> {
        if now < deadline {
            return None;
        }

        // Time elapsed in milliseconds since the last deadline.
        let millis = self.interval.as_millis();
        let since = now.saturating_duration_since(deadline).as_millis();

        let periods = usize::try_from(since / millis + 1).unwrap_or(usize::MAX);

        let tokens = periods
            .checked_mul(self.refill)
            .unwrap_or(MAX_BALANCE)
            .min(MAX_BALANCE);

        let rem = u64::try_from(since % millis).unwrap_or(u64::MAX);

        // Calculated time remaining until the next deadline.
        let deadline = now
            + self
                .interval
                .saturating_sub(time::Duration::from_millis(rem));

        Some((tokens, deadline))
    }
}

impl fmt::Debug for RateLimiter {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("RateLimiter")
            .field("refill", &self.refill)
            .field("interval", &self.interval)
            .field("max", &self.max)
            .field("fair", &self.fair)
            .finish_non_exhaustive()
    }
}

/// Refill the wait queue with the given number of tokens.
#[cfg_attr(feature = "tracing", tracing::instrument(skip_all, level = "trace"))]
fn drain_wait_queue(critical: &mut Guard<'_>, state: &mut State<'_>) {
    trace!(?state, "releasing waiters");

    let mut bump = 0;

    // SAFETY: we're holding the lock guard to all the waiters so we can be
    // sure that we have exclusive access to the wait queue.
    unsafe {
        while state.balance > 0 {
            let mut node = match critical.waiters.pop_back() {
                Some(node) => node,
                None => break,
            };

            let n = node.as_mut();
            n.fill(&mut state.balance);

            trace! {
                ?state,
                remaining = n.remaining,
                "filled node",
            };

            if !n.is_completed() {
                critical.waiters.push_back(node);
                break;
            }

            n.complete.store(true, Ordering::Release);

            if let Some(waker) = n.waker.take() {
                waker.wake();
            }

            bump += 1;

            if bump == BUMP_LIMIT {
                Guard::bump(critical);
                bump = 0;
            }
        }
    }
}

// SAFETY: All the internals of acquire is thread safe and correctly
// synchronized. The embedded waiter queue doesn't have anything inherently
// unsafe in it.
unsafe impl Send for RateLimiter {}
unsafe impl Sync for RateLimiter {}

/// A builder for a [`RateLimiter`].
pub struct Builder {
    /// The max number of tokens.
    max: Option<usize>,
    /// The initial count of tokens.
    initial: usize,
    /// Tokens to add every `per` duration.
    refill: usize,
    /// Interval to add tokens in milliseconds.
    interval: Duration,
    /// If the rate limiter is fair or not.
    fair: bool,
}

impl Builder {
    /// Configure the max number of tokens to use.
    ///
    /// If unspecified, this will default to be 10 times the [`refill`] or the
    /// [`initial`] value, whichever is largest.
    ///
    /// The maximum supported balance is limited to [`isize::MAX`].
    ///
    /// # Examples
    ///
    /// ```
    /// use leaky_bucket::RateLimiter;
    ///
    /// let limiter = RateLimiter::builder()
    ///     .max(10_000)
    ///     .build();
    /// ```
    ///
    /// [`refill`]: Builder::refill
    /// [`initial`]: Builder::initial
    pub fn max(&mut self, max: usize) -> &mut Self {
        self.max = Some(max);
        self
    }

    /// Configure the initial number of tokens to configure. The default value
    /// is `0`.
    ///
    /// # Examples
    ///
    /// ```
    /// use leaky_bucket::RateLimiter;
    ///
    /// let limiter = RateLimiter::builder()
    ///     .initial(10)
    ///     .build();
    /// ```
    pub fn initial(&mut self, initial: usize) -> &mut Self {
        self.initial = initial;
        self
    }

    /// Configure the time duration between which we add [`refill`] number to
    /// the bucket rate limiter.
    ///
    /// This is 100ms by default.
    ///
    /// # Panics
    ///
    /// This panics if the provided interval does not fit within the millisecond
    /// bounds of a [usize] or is zero.
    ///
    /// ```should_panic
    /// use leaky_bucket::RateLimiter;
    /// use tokio::time::Duration;
    ///
    /// let limiter = RateLimiter::builder()
    ///     .interval(Duration::from_secs(u64::MAX))
    ///     .build();
    /// ```
    ///
    /// ```should_panic
    /// use leaky_bucket::RateLimiter;
    /// use tokio::time::Duration;
    ///
    /// let limiter = RateLimiter::builder()
    ///     .interval(Duration::from_millis(0))
    ///     .build();
    /// ```
    ///
    /// # Examples
    ///
    /// ```
    /// use leaky_bucket::RateLimiter;
    /// use tokio::time::Duration;
    ///
    /// let limiter = RateLimiter::builder()
    ///     .interval(Duration::from_millis(100))
    ///     .build();
    /// ```
    ///
    /// [`refill`]: Builder::refill
    pub fn interval(&mut self, interval: Duration) -> &mut Self {
        assert! {
            interval.as_millis() != 0,
            "interval must be non-zero",
        };
        assert! {
            u64::try_from(interval.as_millis()).is_ok(),
            "interval must fit within a 64-bit integer"
        };
        self.interval = interval;
        self
    }

    /// The number of tokens to add at each [`interval`] interval. The default
    /// value is `1`.
    ///
    /// # Panics
    ///
    /// Panics if a refill amount of `0` is specified.
    ///
    /// # Examples
    ///
    /// ```
    /// use leaky_bucket::RateLimiter;
    ///
    /// let limiter = RateLimiter::builder()
    ///     .refill(100)
    ///     .build();
    /// ```
    ///
    /// [`interval`]: Builder::interval
    pub fn refill(&mut self, refill: usize) -> &mut Self {
        assert!(refill > 0, "refill amount cannot be zero");
        self.refill = refill;
        self
    }

    /// Configure the rate limiter to be fair.
    ///
    /// Fairness is enabled by deafult.
    ///
    /// Fairness ensures that tasks make progress in the order that they acquire
    /// even when the rate limiter is under contention. An unfair scheduler
    /// might have a higher total throughput.
    ///
    /// Fair scheduling also affects the behavior of
    /// [`RateLimiter::try_acquire`] which will return `false` if there are any
    /// pending tasks since they should be given priority.
    ///
    /// # Examples
    ///
    /// ```
    /// use leaky_bucket::RateLimiter;
    ///
    /// let limiter = RateLimiter::builder()
    ///     .refill(100)
    ///     .fair(false)
    ///     .build();
    /// ```
    pub fn fair(&mut self, fair: bool) -> &mut Self {
        self.fair = fair;
        self
    }

    /// Construct a new [`RateLimiter`].
    ///
    /// # Examples
    ///
    /// ```
    /// use leaky_bucket::RateLimiter;
    /// use tokio::time::Duration;
    ///
    /// let limiter = RateLimiter::builder()
    ///     .refill(100)
    ///     .interval(Duration::from_millis(200))
    ///     .max(10_000)
    ///     .build();
    /// ```
    pub fn build(&self) -> RateLimiter {
        let deadline = Instant::now() + self.interval;

        let initial = self.initial.min(MAX_BALANCE);
        let refill = self.refill.min(MAX_BALANCE);

        let max = match self.max {
            Some(max) => max.min(MAX_BALANCE),
            None => refill
                .max(initial)
                .saturating_mul(DEFAULT_REFILL_MAX_FACTOR)
                .min(MAX_BALANCE),
        };

        let initial = initial.min(max);

        RateLimiter {
            refill,
            interval: self.interval,
            max,
            fair: self.fair,
            state: AtomicUsize::new(initial << 1 | 1),
            critical: Mutex::new(Critical {
                waiters: LinkedList::new(),
                deadline,
            }),
        }
    }
}

/// Construct a new builder with default options.
///
/// # Examples
///
/// ```
/// use leaky_bucket::Builder;
///
/// let limiter = Builder::default().build();
/// ```
impl Default for Builder {
    fn default() -> Self {
        Self {
            max: None,
            initial: 0,
            refill: 1,
            interval: Duration::from_millis(100),
            fair: true,
        }
    }
}

/// The state of an acquire operation.
#[derive(Debug, Clone, Copy)]
enum AcquireFutState {
    /// Initial unconfigured state.
    Initial,
    /// The acquire is waiting to be released by the core.
    Waiting,
    /// The operation is completed.
    Complete,
    /// The task is currently the core.
    Core,
}

/// Inner state and methods of the acquire.
#[repr(transparent)]
struct AcquireFutInner {
    /// Aliased task state.
    node: UnsafeCell<Node<Task>>,
}

impl AcquireFutInner {
    const fn new() -> AcquireFutInner {
        AcquireFutInner {
            node: UnsafeCell::new(Node::new(Task::new())),
        }
    }

    /// Access the completion flag.
    pub fn complete(&self) -> &AtomicBool {
        // SAFETY: This is always safe to access since it's atomic.
        unsafe { &*ptr::addr_of!((*self.node.get()).complete) }
    }

    /// Get the underlying task mutably.
    ///
    /// We prove that the caller does indeed have mutable access to the node by
    /// passing in a mutable reference to the critical section.
    #[inline]
    pub fn get_task<'crit, C>(
        self: Pin<&'crit mut Self>,
        critical: &'crit mut C,
    ) -> (&'crit mut C, &'crit mut Node<Task>)
    where
        C: IsCritical,
    {
        // SAFETY: Caller has exclusive access to the critical section, since
        // it's passed in as a mutable argument. We can also ensure that none of
        // the borrows outlive the provided closure.
        unsafe { (critical, &mut *self.node.get()) }
    }

    /// Update the waiting state for this acquisition task. This might require
    /// that we update the associated waker.
    #[cfg_attr(
        feature = "tracing",
        tracing::instrument(skip(self, critical, waker), level = "trace")
    )]
    fn update(self: Pin<&mut Self>, critical: &mut Guard<'_>, waker: &Waker) {
        let (critical, task) = self.get_task(critical);

        if !task.is_linked() {
            critical.push_task_front(task);
        }

        let new_waker = match task.waker {
            None => true,
            Some(ref w) => !w.will_wake(waker),
        };

        if new_waker {
            trace!("updating waker");
            task.waker = Some(waker.clone());
        }
    }

    /// Ensure that the current core task is correctly linked up if needed.
    #[cfg_attr(
        feature = "tracing",
        tracing::instrument(skip(self, critical, lim), level = "trace")
    )]
    fn link_core(self: Pin<&mut Self>, critical: &mut Critical, lim: &RateLimiter) {
        let (critical, task) = self.get_task(critical);

        match (lim.fair, task.is_linked()) {
            (true, false) => {
                // Fair scheduling needs to ensure that the core is part of the wait
                // queue, and will be woken up in-order with other tasks.
                critical.push_task(task);
            }
            (false, true) => {
                // Unfair scheduling will not wake the core in order, so
                // don't bother having it linked.
                critical.remove_task(task);
            }
            _ => {}
        }
    }

    /// Release any remaining tokens which are associated with this particular task.
    fn release_remaining(
        self: Pin<&mut Self>,
        critical: &mut Guard<'_>,
        state: &mut State<'_>,
        permits: usize,
    ) {
        let (critical, task) = self.get_task(critical);

        if task.is_linked() {
            critical.remove_task(task);
        }

        // Hand back permits which we've acquired so far.
        let release = permits.saturating_sub(task.remaining);
        state.add_tokens(critical, release, |_, _| ());
    }

    /// Drain the given number of tokens through the core. Returns `true` if the
    /// core has been completed.
    #[cfg_attr(
        feature = "tracing",
        tracing::instrument(skip(self, critical), level = "trace")
    )]
    fn drain_core(
        self: Pin<&mut Self>,
        critical: &mut Guard<'_>,
        state: &mut State<'_>,
        tokens: usize,
    ) -> bool {
        let completed = state.add_tokens(critical, tokens, |critical, state| {
            let (_, task) = self.get_task(critical);

            // If the limiter is not fair, we need to in addition to draining
            // remaining tokens from linked nodes, drain it from ourselves. We
            // fill the current holder of the core last (self). To ensure that
            // it stays around for as long as possible.
            if !state.lim.fair {
                task.fill(&mut state.balance);
            }

            task.is_completed()
        });

        if completed {
            // Everything was drained, including the current core (if
            // appropriate). So we can release it now.
            critical.release(state);
        }

        completed
    }

    /// Assume the current core and calculate how long we must sleep for in
    /// order to do it.
    ///
    /// # Safety
    ///
    /// This might link the current task into the task queue, so the caller must
    /// ensure that it is pinned.
    #[cfg_attr(
        feature = "tracing",
        tracing::instrument(skip(self, critical), level = "trace")
    )]
    fn assume_core(
        mut self: Pin<&mut Self>,
        critical: &mut Guard<'_>,
        state: &mut State<'_>,
        now: Instant,
    ) -> bool {
        self.as_mut().link_core(critical, state.lim);

        let (tokens, deadline) = match state.lim.calculate_drain(critical.deadline, now) {
            Some(tokens) => tokens,
            None => return true,
        };

        // It is appropriate to update the deadline.
        critical.deadline = deadline;
        !self.drain_core(critical, state, tokens)
    }
}

pin_project! {
    /// The future associated with acquiring permits from a rate limiter using
    /// [`RateLimiter::acquire`].
    #[project(!Unpin)]
    pub struct Acquire<'a> {
        #[pin]
        inner: AcquireFut<BorrowedRateLimiter<'a>>,
    }
}

impl Acquire<'_> {
    /// Test if this acquire task is currently coordinating the rate limiter.
    ///
    /// # Examples
    ///
    /// ```
    /// use leaky_bucket::RateLimiter;
    /// use std::future::Future;
    /// use std::sync::Arc;
    /// use std::task::Context;
    ///
    /// struct Waker;
    /// # impl std::task::Wake for Waker { fn wake(self: Arc<Self>) { } }
    ///
    /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() {
    /// let limiter = RateLimiter::builder().build();
    ///
    /// let waker = Arc::new(Waker).into();
    /// let mut cx = Context::from_waker(&waker);
    ///
    /// let a1 = limiter.acquire(1);
    /// tokio::pin!(a1);
    ///
    /// assert!(!a1.is_core());
    /// assert!(a1.as_mut().poll(&mut cx).is_pending());
    /// assert!(a1.is_core());
    ///
    /// a1.as_mut().await;
    ///
    /// // After completion this is no longer a core.
    /// assert!(!a1.is_core());
    /// # }
    /// ```
    pub fn is_core(&self) -> bool {
        self.inner.is_core()
    }
}

impl Future for Acquire<'_> {
    type Output = ();

    #[inline]
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.project().inner.poll(cx)
    }
}

pin_project! {
    /// The future associated with acquiring permits from a rate limiter using
    /// [`RateLimiter::acquire_owned`].
    #[project(!Unpin)]
    pub struct AcquireOwned {
        #[pin]
        inner: AcquireFut<Arc<RateLimiter>>,
    }
}

impl AcquireOwned {
    /// Test if this acquire task is currently coordinating the rate limiter.
    ///
    /// # Examples
    ///
    /// ```
    /// use leaky_bucket::RateLimiter;
    /// use std::future::Future;
    /// use std::sync::Arc;
    /// use std::task::Context;
    ///
    /// struct Waker;
    /// # impl std::task::Wake for Waker { fn wake(self: Arc<Self>) { } }
    ///
    /// # #[tokio::main(flavor="current_thread", start_paused=true)] async fn main() {
    /// let limiter = Arc::new(RateLimiter::builder().build());
    ///
    /// let waker = Arc::new(Waker).into();
    /// let mut cx = Context::from_waker(&waker);
    ///
    /// let a1 = limiter.acquire_owned(1);
    /// tokio::pin!(a1);
    ///
    /// assert!(!a1.is_core());
    /// assert!(a1.as_mut().poll(&mut cx).is_pending());
    /// assert!(a1.is_core());
    ///
    /// a1.as_mut().await;
    ///
    /// // After completion this is no longer a core.
    /// assert!(!a1.is_core());
    /// # }
    /// ```
    pub fn is_core(&self) -> bool {
        self.inner.is_core()
    }
}

impl Future for AcquireOwned {
    type Output = ();

    #[inline]
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.project().inner.poll(cx)
    }
}

pin_project! {
    #[project(!Unpin)]
    #[project = AcquireFutProj]
    struct AcquireFut<T>
    where
        T: Deref<Target = RateLimiter>,
    {
        lim: T,
        permits: usize,
        state: AcquireFutState,
        #[pin]
        sleep: Option<time::Sleep>,
        #[pin]
        inner: AcquireFutInner,
    }

    impl<T> PinnedDrop for AcquireFut<T>
    where
        T: Deref<Target = RateLimiter>,
    {
        fn drop(this: Pin<&mut Self>) {
            let AcquireFutProj { lim, permits, state, inner, .. } = this.project();

            let is_core = match *state {
                AcquireFutState::Waiting => false,
                AcquireFutState::Core { .. } => true,
                _ => return,
            };

            let mut critical = lim.lock();
            let mut s = lim.take();
            inner.release_remaining(&mut critical, &mut s, *permits);

            if is_core {
                critical.release(&mut s);
            }

            *state = AcquireFutState::Complete;
        }
    }
}

impl<T> AcquireFut<T>
where
    T: Deref<Target = RateLimiter>,
{
    #[inline]
    const fn new(lim: T, permits: usize) -> Self {
        Self {
            lim,
            permits,
            state: AcquireFutState::Initial,
            sleep: None,
            inner: AcquireFutInner::new(),
        }
    }

    fn is_core(&self) -> bool {
        matches!(&self.state, AcquireFutState::Core { .. })
    }
}

// SAFETY: All the internals of acquire is thread safe and correctly
// synchronized. The embedded waiter queue doesn't have anything inherently
// unsafe in it.
unsafe impl<T> Send for AcquireFut<T> where T: Send + Deref<Target = RateLimiter> {}
unsafe impl<T> Sync for AcquireFut<T> where T: Sync + Deref<Target = RateLimiter> {}

impl<T> Future for AcquireFut<T>
where
    T: Deref<Target = RateLimiter>,
{
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let AcquireFutProj {
            lim,
            permits,
            state,
            mut sleep,
            inner: mut internal,
            ..
        } = self.project();

        // Hold onto the critical lock for core operations, but only acquire it
        // when strictly necessary.
        let mut critical;

        // Shared state.
        //
        // Once we are holding onto the critical lock, we take the entire state
        // to ensure that any fast-past negotiators do not observe any available
        // permits while potential core work is ongoing.
        let mut s;

        // Hold onto any call to `Instant::now` which we might perform, so we
        // don't have to get the current time multiple times.
        let outer_now;

        match *state {
            AcquireFutState::Complete => {
                return Poll::Ready(());
            }
            AcquireFutState::Initial => {
                // If the rate limiter is not fair, try to oppurtunistically
                // just acquire a permit through the known atomic state.
                //
                // This is known as the fast path, but requires acquire to raise
                // against other tasks when storing the state back.
                if lim.try_fast_path(*permits) {
                    *state = AcquireFutState::Complete;
                    return Poll::Ready(());
                }

                critical = lim.lock();
                s = lim.take();

                let now = Instant::now();

                // If we've hit a deadline, calculate the number of tokens
                // to drain and perform it in line here. This is necessary
                // because the core isn't aware of how long we sleep between
                // each acquire, so we need to perform some of the drain
                // work here in order to avoid acruing a debt that needs to
                // be filled later in.
                //
                // If we didn't do this, and the process slept for a long
                // time, the next time a core is acquired it would be very
                // far removed from the expected deadline and has no idea
                // when permits were acquired, so it would over-eagerly
                // release a lot of acquires and accumulate permits.
                //
                // This is tested for in the `test_idle` suite of tests.
                let tokens =
                    if let Some((tokens, deadline)) = lim.calculate_drain(critical.deadline, now) {
                        trace!(tokens, "inline drain");
                        // We pre-emptively update the deadline of the core
                        // since it might bump, and we don't want other
                        // processes to observe that the deadline has been
                        // reached.
                        critical.deadline = deadline;
                        tokens
                    } else {
                        0
                    };

                let completed = s.add_tokens(&mut critical, tokens, |critical, s| {
                    let (_, task) = internal.as_mut().get_task(critical);
                    task.remaining = *permits;
                    task.fill(&mut s.balance);
                    task.is_completed()
                });

                if completed {
                    *state = AcquireFutState::Complete;
                    return Poll::Ready(());
                }

                // Try to take over as core. If we're unsuccessful we just
                // ensure that we're linked into the wait queue.
                if !mem::take(&mut s.available) {
                    internal.as_mut().update(&mut critical, cx.waker());
                    *state = AcquireFutState::Waiting;
                    return Poll::Pending;
                }

                // SAFETY: This is done in a pinned section, so we know that
                // the linked section stays alive for the duration of this
                // future due to pinning guarantees.
                internal.as_mut().link_core(&mut critical, lim);
                Guard::bump(&mut critical);
                *state = AcquireFutState::Core;
                outer_now = Some(now);
            }
            AcquireFutState::Waiting => {
                // If we are complete, then return as ready.
                //
                // This field is atomic, so we can safely read it under shared
                // access and do not require a lock.
                if internal.complete().load(Ordering::Acquire) {
                    *state = AcquireFutState::Complete;
                    return Poll::Ready(());
                }

                // Note: we need to operate under this lock to ensure that
                // the core acquired here (or elsewhere) observes that the
                // current task has been linked up.
                critical = lim.lock();
                s = lim.take();

                // Try to take over as core. If we're unsuccessful we
                // just ensure that we're linked into the wait queue.
                if !mem::take(&mut s.available) {
                    internal.update(&mut critical, cx.waker());
                    return Poll::Pending;
                }

                let now = Instant::now();

                // This is done in a pinned section, so we know that the linked
                // section stays alive for the duration of this future due to
                // pinning guarantees.
                if !internal.as_mut().assume_core(&mut critical, &mut s, now) {
                    // Marks as completed.
                    *state = AcquireFutState::Complete;
                    return Poll::Ready(());
                }

                Guard::bump(&mut critical);
                *state = AcquireFutState::Core;
                outer_now = Some(now);
            }
            AcquireFutState::Core => {
                critical = lim.lock();
                s = lim.take();
                outer_now = None;
            }
        }

        trace!(until = ?critical.deadline, "taking over core and sleeping");

        let mut sleep = match sleep.as_mut().as_pin_mut() {
            Some(mut sleep) => {
                if sleep.deadline() != critical.deadline {
                    sleep.as_mut().reset(critical.deadline);
                }

                sleep
            }
            None => {
                sleep.set(Some(time::sleep_until(critical.deadline)));
                sleep.as_mut().as_pin_mut().unwrap()
            }
        };

        if sleep.as_mut().poll(cx).is_pending() {
            return Poll::Pending;
        }

        critical.deadline = outer_now.unwrap_or_else(Instant::now) + lim.interval;

        if internal.drain_core(&mut critical, &mut s, lim.refill) {
            *state = AcquireFutState::Complete;
            return Poll::Ready(());
        }

        cx.waker().wake_by_ref();
        Poll::Pending
    }
}

#[cfg(test)]
mod tests {
    use super::{Acquire, AcquireOwned, RateLimiter};

    fn is_send<T: Send>() {}
    fn is_sync<T: Sync>() {}

    #[test]
    fn assert_send_sync() {
        is_send::<AcquireOwned>();
        is_sync::<AcquireOwned>();

        is_send::<RateLimiter>();
        is_sync::<RateLimiter>();

        is_send::<Acquire<'_>>();
        is_sync::<Acquire<'_>>();
    }
}