task_local.rs

#![cfg(all(feature = "full", not(tokio_wasi)))] // Wasi doesn't support threads
#![allow(clippy::declare_interior_mutable_const)]
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::sync::oneshot;

#[tokio::test(flavor = "multi_thread")]
async fn local() {
    tokio::task_local! {
        static REQ_ID: u32;
        pub static FOO: bool;
    }

    let j1 = tokio::spawn(REQ_ID.scope(1, async move {
        assert_eq!(REQ_ID.get(), 1);
        assert_eq!(REQ_ID.get(), 1);
    }));

    let j2 = tokio::spawn(REQ_ID.scope(2, async move {
        REQ_ID.with(|v| {
            assert_eq!(REQ_ID.get(), 2);
            assert_eq!(*v, 2);
        });

        tokio::time::sleep(std::time::Duration::from_millis(10)).await;

        assert_eq!(REQ_ID.get(), 2);
    }));

    let j3 = tokio::spawn(FOO.scope(true, async move {
        assert!(FOO.get());
    }));

    j1.await.unwrap();
    j2.await.unwrap();
    j3.await.unwrap();
}

#[tokio::test]
async fn task_local_available_on_abort() {
    tokio::task_local! {
        static KEY: u32;
    }

    struct MyFuture {
        tx_poll: Option<oneshot::Sender<()>>,
        tx_drop: Option<oneshot::Sender<u32>>,
    }
    impl Future for MyFuture {
        type Output = ();

        fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
            if let Some(tx_poll) = self.tx_poll.take() {
                let _ = tx_poll.send(());
            }
            Poll::Pending
        }
    }
    impl Drop for MyFuture {
        fn drop(&mut self) {
            let _ = self.tx_drop.take().unwrap().send(KEY.get());
        }
    }

    let (tx_drop, rx_drop) = oneshot::channel();
    let (tx_poll, rx_poll) = oneshot::channel();

    let h = tokio::spawn(KEY.scope(
        42,
        MyFuture {
            tx_poll: Some(tx_poll),
            tx_drop: Some(tx_drop),
        },
    ));

    rx_poll.await.unwrap();
    h.abort();
    assert_eq!(rx_drop.await.unwrap(), 42);

    let err = h.await.unwrap_err();
    if !err.is_cancelled() {
        if let Ok(panic) = err.try_into_panic() {
            std::panic::resume_unwind(panic);
        } else {
            panic!();
        }
    }
}

#[tokio::test]
async fn task_local_available_on_completion_drop() {
    tokio::task_local! {
        static KEY: u32;
    }

    struct MyFuture {
        tx: Option<oneshot::Sender<u32>>,
    }
    impl Future for MyFuture {
        type Output = ();

        fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
            Poll::Ready(())
        }
    }
    impl Drop for MyFuture {
        fn drop(&mut self) {
            let _ = self.tx.take().unwrap().send(KEY.get());
        }
    }

    let (tx, rx) = oneshot::channel();

    let h = tokio::spawn(KEY.scope(42, MyFuture { tx: Some(tx) }));

    assert_eq!(rx.await.unwrap(), 42);
    h.await.unwrap();
}

#[tokio::test(flavor = "current_thread")]
async fn task_id_spawn() {
    use tokio::task;

    task::spawn_blocking(|| println!("task id: {}", task::id()))
        .await
        .unwrap();
}

#[tokio::test(flavor = "current_thread")]
async fn task_id_spawn_blocking() {
    use tokio::task;

    tokio::spawn(async { println!("task id: {}", task::id()) })
        .await
        .unwrap();
}

#[cfg(tokio_unstable)]
#[tokio::test(flavor = "multi_thread")]
async fn task_id_collision_multi_thread() {
    use tokio::task;

    let handle1 = tokio::spawn(async { task::id() });
    let handle2 = tokio::spawn(async { task::id() });

    let (id1, id2) = tokio::join!(handle1, handle2);
    assert_ne!(id1.unwrap(), id2.unwrap());
}

#[cfg(tokio_unstable)]
#[tokio::test(flavor = "current_thread")]
async fn task_id_collision_current_thread() {
    use tokio::task;

    let handle1 = tokio::spawn(async { task::id() });
    let handle2 = tokio::spawn(async { task::id() });

    let (id1, id2) = tokio::join!(handle1, handle2);
    assert_ne!(id1.unwrap(), id2.unwrap());
}

#[cfg(tokio_unstable)]
#[tokio::test(flavor = "current_thread")]
async fn task_ids_match_current_thread() {
    use tokio::{sync::oneshot, task};

    let (tx, rx) = oneshot::channel();
    let handle = tokio::spawn(async {
        let id = rx.await.unwrap();
        assert_eq!(id, task::id());
    });
    tx.send(handle.id()).unwrap();
    handle.await.unwrap();
}

#[cfg(tokio_unstable)]
#[tokio::test(flavor = "multi_thread")]
async fn task_ids_match_multi_thread() {
    use tokio::{sync::oneshot, task};

    let (tx, rx) = oneshot::channel();
    let handle = tokio::spawn(async {
        let id = rx.await.unwrap();
        assert_eq!(id, task::id());
    });
    tx.send(handle.id()).unwrap();
    handle.await.unwrap();
}

#[cfg(tokio_unstable)]
#[tokio::test(flavor = "multi_thread")]
async fn task_id_future_destructor_completion() {
    use tokio::task;

    struct MyFuture;
    impl Future for MyFuture {
        type Output = ();

        fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
            Poll::Ready(())
        }
    }
    impl Drop for MyFuture {
        fn drop(&mut self) {
            println!("task id: {}", task::id());
        }
    }

    tokio::spawn(MyFuture).await.unwrap();
}

#[cfg(tokio_unstable)]
#[tokio::test(flavor = "multi_thread")]
async fn task_id_future_destructor_abort() {
    use tokio::task;

    struct MyFuture;
    impl Future for MyFuture {
        type Output = ();

        fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
            Poll::Pending
        }
    }
    impl Drop for MyFuture {
        fn drop(&mut self) {
            println!("task id: {}", task::id());
        }
    }

    tokio::spawn(MyFuture).abort();
}

#[cfg(tokio_unstable)]
#[tokio::test(flavor = "current_thread")]
async fn task_id_output_destructor_handle_dropped_before_completion() {
    use tokio::task;

    struct MyOutput;
    impl Drop for MyOutput {
        fn drop(&mut self) {
            println!("task id: {}", task::id());
        }
    }

    struct MyFuture {
        tx: Option<oneshot::Sender<()>>,
    }
    impl Future for MyFuture {
        type Output = MyOutput;

        fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
            let _ = self.tx.take().unwrap().send(());
            Poll::Ready(MyOutput)
        }
    }
    impl Drop for MyFuture {
        fn drop(&mut self) {
            println!("task id: {}", task::id());
        }
    }

    let (tx, rx) = oneshot::channel();
    let handle = tokio::spawn(MyFuture { tx: Some(tx) });
    drop(handle);
    rx.await.unwrap();
}

#[cfg(tokio_unstable)]
#[tokio::test(flavor = "current_thread")]
async fn task_id_output_destructor_handle_dropped_after_completion() {
    use tokio::task;

    struct MyOutput;
    impl Drop for MyOutput {
        fn drop(&mut self) {
            println!("task id: {}", task::id());
        }
    }

    struct MyFuture {
        tx: Option<oneshot::Sender<()>>,
    }
    impl Future for MyFuture {
        type Output = MyOutput;

        fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
            let _ = self.tx.take().unwrap().send(());
            Poll::Ready(MyOutput)
        }
    }
    impl Drop for MyFuture {
        fn drop(&mut self) {
            println!("task id: {}", task::id());
        }
    }

    let (tx, rx) = oneshot::channel();
    let handle = tokio::spawn(MyFuture { tx: Some(tx) });
    rx.await.unwrap();
    drop(handle);
}