Quelle  task_join_set.rs

#![warn(rust_2018_idioms)]
#![cfg(feature = "full")]

use futures::future::{pending, FutureExt};
use std::panic;
use tokio::sync::oneshot;
use tokio::task::{JoinSet, LocalSet};
use tokio::time::Duration;

fn rt() -> tokio::runtime::Runtime {
    tokio::runtime::Builder::new_current_thread()
        .build()
        .unwrap()
}

// Spawn `N` tasks that return their index (`i`).
fn spawn_index_tasks(set: &mut JoinSet<usize>, n: usize, on: Option<&LocalSet&gt;) {
    for i in 0..n {
        let rc = std::rc::Rc::new(i);
        match on {
            None => set.spawn_local(async move { *rc }),
            Some(local) => set.spawn_local_on(async move { *rc }, local),
        };
    }
}

// Spawn `N` “pending” tasks that own a `oneshot::Sender`.
// When the task is aborted the sender is dropped, which is observed
// via the returned `Receiver`s.
fn spawn_pending_tasks(
    set: &mut JoinSet<()>,
    receivers: &mut Vec<oneshot::Receiver<()>>,
    n: usize,
    on: Option<&LocalSet>,
) {
    for _ in 0..n {
        let (tx, rx) = oneshot::channel::<()>();
        receivers.push(rx);

        let fut = async move {
            pending::<()>().await;
            drop(tx);
        };

        match on {
            None => set.spawn_local(fut),
            Some(local) => set.spawn_local_on(fut, local),
        };
    }
}

// Await every task in a JoinSet and assert every task returns its own index.
async fn drain_joinset_and_assert(mut set: JoinSet<usize>, n: usize) {
    let mut seen = vec![false; n];
    while let Some(res) = set.join_next().await {
        let idx = res.expect("task panicked");
        seen[idx] = true;
    }
    assert!(seen.into_iter().all(|b| b));
    assert!(set.is_empty());
}

// Await every receiver and assert they all return `Err` because the
// corresponding sender (inside an aborted task) was dropped.
async fn await_receivers_and_assert(receivers: Vec<oneshot::Receiver<()>>) {
    for rx in receivers {
        assert!(
            rx.await.is_err(),
            "the task should have been aborted and the sender dropped"
        );
    }
}

#[tokio::test(start_paused = true)]
async fn test_with_sleep() {
    let mut set = JoinSet::new();

    for i in 0..10 {
        set.spawn(async move { i });
        assert_eq!(set.len(), 1 + i);
    }
    set.detach_all();
    assert_eq!(set.len(), 0);

    assert!(set.join_next().await.is_none());

    for i in 0..10 {
        set.spawn(async move {
            tokio::time::sleep(Duration::from_secs(i as u64)).await;
            i
        });
        assert_eq!(set.len(), 1 + i);
    }

    let mut seen = [false; 10];
    while let Some(res) = set.join_next().await.transpose().unwrap() {
        seen[res] = true;
    }

    for was_seen in &seen {
        assert!(was_seen);
    }
    assert!(set.join_next().await.is_none());

    // Do it again.
    for i in 0..10 {
        set.spawn(async move {
            tokio::time::sleep(Duration::from_secs(i as u64)).await;
            i
        });
    }

    let mut seen = [false; 10];
    while let Some(res) = set.join_next().await.transpose().unwrap() {
        seen[res] = true;
    }

    for was_seen in &seen {
        assert!(was_seen);
    }
    assert!(set.join_next().await.is_none());
}

#[tokio::test]
async fn test_abort_on_drop() {
    let mut set = JoinSet::new();

    let mut recvs = Vec::new();

    for _ in 0..16 {
        let (send, recv) = oneshot::channel::<()>();
        recvs.push(recv);

        set.spawn(async {
            // This task will never complete on its own.
            futures::future::pending::<()>().await;
            drop(send);
        });
    }

    drop(set);

    for recv in recvs {
        // The task is aborted soon and we will receive an error.
        assert!(recv.await.is_err());
    }
}

#[tokio::test]
async fn alternating() {
    let mut set = JoinSet::new();

    assert_eq!(set.len(), 0);
    set.spawn(async {});
    assert_eq!(set.len(), 1);
    set.spawn(async {});
    assert_eq!(set.len(), 2);

    for _ in 0..16 {
        let () = set.join_next().await.unwrap().unwrap();
        assert_eq!(set.len(), 1);
        set.spawn(async {});
        assert_eq!(set.len(), 2);
    }
}

#[tokio::test(start_paused = true)]
async fn abort_tasks() {
    let mut set = JoinSet::new();
    let mut num_canceled = 0;
    let mut num_completed = 0;
    for i in 0..16 {
        let abort = set.spawn(async move {
            tokio::time::sleep(Duration::from_secs(i as u64)).await;
            i
        });

        if i % 2 != 0 {
            // abort odd-numbered tasks.
            abort.abort();
        }
    }
    loop {
        match set.join_next().await {
            Some(Ok(res)) => {
                num_completed += 1;
                assert_eq!(res % 2, 0);
            }
            Some(Err(e)) => {
                assert!(e.is_cancelled());
                num_canceled += 1;
            }
            None => break,
        }
    }

    assert_eq!(num_canceled, 8);
    assert_eq!(num_completed, 8);
}

#[test]
fn runtime_gone() {
    let mut set = JoinSet::new();
    {
        let rt = rt();
        set.spawn_on(async { 1 }, rt.handle());
        drop(rt);
    }

    assert!(rt()
        .block_on(set.join_next())
        .unwrap()
        .unwrap_err()
        .is_cancelled());
}

#[tokio::test]
async fn join_all() {
    let mut set: JoinSet<i32> = JoinSet::new();

    for _ in 0..5 {
        set.spawn(async { 1 });
    }
    let res: Vec<i32> = set.join_all().await;

    assert_eq!(res.len(), 5);
    for itm in res.into_iter() {
        assert_eq!(itm, 1)
    }
}

#[cfg(panic = "unwind")]
#[tokio::test(start_paused = true)]
async fn task_panics() {
    let mut set: JoinSet<()> = JoinSet::new();

    let (tx, mut rx) = oneshot::channel();
    assert_eq!(set.len(), 0);

    set.spawn(async move {
        tokio::time::sleep(Duration::from_secs(2)).await;
        tx.send(()).unwrap();
    });
    assert_eq!(set.len(), 1);

    set.spawn(async {
        tokio::time::sleep(Duration::from_secs(1)).await;
        panic!();
    });
    assert_eq!(set.len(), 2);

    let panic = tokio::spawn(set.join_all()).await.unwrap_err();
    assert!(rx.try_recv().is_err());
    assert!(panic.is_panic());
}

#[tokio::test(start_paused = true)]
async fn abort_all() {
    let mut set: JoinSet<()> = JoinSet::new();

    for _ in 0..5 {
        set.spawn(futures::future::pending());
    }
    for _ in 0..5 {
        set.spawn(async {
            tokio::time::sleep(Duration::from_secs(1)).await;
        });
    }

    // The join set will now have 5 pending tasks and 5 ready tasks.
    tokio::time::sleep(Duration::from_secs(2)).await;

    set.abort_all();
    assert_eq!(set.len(), 10);

    let mut count = 0;
    while let Some(res) = set.join_next().await {
        if let Err(err) = res {
            assert!(err.is_cancelled());
        }
        count += 1;
    }
    assert_eq!(count, 10);
    assert_eq!(set.len(), 0);
}

// This ensures that `join_next` works correctly when the coop budget is
// exhausted.
#[tokio::test(flavor = "current_thread")]
async fn join_set_coop() {
    // Large enough to trigger coop.
    const TASK_NUM: u32 = 1000;

    static SEM: tokio::sync::Semaphore = tokio::sync::Semaphore::const_new(0);

    let mut set = JoinSet::new();

    for _ in 0..TASK_NUM {
        set.spawn(async {
            SEM.add_permits(1);
        });
    }

    // Wait for all tasks to complete.
    //
    // Since this is a `current_thread` runtime, there's no race condition
    // between the last permit being added and the task completing.
    let _ = SEM.acquire_many(TASK_NUM).await.unwrap();

    let mut count = 0;
    let mut coop_count = 0;
    loop {
        match set.join_next().now_or_never() {
            Some(Some(Ok(()))) => {}
            Some(Some(Err(err))) => panic!("failed: {err}"),
            None => {
                coop_count += 1;
                tokio::task::yield_now().await;
                continue;
            }
            Some(None) => break,
        }

        count += 1;
    }
    assert!(coop_count >= 1);
    assert_eq!(count, TASK_NUM);
}

#[tokio::test(flavor = "current_thread")]
async fn try_join_next() {
    const TASK_NUM: u32 = 1000;

    let (send, recv) = tokio::sync::watch::channel(());

    let mut set = JoinSet::new();

    for _ in 0..TASK_NUM {
        let mut recv = recv.clone();
        set.spawn(async move { recv.changed().await.unwrap() });
    }
    drop(recv);

    assert!(set.try_join_next().is_none());

    send.send_replace(());
    send.closed().await;

    let mut count = 0;
    loop {
        match set.try_join_next() {
            Some(Ok(())) => {
                count += 1;
            }
            Some(Err(err)) => panic!("failed: {err}"),
            None => {
                break;
            }
        }
    }

    assert_eq!(count, TASK_NUM);
}

#[cfg(tokio_unstable)]
#[tokio::test(flavor = "current_thread")]
async fn try_join_next_with_id() {
    const TASK_NUM: u32 = 1000;

    let (send, recv) = tokio::sync::watch::channel(());

    let mut set = JoinSet::new();
    let mut spawned = std::collections::HashSet::with_capacity(TASK_NUM as usize);

    for _ in 0..TASK_NUM {
        let mut recv = recv.clone();
        let handle = set.spawn(async move { recv.changed().await.unwrap() });

        spawned.insert(handle.id());
    }
    drop(recv);

    assert!(set.try_join_next_with_id().is_none());

    send.send_replace(());
    send.closed().await;

    let mut count = 0;
    let mut joined = std::collections::HashSet::with_capacity(TASK_NUM as usize);
    loop {
        match set.try_join_next_with_id() {
            Some(Ok((id, ()))) => {
                count += 1;
                joined.insert(id);
            }
            Some(Err(err)) => panic!("failed: {err}"),
            None => {
                break;
            }
        }
    }

    assert_eq!(count, TASK_NUM);
    assert_eq!(joined, spawned);
}

#[tokio::test]
async fn extend() {
    let mut set: JoinSet<_> = (0..5).map(|i| async move { i }).collect();

    set.extend((5..10).map(|i| async move { i }));

    let mut seen = [false; 10];
    while let Some(res) = set.join_next().await {
        let idx = res.unwrap();
        seen[idx] = true;
    }

    for s in &seen {
        assert!(s);
    }
}

mod spawn_local {
    use super::*;

    #[test]
    #[should_panic(
        expected = "`spawn_local` called from outside of a `task::LocalSet` or `runtime::LocalRuntime`"
    )]
    fn panic_outside_any_runtime() {
        let mut set = JoinSet::new();
        set.spawn_local(async {});
    }

    #[tokio::test(flavor = "multi_thread")]
    #[should_panic(
        expected = "`spawn_local` called from outside of a `task::LocalSet` or `runtime::LocalRuntime`"
    )]
    async fn panic_in_multi_thread_runtime() {
        let mut set = JoinSet::new();
        set.spawn_local(async {});
    }

    #[cfg(tokio_unstable)]
    mod local_runtime {
        use super::*;

        /// Spawn several tasks, and then join all tasks.
        #[tokio::test(flavor = "local")]
        async fn spawn_then_join_next() {
            const N: usize = 8;

            let mut set = JoinSet::new();
            spawn_index_tasks(&mut set, N, None);

            assert!(set.try_join_next().is_none());
            drain_joinset_and_assert(set, N).await;
        }

        /// Spawn several pending-forever tasks, and then shutdown the [`JoinSet`].
        #[tokio::test(flavor = "local")]
        async fn spawn_then_shutdown() {
            const N: usize = 8;

            let mut set = JoinSet::new();
            let mut receivers = Vec::new();

            spawn_pending_tasks(&mut set, &mut receivers, N, None);

            assert!(set.try_join_next().is_none());
            set.shutdown().await;
            assert!(set.is_empty());

            await_receivers_and_assert(receivers).await;
        }

        /// Spawn several pending-forever tasks, and then drop the [`JoinSet`].
        #[tokio::test(flavor = "local")]
        async fn spawn_then_drop() {
            const N: usize = 8;
            let mut set = JoinSet::new();
            let mut receivers = Vec::new();

            spawn_pending_tasks(&mut set, &mut receivers, N, None);

            assert!(set.try_join_next().is_none());
            drop(set);

            await_receivers_and_assert(receivers).await;
        }
    }

    mod local_set {
        use super::*;

        /// Spawn several tasks, and then join all tasks.
        #[tokio::test(flavor = "current_thread")]
        async fn spawn_then_join_next() {
            const N: usize = 8;
            let local = LocalSet::new();

            local
                .run_until(async move {
                    let mut set = JoinSet::new();
                    spawn_index_tasks(&mut set, N, None);
                    drain_joinset_and_assert(set, N).await;
                })
                .await;
        }

        /// Spawn several pending-forever tasks, and then shutdown the [`JoinSet`].
        #[tokio::test(flavor = "current_thread")]
        async fn spawn_then_shutdown() {
            const N: usize = 8;
            let local = LocalSet::new();

            local
                .run_until(async {
                    let mut set = JoinSet::new();
                    let mut receivers = Vec::new();

                    spawn_pending_tasks(&mut set, &mut receivers, N, None);
                    assert!(set.try_join_next().is_none());

                    set.shutdown().await;
                    assert!(set.is_empty());

                    await_receivers_and_assert(receivers).await;
                })
                .await;
        }

        /// Spawn several pending-forever tasks, and then drop the [`JoinSet`].
        #[tokio::test(flavor = "current_thread")]
        async fn spawn_then_drop() {
            const N: usize = 8;
            let local = LocalSet::new();

            local
                .run_until(async {
                    let mut set = JoinSet::new();
                    let mut receivers = Vec::new();

                    spawn_pending_tasks(&mut set, &mut receivers, N, None);
                    assert!(set.try_join_next().is_none());

                    drop(set);
                    await_receivers_and_assert(receivers).await;
                })
                .await;
        }
    }
}

mod spawn_local_on {
    use super::*;

    #[cfg(tokio_unstable)]
    mod local_runtime {
        use super::*;

        /// Spawn several tasks, and then join all tasks.
        #[tokio::test(flavor = "local")]
        async fn spawn_then_join_next() {
            const N: usize = 8;

            let local = LocalSet::new();
            let mut set = JoinSet::new();

            spawn_index_tasks(&mut set, N, Some(&local));
            assert!(set.try_join_next().is_none());

            local
                .run_until(async move {
                    drain_joinset_and_assert(set, N).await;
                })
                .await;
        }
    }

    mod local_set {
        use super::*;

        /// Spawn several tasks, and then join all tasks.
        #[tokio::test(flavor = "current_thread")]
        async fn spawn_then_join_next() {
            const N: usize = 8;
            let local = LocalSet::new();
            let mut pending_set = JoinSet::new();

            spawn_index_tasks(&mut pending_set, N, Some(&local));
            assert!(pending_set.try_join_next().is_none());

            local
                .run_until(async move {
                    drain_joinset_and_assert(pending_set, N).await;
                })
                .await;
        }

        /// Spawn several pending-forever tasks, and then shutdown the [`JoinSet`].
        #[tokio::test(flavor = "current_thread")]
        async fn spawn_then_shutdown() {
            const N: usize = 8;
            let local = LocalSet::new();
            let mut set = JoinSet::new();
            let mut receivers = Vec::new();

            spawn_pending_tasks(&mut set, &mut receivers, N, Some(&local));
            assert!(set.try_join_next().is_none());

            local
                .run_until(async move {
                    set.shutdown().await;
                    assert!(set.is_empty());
                    await_receivers_and_assert(receivers).await;
                })
                .await;
        }

        /// Spawn several pending-forever tasks and then drop the [`JoinSet`]
        /// before the `LocalSet` is driven and while the `LocalSet` is already driven.
        #[tokio::test(flavor = "current_thread")]
        async fn spawn_then_drop() {
            const N: usize = 8;

            {
                let local = LocalSet::new();
                let mut set = JoinSet::new();
                let mut receivers = Vec::new();

                spawn_pending_tasks(&mut set, &mut receivers, N, Some(&local));
                assert!(set.try_join_next().is_none());

                drop(set);

                local
                    .run_until(async move {
                        await_receivers_and_assert(receivers).await;
                    })
                    .await;
            }

            {
                let local = LocalSet::new();
                let mut set = JoinSet::new();
                let mut receivers = Vec::new();

                spawn_pending_tasks(&mut set, &mut receivers, N, Some(&local));
                assert!(set.try_join_next().is_none());

                local
                    .run_until(async move {
                        drop(set);
                        await_receivers_and_assert(receivers).await;
                    })
                    .await;
            }
        }
    }
}