mirror of
				https://github.com/mozilla/gecko-dev.git
				synced 2025-11-04 10:18:41 +02:00 
			
		
		
		
	
		
			
				
	
	
		
			693 lines
		
	
	
	
		
			20 KiB
		
	
	
	
		
			Rust
		
	
	
	
	
	
			
		
		
	
	
			693 lines
		
	
	
	
		
			20 KiB
		
	
	
	
		
			Rust
		
	
	
	
	
	
#![warn(rust_2018_idioms)]
 | 
						|
#![cfg(feature = "full")]
 | 
						|
 | 
						|
use futures::{
 | 
						|
    future::{pending, ready},
 | 
						|
    FutureExt,
 | 
						|
};
 | 
						|
 | 
						|
use tokio::runtime;
 | 
						|
use tokio::sync::{mpsc, oneshot};
 | 
						|
use tokio::task::{self, LocalSet};
 | 
						|
use tokio::time;
 | 
						|
 | 
						|
#[cfg(not(tokio_wasi))]
 | 
						|
use std::cell::Cell;
 | 
						|
use std::sync::atomic::AtomicBool;
 | 
						|
#[cfg(not(tokio_wasi))]
 | 
						|
use std::sync::atomic::AtomicUsize;
 | 
						|
use std::sync::atomic::Ordering;
 | 
						|
#[cfg(not(tokio_wasi))]
 | 
						|
use std::sync::atomic::Ordering::SeqCst;
 | 
						|
use std::time::Duration;
 | 
						|
 | 
						|
#[tokio::test(flavor = "current_thread")]
 | 
						|
async fn local_current_thread_scheduler() {
 | 
						|
    LocalSet::new()
 | 
						|
        .run_until(async {
 | 
						|
            task::spawn_local(async {}).await.unwrap();
 | 
						|
        })
 | 
						|
        .await;
 | 
						|
}
 | 
						|
 | 
						|
#[cfg(not(tokio_wasi))] // Wasi doesn't support threads
 | 
						|
#[tokio::test(flavor = "multi_thread")]
 | 
						|
async fn local_threadpool() {
 | 
						|
    thread_local! {
 | 
						|
        static ON_RT_THREAD: Cell<bool> = Cell::new(false);
 | 
						|
    }
 | 
						|
 | 
						|
    ON_RT_THREAD.with(|cell| cell.set(true));
 | 
						|
 | 
						|
    LocalSet::new()
 | 
						|
        .run_until(async {
 | 
						|
            assert!(ON_RT_THREAD.with(|cell| cell.get()));
 | 
						|
            task::spawn_local(async {
 | 
						|
                assert!(ON_RT_THREAD.with(|cell| cell.get()));
 | 
						|
            })
 | 
						|
            .await
 | 
						|
            .unwrap();
 | 
						|
        })
 | 
						|
        .await;
 | 
						|
}
 | 
						|
 | 
						|
#[cfg(not(tokio_wasi))] // Wasi doesn't support threads
 | 
						|
#[tokio::test(flavor = "multi_thread")]
 | 
						|
async fn localset_future_threadpool() {
 | 
						|
    thread_local! {
 | 
						|
        static ON_LOCAL_THREAD: Cell<bool> = Cell::new(false);
 | 
						|
    }
 | 
						|
 | 
						|
    ON_LOCAL_THREAD.with(|cell| cell.set(true));
 | 
						|
 | 
						|
    let local = LocalSet::new();
 | 
						|
    local.spawn_local(async move {
 | 
						|
        assert!(ON_LOCAL_THREAD.with(|cell| cell.get()));
 | 
						|
    });
 | 
						|
    local.await;
 | 
						|
}
 | 
						|
 | 
						|
#[cfg(not(tokio_wasi))] // Wasi doesn't support threads
 | 
						|
#[tokio::test(flavor = "multi_thread")]
 | 
						|
async fn localset_future_timers() {
 | 
						|
    static RAN1: AtomicBool = AtomicBool::new(false);
 | 
						|
    static RAN2: AtomicBool = AtomicBool::new(false);
 | 
						|
 | 
						|
    let local = LocalSet::new();
 | 
						|
    local.spawn_local(async move {
 | 
						|
        time::sleep(Duration::from_millis(5)).await;
 | 
						|
        RAN1.store(true, Ordering::SeqCst);
 | 
						|
    });
 | 
						|
    local.spawn_local(async move {
 | 
						|
        time::sleep(Duration::from_millis(10)).await;
 | 
						|
        RAN2.store(true, Ordering::SeqCst);
 | 
						|
    });
 | 
						|
    local.await;
 | 
						|
    assert!(RAN1.load(Ordering::SeqCst));
 | 
						|
    assert!(RAN2.load(Ordering::SeqCst));
 | 
						|
}
 | 
						|
 | 
						|
#[tokio::test]
 | 
						|
async fn localset_future_drives_all_local_futs() {
 | 
						|
    static RAN1: AtomicBool = AtomicBool::new(false);
 | 
						|
    static RAN2: AtomicBool = AtomicBool::new(false);
 | 
						|
    static RAN3: AtomicBool = AtomicBool::new(false);
 | 
						|
 | 
						|
    let local = LocalSet::new();
 | 
						|
    local.spawn_local(async move {
 | 
						|
        task::spawn_local(async {
 | 
						|
            task::yield_now().await;
 | 
						|
            RAN3.store(true, Ordering::SeqCst);
 | 
						|
        });
 | 
						|
        task::yield_now().await;
 | 
						|
        RAN1.store(true, Ordering::SeqCst);
 | 
						|
    });
 | 
						|
    local.spawn_local(async move {
 | 
						|
        task::yield_now().await;
 | 
						|
        RAN2.store(true, Ordering::SeqCst);
 | 
						|
    });
 | 
						|
    local.await;
 | 
						|
    assert!(RAN1.load(Ordering::SeqCst));
 | 
						|
    assert!(RAN2.load(Ordering::SeqCst));
 | 
						|
    assert!(RAN3.load(Ordering::SeqCst));
 | 
						|
}
 | 
						|
 | 
						|
#[cfg(not(tokio_wasi))] // Wasi doesn't support threads
 | 
						|
#[tokio::test(flavor = "multi_thread")]
 | 
						|
async fn local_threadpool_timer() {
 | 
						|
    // This test ensures that runtime services like the timer are properly
 | 
						|
    // set for the local task set.
 | 
						|
    thread_local! {
 | 
						|
        static ON_RT_THREAD: Cell<bool> = Cell::new(false);
 | 
						|
    }
 | 
						|
 | 
						|
    ON_RT_THREAD.with(|cell| cell.set(true));
 | 
						|
 | 
						|
    LocalSet::new()
 | 
						|
        .run_until(async {
 | 
						|
            assert!(ON_RT_THREAD.with(|cell| cell.get()));
 | 
						|
            let join = task::spawn_local(async move {
 | 
						|
                assert!(ON_RT_THREAD.with(|cell| cell.get()));
 | 
						|
                time::sleep(Duration::from_millis(10)).await;
 | 
						|
                assert!(ON_RT_THREAD.with(|cell| cell.get()));
 | 
						|
            });
 | 
						|
            join.await.unwrap();
 | 
						|
        })
 | 
						|
        .await;
 | 
						|
}
 | 
						|
#[test]
 | 
						|
fn enter_guard_spawn() {
 | 
						|
    let local = LocalSet::new();
 | 
						|
    let _guard = local.enter();
 | 
						|
    // Run the local task set.
 | 
						|
 | 
						|
    let join = task::spawn_local(async { true });
 | 
						|
    let rt = runtime::Builder::new_current_thread()
 | 
						|
        .enable_all()
 | 
						|
        .build()
 | 
						|
        .unwrap();
 | 
						|
    local.block_on(&rt, async move {
 | 
						|
        assert!(join.await.unwrap());
 | 
						|
    });
 | 
						|
}
 | 
						|
 | 
						|
#[cfg(not(tokio_wasi))] // Wasi doesn't support panic recovery
 | 
						|
#[test]
 | 
						|
// This will panic, since the thread that calls `block_on` cannot use
 | 
						|
// in-place blocking inside of `block_on`.
 | 
						|
#[should_panic]
 | 
						|
fn local_threadpool_blocking_in_place() {
 | 
						|
    thread_local! {
 | 
						|
        static ON_RT_THREAD: Cell<bool> = Cell::new(false);
 | 
						|
    }
 | 
						|
 | 
						|
    ON_RT_THREAD.with(|cell| cell.set(true));
 | 
						|
 | 
						|
    let rt = runtime::Builder::new_current_thread()
 | 
						|
        .enable_all()
 | 
						|
        .build()
 | 
						|
        .unwrap();
 | 
						|
    LocalSet::new().block_on(&rt, async {
 | 
						|
        assert!(ON_RT_THREAD.with(|cell| cell.get()));
 | 
						|
        let join = task::spawn_local(async move {
 | 
						|
            assert!(ON_RT_THREAD.with(|cell| cell.get()));
 | 
						|
            task::block_in_place(|| {});
 | 
						|
            assert!(ON_RT_THREAD.with(|cell| cell.get()));
 | 
						|
        });
 | 
						|
        join.await.unwrap();
 | 
						|
    });
 | 
						|
}
 | 
						|
 | 
						|
#[cfg(not(tokio_wasi))] // Wasi doesn't support threads
 | 
						|
#[tokio::test(flavor = "multi_thread")]
 | 
						|
async fn local_threadpool_blocking_run() {
 | 
						|
    thread_local! {
 | 
						|
        static ON_RT_THREAD: Cell<bool> = Cell::new(false);
 | 
						|
    }
 | 
						|
 | 
						|
    ON_RT_THREAD.with(|cell| cell.set(true));
 | 
						|
 | 
						|
    LocalSet::new()
 | 
						|
        .run_until(async {
 | 
						|
            assert!(ON_RT_THREAD.with(|cell| cell.get()));
 | 
						|
            let join = task::spawn_local(async move {
 | 
						|
                assert!(ON_RT_THREAD.with(|cell| cell.get()));
 | 
						|
                task::spawn_blocking(|| {
 | 
						|
                    assert!(
 | 
						|
                        !ON_RT_THREAD.with(|cell| cell.get()),
 | 
						|
                        "blocking must not run on the local task set's thread"
 | 
						|
                    );
 | 
						|
                })
 | 
						|
                .await
 | 
						|
                .unwrap();
 | 
						|
                assert!(ON_RT_THREAD.with(|cell| cell.get()));
 | 
						|
            });
 | 
						|
            join.await.unwrap();
 | 
						|
        })
 | 
						|
        .await;
 | 
						|
}
 | 
						|
 | 
						|
#[cfg(not(tokio_wasi))] // Wasi doesn't support threads
 | 
						|
#[tokio::test(flavor = "multi_thread")]
 | 
						|
async fn all_spawns_are_local() {
 | 
						|
    use futures::future;
 | 
						|
    thread_local! {
 | 
						|
        static ON_RT_THREAD: Cell<bool> = Cell::new(false);
 | 
						|
    }
 | 
						|
 | 
						|
    ON_RT_THREAD.with(|cell| cell.set(true));
 | 
						|
 | 
						|
    LocalSet::new()
 | 
						|
        .run_until(async {
 | 
						|
            assert!(ON_RT_THREAD.with(|cell| cell.get()));
 | 
						|
            let handles = (0..128)
 | 
						|
                .map(|_| {
 | 
						|
                    task::spawn_local(async {
 | 
						|
                        assert!(ON_RT_THREAD.with(|cell| cell.get()));
 | 
						|
                    })
 | 
						|
                })
 | 
						|
                .collect::<Vec<_>>();
 | 
						|
            for joined in future::join_all(handles).await {
 | 
						|
                joined.unwrap();
 | 
						|
            }
 | 
						|
        })
 | 
						|
        .await;
 | 
						|
}
 | 
						|
 | 
						|
#[cfg(not(tokio_wasi))] // Wasi doesn't support threads
 | 
						|
#[tokio::test(flavor = "multi_thread")]
 | 
						|
async fn nested_spawn_is_local() {
 | 
						|
    thread_local! {
 | 
						|
        static ON_RT_THREAD: Cell<bool> = Cell::new(false);
 | 
						|
    }
 | 
						|
 | 
						|
    ON_RT_THREAD.with(|cell| cell.set(true));
 | 
						|
 | 
						|
    LocalSet::new()
 | 
						|
        .run_until(async {
 | 
						|
            assert!(ON_RT_THREAD.with(|cell| cell.get()));
 | 
						|
            task::spawn_local(async {
 | 
						|
                assert!(ON_RT_THREAD.with(|cell| cell.get()));
 | 
						|
                task::spawn_local(async {
 | 
						|
                    assert!(ON_RT_THREAD.with(|cell| cell.get()));
 | 
						|
                    task::spawn_local(async {
 | 
						|
                        assert!(ON_RT_THREAD.with(|cell| cell.get()));
 | 
						|
                        task::spawn_local(async {
 | 
						|
                            assert!(ON_RT_THREAD.with(|cell| cell.get()));
 | 
						|
                        })
 | 
						|
                        .await
 | 
						|
                        .unwrap();
 | 
						|
                    })
 | 
						|
                    .await
 | 
						|
                    .unwrap();
 | 
						|
                })
 | 
						|
                .await
 | 
						|
                .unwrap();
 | 
						|
            })
 | 
						|
            .await
 | 
						|
            .unwrap();
 | 
						|
        })
 | 
						|
        .await;
 | 
						|
}
 | 
						|
 | 
						|
#[cfg(not(tokio_wasi))] // Wasi doesn't support threads
 | 
						|
#[test]
 | 
						|
fn join_local_future_elsewhere() {
 | 
						|
    thread_local! {
 | 
						|
        static ON_RT_THREAD: Cell<bool> = Cell::new(false);
 | 
						|
    }
 | 
						|
 | 
						|
    ON_RT_THREAD.with(|cell| cell.set(true));
 | 
						|
 | 
						|
    let rt = runtime::Runtime::new().unwrap();
 | 
						|
    let local = LocalSet::new();
 | 
						|
    local.block_on(&rt, async move {
 | 
						|
        let (tx, rx) = oneshot::channel();
 | 
						|
        let join = task::spawn_local(async move {
 | 
						|
            assert!(
 | 
						|
                ON_RT_THREAD.with(|cell| cell.get()),
 | 
						|
                "local task must run on local thread, no matter where it is awaited"
 | 
						|
            );
 | 
						|
            rx.await.unwrap();
 | 
						|
 | 
						|
            "hello world"
 | 
						|
        });
 | 
						|
        let join2 = task::spawn(async move {
 | 
						|
            assert!(
 | 
						|
                !ON_RT_THREAD.with(|cell| cell.get()),
 | 
						|
                "spawned task should be on a worker"
 | 
						|
            );
 | 
						|
 | 
						|
            tx.send(()).expect("task shouldn't have ended yet");
 | 
						|
 | 
						|
            join.await.expect("task should complete successfully");
 | 
						|
        });
 | 
						|
        join2.await.unwrap()
 | 
						|
    });
 | 
						|
}
 | 
						|
 | 
						|
// Tests for <https://github.com/tokio-rs/tokio/issues/4973>
 | 
						|
#[cfg(not(tokio_wasi))] // Wasi doesn't support threads
 | 
						|
#[tokio::test(flavor = "multi_thread")]
 | 
						|
async fn localset_in_thread_local() {
 | 
						|
    thread_local! {
 | 
						|
        static LOCAL_SET: LocalSet = LocalSet::new();
 | 
						|
    }
 | 
						|
 | 
						|
    // holds runtime thread until end of main fn.
 | 
						|
    let (_tx, rx) = oneshot::channel::<()>();
 | 
						|
    let handle = tokio::runtime::Handle::current();
 | 
						|
 | 
						|
    std::thread::spawn(move || {
 | 
						|
        LOCAL_SET.with(|local_set| {
 | 
						|
            handle.block_on(local_set.run_until(async move {
 | 
						|
                let _ = rx.await;
 | 
						|
            }))
 | 
						|
        });
 | 
						|
    });
 | 
						|
}
 | 
						|
 | 
						|
#[test]
 | 
						|
fn drop_cancels_tasks() {
 | 
						|
    use std::rc::Rc;
 | 
						|
 | 
						|
    // This test reproduces issue #1842
 | 
						|
    let rt = rt();
 | 
						|
    let rc1 = Rc::new(());
 | 
						|
    let rc2 = rc1.clone();
 | 
						|
 | 
						|
    let (started_tx, started_rx) = oneshot::channel();
 | 
						|
 | 
						|
    let local = LocalSet::new();
 | 
						|
    local.spawn_local(async move {
 | 
						|
        // Move this in
 | 
						|
        let _rc2 = rc2;
 | 
						|
 | 
						|
        started_tx.send(()).unwrap();
 | 
						|
        futures::future::pending::<()>().await;
 | 
						|
    });
 | 
						|
 | 
						|
    local.block_on(&rt, async {
 | 
						|
        started_rx.await.unwrap();
 | 
						|
    });
 | 
						|
    drop(local);
 | 
						|
    drop(rt);
 | 
						|
 | 
						|
    assert_eq!(1, Rc::strong_count(&rc1));
 | 
						|
}
 | 
						|
 | 
						|
/// Runs a test function in a separate thread, and panics if the test does not
 | 
						|
/// complete within the specified timeout, or if the test function panics.
 | 
						|
///
 | 
						|
/// This is intended for running tests whose failure mode is a hang or infinite
 | 
						|
/// loop that cannot be detected otherwise.
 | 
						|
fn with_timeout(timeout: Duration, f: impl FnOnce() + Send + 'static) {
 | 
						|
    use std::sync::mpsc::RecvTimeoutError;
 | 
						|
 | 
						|
    let (done_tx, done_rx) = std::sync::mpsc::channel();
 | 
						|
    let thread = std::thread::spawn(move || {
 | 
						|
        f();
 | 
						|
 | 
						|
        // Send a message on the channel so that the test thread can
 | 
						|
        // determine if we have entered an infinite loop:
 | 
						|
        done_tx.send(()).unwrap();
 | 
						|
    });
 | 
						|
 | 
						|
    // Since the failure mode of this test is an infinite loop, rather than
 | 
						|
    // something we can easily make assertions about, we'll run it in a
 | 
						|
    // thread. When the test thread finishes, it will send a message on a
 | 
						|
    // channel to this thread. We'll wait for that message with a fairly
 | 
						|
    // generous timeout, and if we don't receive it, we assume the test
 | 
						|
    // thread has hung.
 | 
						|
    //
 | 
						|
    // Note that it should definitely complete in under a minute, but just
 | 
						|
    // in case CI is slow, we'll give it a long timeout.
 | 
						|
    match done_rx.recv_timeout(timeout) {
 | 
						|
        Err(RecvTimeoutError::Timeout) => panic!(
 | 
						|
            "test did not complete within {:?} seconds, \
 | 
						|
             we have (probably) entered an infinite loop!",
 | 
						|
            timeout,
 | 
						|
        ),
 | 
						|
        // Did the test thread panic? We'll find out for sure when we `join`
 | 
						|
        // with it.
 | 
						|
        Err(RecvTimeoutError::Disconnected) => {}
 | 
						|
        // Test completed successfully!
 | 
						|
        Ok(()) => {}
 | 
						|
    }
 | 
						|
 | 
						|
    thread.join().expect("test thread should not panic!")
 | 
						|
}
 | 
						|
 | 
						|
#[cfg_attr(tokio_wasi, ignore = "`unwrap()` in `with_timeout()` panics on Wasi")]
 | 
						|
#[test]
 | 
						|
fn drop_cancels_remote_tasks() {
 | 
						|
    // This test reproduces issue #1885.
 | 
						|
    with_timeout(Duration::from_secs(60), || {
 | 
						|
        let (tx, mut rx) = mpsc::channel::<()>(1024);
 | 
						|
 | 
						|
        let rt = rt();
 | 
						|
 | 
						|
        let local = LocalSet::new();
 | 
						|
        local.spawn_local(async move { while rx.recv().await.is_some() {} });
 | 
						|
        local.block_on(&rt, async {
 | 
						|
            time::sleep(Duration::from_millis(1)).await;
 | 
						|
        });
 | 
						|
 | 
						|
        drop(tx);
 | 
						|
 | 
						|
        // This enters an infinite loop if the remote notified tasks are not
 | 
						|
        // properly cancelled.
 | 
						|
        drop(local);
 | 
						|
    });
 | 
						|
}
 | 
						|
 | 
						|
#[cfg_attr(
 | 
						|
    tokio_wasi,
 | 
						|
    ignore = "FIXME: `task::spawn_local().await.unwrap()` panics on Wasi"
 | 
						|
)]
 | 
						|
#[test]
 | 
						|
fn local_tasks_wake_join_all() {
 | 
						|
    // This test reproduces issue #2460.
 | 
						|
    with_timeout(Duration::from_secs(60), || {
 | 
						|
        use futures::future::join_all;
 | 
						|
        use tokio::task::LocalSet;
 | 
						|
 | 
						|
        let rt = rt();
 | 
						|
        let set = LocalSet::new();
 | 
						|
        let mut handles = Vec::new();
 | 
						|
 | 
						|
        for _ in 1..=128 {
 | 
						|
            handles.push(set.spawn_local(async move {
 | 
						|
                tokio::task::spawn_local(async move {}).await.unwrap();
 | 
						|
            }));
 | 
						|
        }
 | 
						|
 | 
						|
        rt.block_on(set.run_until(join_all(handles)));
 | 
						|
    });
 | 
						|
}
 | 
						|
 | 
						|
#[cfg(not(tokio_wasi))] // Wasi doesn't support panic recovery
 | 
						|
#[test]
 | 
						|
fn local_tasks_are_polled_after_tick() {
 | 
						|
    // This test depends on timing, so we run it up to five times.
 | 
						|
    for _ in 0..4 {
 | 
						|
        let res = std::panic::catch_unwind(local_tasks_are_polled_after_tick_inner);
 | 
						|
        if res.is_ok() {
 | 
						|
            // success
 | 
						|
            return;
 | 
						|
        }
 | 
						|
    }
 | 
						|
 | 
						|
    // Test failed 4 times. Try one more time without catching panics. If it
 | 
						|
    // fails again, the test fails.
 | 
						|
    local_tasks_are_polled_after_tick_inner();
 | 
						|
}
 | 
						|
 | 
						|
#[cfg(not(tokio_wasi))] // Wasi doesn't support panic recovery
 | 
						|
#[tokio::main(flavor = "current_thread")]
 | 
						|
async fn local_tasks_are_polled_after_tick_inner() {
 | 
						|
    // Reproduces issues #1899 and #1900
 | 
						|
 | 
						|
    static RX1: AtomicUsize = AtomicUsize::new(0);
 | 
						|
    static RX2: AtomicUsize = AtomicUsize::new(0);
 | 
						|
    const EXPECTED: usize = 500;
 | 
						|
 | 
						|
    RX1.store(0, SeqCst);
 | 
						|
    RX2.store(0, SeqCst);
 | 
						|
 | 
						|
    let (tx, mut rx) = mpsc::unbounded_channel();
 | 
						|
 | 
						|
    let local = LocalSet::new();
 | 
						|
 | 
						|
    local
 | 
						|
        .run_until(async {
 | 
						|
            let task2 = task::spawn(async move {
 | 
						|
                // Wait a bit
 | 
						|
                time::sleep(Duration::from_millis(10)).await;
 | 
						|
 | 
						|
                let mut oneshots = Vec::with_capacity(EXPECTED);
 | 
						|
 | 
						|
                // Send values
 | 
						|
                for _ in 0..EXPECTED {
 | 
						|
                    let (oneshot_tx, oneshot_rx) = oneshot::channel();
 | 
						|
                    oneshots.push(oneshot_tx);
 | 
						|
                    tx.send(oneshot_rx).unwrap();
 | 
						|
                }
 | 
						|
 | 
						|
                time::sleep(Duration::from_millis(10)).await;
 | 
						|
 | 
						|
                for tx in oneshots.drain(..) {
 | 
						|
                    tx.send(()).unwrap();
 | 
						|
                }
 | 
						|
 | 
						|
                loop {
 | 
						|
                    time::sleep(Duration::from_millis(20)).await;
 | 
						|
                    let rx1 = RX1.load(SeqCst);
 | 
						|
                    let rx2 = RX2.load(SeqCst);
 | 
						|
 | 
						|
                    if rx1 == EXPECTED && rx2 == EXPECTED {
 | 
						|
                        break;
 | 
						|
                    }
 | 
						|
                }
 | 
						|
            });
 | 
						|
 | 
						|
            while let Some(oneshot) = rx.recv().await {
 | 
						|
                RX1.fetch_add(1, SeqCst);
 | 
						|
 | 
						|
                task::spawn_local(async move {
 | 
						|
                    oneshot.await.unwrap();
 | 
						|
                    RX2.fetch_add(1, SeqCst);
 | 
						|
                });
 | 
						|
            }
 | 
						|
 | 
						|
            task2.await.unwrap();
 | 
						|
        })
 | 
						|
        .await;
 | 
						|
}
 | 
						|
 | 
						|
#[tokio::test]
 | 
						|
async fn acquire_mutex_in_drop() {
 | 
						|
    use futures::future::pending;
 | 
						|
 | 
						|
    let (tx1, rx1) = oneshot::channel();
 | 
						|
    let (tx2, rx2) = oneshot::channel();
 | 
						|
    let local = LocalSet::new();
 | 
						|
 | 
						|
    local.spawn_local(async move {
 | 
						|
        let _ = rx2.await;
 | 
						|
        unreachable!();
 | 
						|
    });
 | 
						|
 | 
						|
    local.spawn_local(async move {
 | 
						|
        let _ = rx1.await;
 | 
						|
        tx2.send(()).unwrap();
 | 
						|
        unreachable!();
 | 
						|
    });
 | 
						|
 | 
						|
    // Spawn a task that will never notify
 | 
						|
    local.spawn_local(async move {
 | 
						|
        pending::<()>().await;
 | 
						|
        tx1.send(()).unwrap();
 | 
						|
    });
 | 
						|
 | 
						|
    // Tick the loop
 | 
						|
    local
 | 
						|
        .run_until(async {
 | 
						|
            task::yield_now().await;
 | 
						|
        })
 | 
						|
        .await;
 | 
						|
 | 
						|
    // Drop the LocalSet
 | 
						|
    drop(local);
 | 
						|
}
 | 
						|
 | 
						|
#[tokio::test]
 | 
						|
async fn spawn_wakes_localset() {
 | 
						|
    let local = LocalSet::new();
 | 
						|
    futures::select! {
 | 
						|
        _ = local.run_until(pending::<()>()).fuse() => unreachable!(),
 | 
						|
        ret = async { local.spawn_local(ready(())).await.unwrap()}.fuse() => ret
 | 
						|
    }
 | 
						|
}
 | 
						|
 | 
						|
#[test]
 | 
						|
fn store_local_set_in_thread_local_with_runtime() {
 | 
						|
    use tokio::runtime::Runtime;
 | 
						|
 | 
						|
    thread_local! {
 | 
						|
        static CURRENT: RtAndLocalSet = RtAndLocalSet::new();
 | 
						|
    }
 | 
						|
 | 
						|
    struct RtAndLocalSet {
 | 
						|
        rt: Runtime,
 | 
						|
        local: LocalSet,
 | 
						|
    }
 | 
						|
 | 
						|
    impl RtAndLocalSet {
 | 
						|
        fn new() -> RtAndLocalSet {
 | 
						|
            RtAndLocalSet {
 | 
						|
                rt: tokio::runtime::Builder::new_current_thread()
 | 
						|
                    .enable_all()
 | 
						|
                    .build()
 | 
						|
                    .unwrap(),
 | 
						|
                local: LocalSet::new(),
 | 
						|
            }
 | 
						|
        }
 | 
						|
 | 
						|
        async fn inner_method(&self) {
 | 
						|
            self.local
 | 
						|
                .run_until(async move {
 | 
						|
                    tokio::task::spawn_local(async {});
 | 
						|
                })
 | 
						|
                .await
 | 
						|
        }
 | 
						|
 | 
						|
        fn method(&self) {
 | 
						|
            self.rt.block_on(self.inner_method());
 | 
						|
        }
 | 
						|
    }
 | 
						|
 | 
						|
    CURRENT.with(|f| {
 | 
						|
        f.method();
 | 
						|
    });
 | 
						|
}
 | 
						|
 | 
						|
#[cfg(tokio_unstable)]
 | 
						|
mod unstable {
 | 
						|
    use tokio::runtime::UnhandledPanic;
 | 
						|
    use tokio::task::LocalSet;
 | 
						|
 | 
						|
    #[tokio::test]
 | 
						|
    #[should_panic(
 | 
						|
        expected = "a spawned task panicked and the LocalSet is configured to shutdown on unhandled panic"
 | 
						|
    )]
 | 
						|
    async fn shutdown_on_panic() {
 | 
						|
        LocalSet::new()
 | 
						|
            .unhandled_panic(UnhandledPanic::ShutdownRuntime)
 | 
						|
            .run_until(async {
 | 
						|
                tokio::task::spawn_local(async {
 | 
						|
                    panic!("boom");
 | 
						|
                });
 | 
						|
 | 
						|
                futures::future::pending::<()>().await;
 | 
						|
            })
 | 
						|
            .await;
 | 
						|
    }
 | 
						|
 | 
						|
    // This test compares that, when the task driving `run_until` has already
 | 
						|
    // consumed budget, the `run_until` future has less budget than a "spawned"
 | 
						|
    // task.
 | 
						|
    //
 | 
						|
    // "Budget" is a fuzzy metric as the Tokio runtime is able to change values
 | 
						|
    // internally. This is why the test uses indirection to test this.
 | 
						|
    #[tokio::test]
 | 
						|
    async fn run_until_does_not_get_own_budget() {
 | 
						|
        // Consume some budget
 | 
						|
        tokio::task::consume_budget().await;
 | 
						|
 | 
						|
        LocalSet::new()
 | 
						|
            .run_until(async {
 | 
						|
                let spawned = tokio::spawn(async {
 | 
						|
                    let mut spawned_n = 0;
 | 
						|
 | 
						|
                    {
 | 
						|
                        let mut spawned = tokio_test::task::spawn(async {
 | 
						|
                            loop {
 | 
						|
                                spawned_n += 1;
 | 
						|
                                tokio::task::consume_budget().await;
 | 
						|
                            }
 | 
						|
                        });
 | 
						|
                        // Poll once
 | 
						|
                        assert!(!spawned.poll().is_ready());
 | 
						|
                    }
 | 
						|
 | 
						|
                    spawned_n
 | 
						|
                });
 | 
						|
 | 
						|
                let mut run_until_n = 0;
 | 
						|
                {
 | 
						|
                    let mut run_until = tokio_test::task::spawn(async {
 | 
						|
                        loop {
 | 
						|
                            run_until_n += 1;
 | 
						|
                            tokio::task::consume_budget().await;
 | 
						|
                        }
 | 
						|
                    });
 | 
						|
                    // Poll once
 | 
						|
                    assert!(!run_until.poll().is_ready());
 | 
						|
                }
 | 
						|
 | 
						|
                let spawned_n = spawned.await.unwrap();
 | 
						|
                assert_ne!(spawned_n, 0);
 | 
						|
                assert_ne!(run_until_n, 0);
 | 
						|
                assert!(spawned_n > run_until_n);
 | 
						|
            })
 | 
						|
            .await
 | 
						|
    }
 | 
						|
}
 | 
						|
 | 
						|
fn rt() -> runtime::Runtime {
 | 
						|
    tokio::runtime::Builder::new_current_thread()
 | 
						|
        .enable_all()
 | 
						|
        .build()
 | 
						|
        .unwrap()
 | 
						|
}
 |