async-safe-defer 0.2.0

use async_safe_defer::fixed::{FixedAsyncScope, FixedSendAsyncScope};
use core::{future::Future, pin::pin, task::Poll};
use std::{
    cell::{Cell, RefCell},
    panic::{catch_unwind, AssertUnwindSafe},
    pin::Pin,
    sync::{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) {}

#[tokio::test]
async fn fixed_scope_runs_heterogeneous_futures_in_lifo_order() {
    let order = RefCell::new(Vec::new());
    let mut first = pin!(async { order.borrow_mut().push(1) });
    let mut second = pin!(async {
        core::future::ready(()).await;
        order.borrow_mut().push(2);
    });
    let mut scope = FixedAsyncScope::<2>::new();

    scope.try_defer(first.as_mut()).unwrap();
    scope.try_defer(second.as_mut()).unwrap();
    assert_eq!(scope.len(), 2);
    assert_eq!(scope.capacity(), 2);
    assert!(scope.is_full());

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

#[tokio::test]
async fn finish_consumes_and_drains_the_fixed_scope() {
    let completed = Cell::new(false);
    let mut cleanup = pin!(async { completed.set(true) });
    let mut scope = FixedAsyncScope::<1>::new();
    scope.try_defer(cleanup.as_mut()).unwrap();

    scope.finish().await;
    assert!(completed.get());
}

#[tokio::test]
async fn capacity_error_returns_rejected_future() {
    let ran = Cell::new(0);
    let mut first = pin!(async { ran.set(ran.get() + 1) });
    let mut rejected = pin!(async { ran.set(ran.get() + 10) });
    let mut scope = FixedAsyncScope::<1>::new();

    scope.try_defer(first.as_mut()).unwrap();
    let rejected = scope.try_defer(rejected.as_mut()).unwrap_err().into_inner();
    scope.clear();
    scope.try_defer(rejected).unwrap();
    scope.run().await;
    drop(scope);

    assert_eq!(ran.get(), 10);
}

#[test]
fn zero_capacity_is_non_panicking() {
    let mut future = pin!(async {});
    let mut scope = FixedAsyncScope::<0>::new();

    assert!(scope.is_empty());
    assert!(scope.is_full());
    assert_eq!(scope.capacity(), 0);
    assert!(scope.try_defer(future.as_mut()).is_err());
}

#[tokio::test]
async fn dropping_run_preserves_current_and_runs_newer_future_first() {
    let polls = Cell::new(0);
    let completed = Cell::new(false);
    let order = RefCell::new(Vec::new());
    let mut cleanup = pin!(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 newer = pin!(async { order.borrow_mut().push(2) });
    let mut scope = FixedAsyncScope::<2>::new();
    scope.try_defer(cleanup.as_mut()).unwrap();

    let mut first_run = scope.run();
    assert!(poll_once(Pin::new(&mut first_run)).is_pending());
    drop(first_run);
    assert_eq!(scope.len(), 1);

    scope.try_defer(newer.as_mut()).unwrap();
    scope.run().await;
    assert!(scope.is_empty());
    drop(scope);
    assert_eq!(polls.get(), 2);
    assert!(completed.get());
    assert_eq!(*order.borrow(), vec![2, 1]);
}

#[tokio::test]
async fn fixed_send_scope_runs_send_futures_in_lifo_order() {
    let order = Arc::new(Mutex::new(Vec::new()));
    let first_order = Arc::clone(&order);
    let second_order = Arc::clone(&order);
    let mut first = pin!(async move { first_order.lock().unwrap().push(1) });
    let mut second = pin!(async move { second_order.lock().unwrap().push(2) });
    let mut scope = FixedSendAsyncScope::<2>::new();

    scope.try_defer(first.as_mut()).unwrap();
    scope.try_defer(second.as_mut()).unwrap();
    assert_send(&scope);

    let run = scope.run();
    assert_send(&run);
    run.await;

    assert_eq!(*order.lock().unwrap(), vec![2, 1]);
}

#[test]
fn fixed_send_scope_runner_and_finish_future_are_send() {
    let mut cleanup = pin!(async {});
    let mut scope = FixedSendAsyncScope::<1>::new();
    scope.try_defer(cleanup.as_mut()).unwrap();
    assert_send(&scope);

    let run = scope.run();
    assert_send(&run);
    drop(run);

    let finish = scope.finish();
    assert_send(&finish);
    drop(finish);
}

#[test]
fn dropping_finish_releases_the_caller_owned_future() {
    let polls = Cell::new(0);
    let completed = Cell::new(false);
    let mut cleanup = pin!(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;
        completed.set(true);
    });
    let mut scope = FixedAsyncScope::<1>::new();
    scope.try_defer(cleanup.as_mut()).unwrap();
    let mut finish = Box::pin(scope.finish());

    assert!(poll_once(finish.as_mut()).is_pending());
    drop(finish);
    assert!(poll_once(cleanup.as_mut()).is_ready());
    assert!(completed.get());
}

#[tokio::test]
async fn fixed_scope_can_continue_after_cleanup_panics() {
    let completed = Cell::new(false);
    let mut older = pin!(async { completed.set(true) });
    let mut panicking = pin!(async { panic!("cleanup failed") });
    let mut scope = FixedAsyncScope::<2>::new();
    scope.try_defer(older.as_mut()).unwrap();
    scope.try_defer(panicking.as_mut()).unwrap();

    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;
    drop(scope);
    assert!(completed.get());
}

#[test]
fn fixed_scope_supports_default_clear_and_debug() {
    let mut scope = FixedAsyncScope::<3>::default();
    assert_eq!(
        format!("{scope:?}"),
        "FixedAsyncScope { pending: 0, capacity: 3 }"
    );
    scope.clear();
    assert!(scope.is_empty());
}