async-safe-defer 0.2.0

#![cfg(feature = "alloc")]

use async_safe_defer::{async_scope, send_async_scope, LocalAsyncScope, SendAsyncScope};
use core::{future::Future, task::Poll};
use std::{
    cell::{Cell, RefCell},
    panic::{catch_unwind, AssertUnwindSafe},
    pin::Pin,
    rc::Rc,
    sync::{
        atomic::{AtomicUsize, Ordering},
        Arc, Mutex,
    },
    task::{Context, Wake, Waker},
};

struct NoopWake;

impl Wake for NoopWake {
    fn wake(self: Arc<Self>) {}
}

fn poll_once<F: Future>(future: Pin<&mut F>) -> Poll<F::Output> {
    let waker = Waker::from(Arc::new(NoopWake));
    future.poll(&mut Context::from_waker(&waker))
}

fn assert_send<T: Send>(_: &T) {}

struct DropProbe<'a>(&'a Cell<usize>);

impl Drop for DropProbe<'_> {
    fn drop(&mut self) {
        self.0.set(self.0.get() + 1);
    }
}

#[tokio::test]
async fn local_scope_is_lifo_and_accepts_non_send_borrows() {
    let order = Rc::new(RefCell::new(Vec::new()));
    let mut scope = LocalAsyncScope::new();

    let first = Rc::clone(&order);
    scope.defer(move || async move { first.borrow_mut().push(1) });
    let second = Rc::clone(&order);
    scope.defer(move || async move { second.borrow_mut().push(2) });

    assert_eq!(scope.len(), 2);
    scope.run().await;
    assert!(scope.is_empty());
    assert_eq!(*order.borrow(), vec![2, 1]);
}

#[tokio::test]
async fn cleanup_factory_is_lazy_until_run() {
    let factory_calls = Cell::new(0);
    let mut scope = LocalAsyncScope::with_capacity(1);
    assert!(scope.capacity() >= 1);
    scope.defer(|| {
        factory_calls.set(factory_calls.get() + 1);
        async {}
    });

    assert_eq!(factory_calls.get(), 0);
    scope.run().await;
    assert_eq!(factory_calls.get(), 1);
}

#[tokio::test]
async fn finish_consumes_and_drains_the_scope() {
    let local_completed = Cell::new(false);
    let mut local = LocalAsyncScope::new();
    local.defer(|| async { local_completed.set(true) });
    local.finish().await;
    assert!(local_completed.get());

    let send_completed = Arc::new(Mutex::new(false));
    let captured = Arc::clone(&send_completed);
    let mut send = SendAsyncScope::new();
    send.defer(move || async move { *captured.lock().unwrap() = true });
    let finish = send.finish();
    assert_send(&finish);
    finish.await;
    assert!(*send_completed.lock().unwrap());
}

#[tokio::test]
async fn macro_preserves_body_value_and_runs_on_inner_return() {
    let cleaned = Cell::new(false);
    let value = async_scope!(scope, {
        let cleaned = &cleaned;
        scope.defer(move || async move { cleaned.set(true) });
        return 42;
    })
    .await;

    assert_eq!(value, 42);
    assert!(cleaned.get());
}

async fn fail_after_registering(cleaned: &Cell<bool>) -> Result<(), &'static str> {
    async_scope!(scope, {
        scope.defer(|| async { cleaned.set(true) });
        Err::<(), _>("body failed")?;
        Ok(())
    })
    .await
}

#[tokio::test]
async fn macro_runs_cleanup_before_propagating_question_mark() {
    let cleaned = Cell::new(false);
    assert_eq!(fail_after_registering(&cleaned).await, Err("body failed"));
    assert!(cleaned.get());
}

#[tokio::test]
async fn macro_can_move_body_local_state_into_cleanup() {
    let output = Rc::new(RefCell::new(Vec::new()));
    let captured = Rc::clone(&output);

    async_scope!(scope, {
        let value = String::from("body-owned");
        scope.defer(move || async move {
            captured.borrow_mut().push(value);
        });
    })
    .await;

    assert_eq!(&*output.borrow(), &["body-owned"]);
}

#[tokio::test]
async fn send_macro_builds_a_send_future() {
    let output = Arc::new(Mutex::new(Vec::new()));
    let captured = Arc::clone(&output);
    let future = send_async_scope!(scope, {
        scope.defer(move || async move {
            captured.lock().unwrap().push(1);
        });
        7
    });

    assert_send(&future);
    assert_eq!(future.await, 7);
    assert_eq!(*output.lock().unwrap(), vec![1]);
}

#[tokio::test]
async fn dropping_run_future_preserves_in_progress_cleanup() {
    let polls = Cell::new(0);
    let completed = Cell::new(false);
    let order = RefCell::new(Vec::new());
    let mut scope = LocalAsyncScope::new();
    scope.defer(|| async {
        core::future::poll_fn(|context| {
            let count = polls.get();
            polls.set(count + 1);
            if count == 0 {
                context.waker().wake_by_ref();
                Poll::Pending
            } else {
                Poll::Ready(())
            }
        })
        .await;
        order.borrow_mut().push(1);
        completed.set(true);
    });

    let mut first_run = Box::pin(scope.run());
    assert!(poll_once(first_run.as_mut()).is_pending());
    drop(first_run);
    assert_eq!(scope.len(), 1);

    scope.defer(|| async { order.borrow_mut().push(2) });
    scope.run().await;
    assert!(scope.is_empty());
    assert_eq!(polls.get(), 2);
    assert!(completed.get());
    assert_eq!(*order.borrow(), vec![2, 1]);
}

#[tokio::test]
async fn manual_scope_can_continue_after_cleanup_panics() {
    let completed = Cell::new(false);
    let mut scope = LocalAsyncScope::new();
    scope.defer(|| async { completed.set(true) });
    scope.defer(|| async { panic!("cleanup failed") });

    let outcome = catch_unwind(AssertUnwindSafe(|| {
        let mut run = scope.run();
        let _ = poll_once(Pin::new(&mut run));
    }));

    assert!(outcome.is_err());
    assert_eq!(scope.len(), 1);
    scope.run().await;
    assert!(completed.get());
}

#[tokio::test]
async fn send_run_is_send_and_resumes_after_newer_cleanup() {
    let polls = Arc::new(AtomicUsize::new(0));
    let order = Arc::new(Mutex::new(Vec::new()));
    let mut scope = SendAsyncScope::new();

    let current_polls = Arc::clone(&polls);
    let current_order = Arc::clone(&order);
    scope.defer(move || async move {
        core::future::poll_fn(|context| {
            if current_polls.fetch_add(1, Ordering::SeqCst) == 0 {
                context.waker().wake_by_ref();
                Poll::Pending
            } else {
                Poll::Ready(())
            }
        })
        .await;
        current_order.lock().unwrap().push(1);
    });

    let mut first_run = Box::pin(scope.run());
    assert_send(&first_run);
    assert!(poll_once(first_run.as_mut()).is_pending());
    drop(first_run);

    let newer_order = Arc::clone(&order);
    scope.defer(move || async move { newer_order.lock().unwrap().push(2) });
    let run = scope.run();
    assert_send(&run);
    run.await;

    assert_eq!(polls.load(Ordering::SeqCst), 2);
    assert_eq!(*order.lock().unwrap(), vec![2, 1]);
}

#[test]
fn dropping_macro_future_discards_registered_cleanup() {
    let cleaned = Cell::new(false);
    let future = async_scope!(scope, {
        scope.defer(|| async { cleaned.set(true) });
        core::future::pending::<()>().await;
    });
    let mut future = Box::pin(future);

    assert!(poll_once(future.as_mut()).is_pending());
    drop(future);
    assert!(!cleaned.get());
}

#[test]
fn body_panic_discards_registered_macro_cleanup() {
    let cleanup_ran = Cell::new(false);
    let action_drops = Cell::new(0);
    let future = async_scope!(scope, {
        let cleanup_ran = &cleanup_ran;
        let probe = DropProbe(&action_drops);
        scope.defer(move || async move {
            drop(probe);
            cleanup_ran.set(true);
        });
        panic!("body failed");
    });
    let mut future = Box::pin(future);

    let outcome = catch_unwind(AssertUnwindSafe(|| {
        let _ = poll_once(future.as_mut());
    }));
    assert!(outcome.is_err());
    drop(future);
    assert!(!cleanup_ran.get());
    assert_eq!(action_drops.get(), 1);
}

#[test]
fn cleanup_panic_discards_older_macro_actions() {
    let older_ran = Cell::new(false);
    let action_drops = Cell::new(0);
    let future = async_scope!(scope, {
        let older_ran = &older_ran;
        let probe = DropProbe(&action_drops);
        scope.defer(move || async move {
            drop(probe);
            older_ran.set(true);
        });
        scope.defer(|| async { panic!("cleanup failed") });
    });
    let mut future = Box::pin(future);

    let outcome = catch_unwind(AssertUnwindSafe(|| {
        let _ = poll_once(future.as_mut());
    }));
    assert!(outcome.is_err());
    drop(future);
    assert!(!older_ran.get());
    assert_eq!(action_drops.get(), 1);
}

#[test]
fn dropping_finish_discards_current_and_older_cleanup() {
    let current_ran = Cell::new(false);
    let older_ran = Cell::new(false);
    let action_drops = Cell::new(0);
    let mut scope = LocalAsyncScope::new();
    let probe = DropProbe(&action_drops);
    let older_ran_ref = &older_ran;
    scope.defer(|| async move {
        drop(probe);
        older_ran_ref.set(true);
    });
    scope.defer(|| async {
        core::future::pending::<()>().await;
        current_ran.set(true);
    });
    let mut finish = Box::pin(scope.finish());

    assert!(poll_once(finish.as_mut()).is_pending());
    drop(finish);
    assert!(!current_ran.get());
    assert!(!older_ran.get());
    assert_eq!(action_drops.get(), 1);
}

#[test]
fn scopes_support_default_clear_and_debug() {
    let mut local = LocalAsyncScope::default();
    local.defer(|| async {});
    assert_eq!(format!("{local:?}"), "LocalAsyncScope { pending: 1 }");
    local.clear();
    assert!(local.is_empty());

    let mut send = SendAsyncScope::default();
    assert_eq!(send.capacity(), 0);
    send.defer(|| async {});
    assert_eq!(format!("{send:?}"), "SendAsyncScope { pending: 1 }");
    send.clear();
    assert!(send.is_empty());
}