core/test/rust/runtime/task/group.rs

use super::{DropNotice, runtime};
use dever_runtime::task::{Error, Group, Limits, Runtime, Task};
use std::{
    future::{Future, poll_fn},
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
    },
    task::Poll,
    time::Duration,
};

type Tasks = Group<Task<Result<(), &'static str>>, &'static str>;

#[test]
fn group_limits_are_checked_and_empty_groups_complete() {
    for limit in [-1, 0, 65537] {
        assert!(matches!(Tasks::new(limit), Err(Error::InvalidGroupLimit)));
    }
    runtime()
        .run(async |_| {
            Tasks::new(65536).unwrap().wait().await.unwrap().unwrap();
            Tasks::new(1).unwrap().stop().await.unwrap().unwrap();
        })
        .unwrap();
}

#[test]
fn group_capacity_defers_the_next_factory() {
    let called = Arc::new(AtomicBool::new(false));
    runtime()
        .run(async |scope| {
            let mut group = Tasks::new(1).unwrap();
            let (release, ready) = tokio::sync::oneshot::channel();
            group
                .run(|| {
                    scope.spawn_leaf(|| async {
                        ready.await.unwrap();
                        Ok(())
                    })
                })
                .await
                .unwrap()
                .unwrap();
            {
                let mut next = std::pin::pin!(group.run(|| {
                    called.store(true, Ordering::SeqCst);
                    scope.spawn_leaf(|| async { Ok(()) })
                }));
                poll_fn(|context| {
                    assert!(next.as_mut().poll(context).is_pending());
                    Poll::Ready(())
                })
                .await;
                assert!(!called.load(Ordering::SeqCst));
                release.send(()).unwrap();
                next.await.unwrap().unwrap();
            }
            assert!(called.load(Ordering::SeqCst));
            group.wait().await.unwrap().unwrap();
        })
        .unwrap();
}

#[test]
fn first_group_failure_cancels_and_drains_other_tasks() {
    let dropped = Arc::new(AtomicBool::new(false));
    let observed = dropped.clone();
    runtime()
        .run(async |scope| {
            let mut group = Tasks::new(2).unwrap();
            let owner = DropNotice(dropped);
            group
                .run(|| {
                    scope.spawn_leaf(move || async move {
                        let _owner = owner;
                        std::future::pending::<Result<(), &'static str>>().await
                    })
                })
                .await
                .unwrap()
                .unwrap();
            group
                .run(|| scope.spawn_leaf(|| async { Err("first failure") }))
                .await
                .unwrap()
                .unwrap();
            assert_eq!(group.wait().await, Ok(Err("first failure")));
            assert!(observed.load(Ordering::SeqCst));
        })
        .unwrap();
}

#[test]
fn failed_group_never_invokes_a_waiting_factory() {
    let called = Arc::new(AtomicBool::new(false));
    runtime()
        .run(async |scope| {
            let mut group = Tasks::new(1).unwrap();
            group
                .run(|| scope.spawn_leaf(|| async { Err("first failure") }))
                .await
                .unwrap()
                .unwrap();
            let result = group
                .run(|| {
                    called.store(true, Ordering::SeqCst);
                    scope.spawn_leaf(|| async { Ok(()) })
                })
                .await;
            assert_eq!(result, Ok(Err("first failure")));
            assert!(!called.load(Ordering::SeqCst));
            assert_eq!(
                group.run(|| scope.spawn_leaf(|| async { Ok(()) })).await,
                Err(Error::Cancelled)
            );
            group.wait().await.unwrap().unwrap();
        })
        .unwrap();
}

#[test]
fn stopped_group_preserves_late_blocking_failures() {
    runtime()
        .run(async |scope| {
            let mut group = Tasks::new(1).unwrap();
            let (started, ready) = tokio::sync::oneshot::channel();
            let (release, blocked) = std::sync::mpsc::channel();
            group
                .run(|| {
                    scope.spawn_sync(move || {
                        started.send(()).unwrap();
                        blocked.recv_timeout(Duration::from_secs(2)).unwrap();
                        Err("late failure")
                    })
                })
                .await
                .unwrap()
                .unwrap();
            ready.await.unwrap();
            let mut stopped = std::pin::pin!(group.stop());
            poll_fn(|context| {
                assert!(stopped.as_mut().poll(context).is_pending());
                Poll::Ready(())
            })
            .await;
            release.send(()).unwrap();
            assert_eq!(stopped.await, Ok(Err("late failure")));
        })
        .unwrap();
}

#[test]
fn admission_failure_drains_existing_group_work() {
    let runtime = Runtime::new(Limits {
        threads: 2,
        workers: 1,
        tasks: 1,
    })
    .unwrap();
    let dropped = Arc::new(AtomicBool::new(false));
    let observed = dropped.clone();
    runtime
        .run(async |scope| {
            let mut group = Tasks::new(2).unwrap();
            let owner = DropNotice(dropped);
            group
                .run(|| {
                    scope.spawn_leaf(move || async move {
                        let _owner = owner;
                        std::future::pending::<Result<(), &'static str>>().await
                    })
                })
                .await
                .unwrap()
                .unwrap();
            assert_eq!(
                group.run(|| scope.spawn_leaf(|| async { Ok(()) })).await,
                Err(Error::Capacity)
            );
            assert!(observed.load(Ordering::SeqCst));
            assert_eq!(
                scope
                    .spawn_leaf(|| async { 42 })
                    .unwrap()
                    .wait()
                    .await
                    .unwrap(),
                42
            );
            group.wait().await.unwrap().unwrap();
        })
        .unwrap();
}

#[test]
fn dropping_a_group_cancels_unclaimed_work() {
    let dropped = Arc::new(AtomicBool::new(false));
    let observed = dropped.clone();
    runtime()
        .run(async |scope| {
            let mut group = Tasks::new(2).unwrap();
            let owner = DropNotice(dropped);
            group
                .run(|| {
                    scope.spawn_leaf(move || async move {
                        let _owner = owner;
                        std::future::pending::<Result<(), &'static str>>().await
                    })
                })
                .await
                .unwrap()
                .unwrap();
            drop(group);
            scope.finish().await;
            assert!(observed.load(Ordering::SeqCst));
        })
        .unwrap();
}