async_safe_defer/

async_scope.rs

use alloc::{boxed::Box, vec::Vec};
use core::{
    fmt,
    future::Future,
    pin::Pin,
    task::{Context, Poll},
};

type LocalFuture<'a> = Pin<Box<dyn Future<Output = ()> + 'a>>;
type SendFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;

fn poll_tasks<F>(tasks: &mut Vec<Pin<Box<F>>>, context: &mut Context<'_>) -> Poll<()>
where
    F: Future<Output = ()> + ?Sized,
{
    loop {
        let Some(mut future) = tasks.pop() else {
            return Poll::Ready(());
        };

        match future.as_mut().poll(context) {
            Poll::Pending => {
                tasks.push(future);
                return Poll::Pending;
            }
            Poll::Ready(()) => {}
        }
    }
}

/// A LIFO stack of local asynchronous cleanup actions.
///
/// Each registration stores one boxed wrapper future. The cleanup factory is
/// invoked only when that wrapper is first polled, and actions are awaited
/// sequentially.
///
/// Dropping a pending [`LocalRun`] leaves its current action in the stack. A
/// later run resumes it after any newer LIFO entries. Dropping the scope itself
/// drops all remaining actions without polling them further.
#[must_use = "call run().await or finish().await to execute registered cleanup actions"]
pub struct LocalAsyncScope<'a> {
    tasks: Vec<LocalFuture<'a>>,
}

/// Drains a [`LocalAsyncScope`] without consuming it.
#[must_use = "futures do nothing unless polled or awaited"]
pub struct LocalRun<'scope, 'task> {
    scope: &'scope mut LocalAsyncScope<'task>,
}

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

    fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
        poll_tasks(&mut self.scope.tasks, context)
    }
}

impl<'a> LocalAsyncScope<'a> {
    /// Creates an empty scope.
    #[inline]
    pub const fn new() -> Self {
        Self { tasks: Vec::new() }
    }

    /// Creates an empty scope with space for at least `capacity` actions.
    #[inline]
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            tasks: Vec::with_capacity(capacity),
        }
    }

    /// Queues a cleanup factory without invoking it.
    #[inline]
    pub fn defer<F, Fut>(&mut self, action: F)
    where
        F: FnOnce() -> Fut + 'a,
        Fut: Future<Output = ()> + 'a,
    {
        self.tasks.push(Box::pin(async move { action().await }));
    }

    /// Returns the number of cleanup actions that have not completed.
    #[inline]
    pub fn len(&self) -> usize {
        self.tasks.len()
    }

    /// Returns the number of actions the scope can hold without reallocating.
    #[inline]
    pub fn capacity(&self) -> usize {
        self.tasks.capacity()
    }

    /// Returns `true` when no cleanup actions remain.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.tasks.is_empty()
    }

    /// Drops all queued actions without polling them further.
    #[inline]
    pub fn clear(&mut self) {
        self.tasks.clear();
    }

    /// Borrows the scope and drains its actions in LIFO order.
    #[inline]
    pub fn run(&mut self) -> LocalRun<'_, 'a> {
        LocalRun { scope: self }
    }

    /// Consumes the scope and runs its actions in LIFO order.
    ///
    /// Dropping the returned future also drops every remaining action.
    #[inline]
    pub async fn finish(mut self) {
        self.run().await;
    }
}

impl Default for LocalAsyncScope<'_> {
    fn default() -> Self {
        Self::new()
    }
}

impl fmt::Debug for LocalAsyncScope<'_> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("LocalAsyncScope")
            .field("pending", &self.tasks.len())
            .finish()
    }
}

/// Convenience name for [`LocalAsyncScope`].
///
/// Use [`SendAsyncScope`] when the scope must cross thread boundaries.
pub use LocalAsyncScope as AsyncScope;

/// A LIFO stack of `Send` asynchronous cleanup actions.
///
/// Deferred closures and the futures they return must implement [`Send`].
/// Execution and cancellation behavior match [`LocalAsyncScope`].
#[must_use = "call run().await or finish().await to execute registered cleanup actions"]
pub struct SendAsyncScope<'a> {
    tasks: Vec<SendFuture<'a>>,
}

/// The `Send` counterpart to [`LocalRun`].
#[must_use = "futures do nothing unless polled or awaited"]
pub struct SendRun<'scope, 'task> {
    scope: &'scope mut SendAsyncScope<'task>,
}

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

    fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
        poll_tasks(&mut self.scope.tasks, context)
    }
}

impl<'a> SendAsyncScope<'a> {
    /// Creates an empty scope.
    #[inline]
    pub const fn new() -> Self {
        Self { tasks: Vec::new() }
    }

    /// Creates an empty scope with space for at least `capacity` actions.
    #[inline]
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            tasks: Vec::with_capacity(capacity),
        }
    }

    /// Queues a `Send` cleanup closure without invoking it.
    #[inline]
    pub fn defer<F, Fut>(&mut self, action: F)
    where
        F: FnOnce() -> Fut + Send + 'a,
        Fut: Future<Output = ()> + Send + 'a,
    {
        self.tasks.push(Box::pin(async move { action().await }));
    }

    /// Returns the number of cleanup actions that have not completed.
    #[inline]
    pub fn len(&self) -> usize {
        self.tasks.len()
    }

    /// Returns the number of actions the scope can hold without reallocating.
    #[inline]
    pub fn capacity(&self) -> usize {
        self.tasks.capacity()
    }

    /// Returns `true` when no cleanup actions remain.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.tasks.is_empty()
    }

    /// Drops all queued actions without polling them further.
    #[inline]
    pub fn clear(&mut self) {
        self.tasks.clear();
    }

    /// Borrows the scope and drains its actions in LIFO order.
    #[inline]
    pub fn run(&mut self) -> SendRun<'_, 'a> {
        SendRun { scope: self }
    }

    /// Consumes the scope and runs its actions in LIFO order.
    ///
    /// Dropping the returned future also drops every remaining action.
    #[inline]
    pub async fn finish(mut self) {
        self.run().await;
    }
}

impl Default for SendAsyncScope<'_> {
    fn default() -> Self {
        Self::new()
    }
}

impl fmt::Debug for SendAsyncScope<'_> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("SendAsyncScope")
            .field("pending", &self.tasks.len())
            .finish()
    }
}