93 lines
2.7 KiB
Rust
93 lines
2.7 KiB
Rust
|
#![warn(rust_2018_idioms)]
|
||
|
|
#![cfg(feature = "full")]
|
||
|
|
#![cfg(not(target_os = "wasi"))] // Wasi doesn't support threads
|
||
|
|
|
||
|
|
use std::sync::atomic::{AtomicBool, Ordering};
|
||
|
|
use std::sync::Arc;
|
||
|
|
use tokio::runtime::Builder;
|
||
|
|
use tokio::sync::Notify;
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn before_park_wakes_block_on_task() {
|
||
|
|
let notify = Arc::new(Notify::new());
|
||
|
|
let notify2 = notify.clone();
|
||
|
|
let woken = Arc::new(AtomicBool::new(false));
|
||
|
|
let woken2 = woken.clone();
|
||
|
|
|
||
|
|
let rt = Builder::new_current_thread()
|
||
|
|
.enable_all()
|
||
|
|
.on_thread_park(move || {
|
||
|
|
// Only wake once to avoid busy loop if something goes wrong,
|
||
|
|
// though in this test we expect it to unpark immediately.
|
||
|
|
if !woken2.swap(true, Ordering::SeqCst) {
|
||
|
|
notify2.notify_one();
|
||
|
|
}
|
||
|
|
})
|
||
|
|
.build()
|
||
|
|
.unwrap();
|
||
|
|
|
||
|
|
rt.block_on(async {
|
||
|
|
// This will block until `notify` is notified.
|
||
|
|
// `before_park` should run when the runtime is about to park.
|
||
|
|
// It will notify `notify`, which should wake this task.
|
||
|
|
// The runtime should then see the task is woken and NOT park.
|
||
|
|
notify.notified().await;
|
||
|
|
});
|
||
|
|
|
||
|
|
assert!(woken.load(Ordering::SeqCst));
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn before_park_spawns_task() {
|
||
|
|
let notify = Arc::new(Notify::new());
|
||
|
|
let notify2 = notify.clone();
|
||
|
|
let woken = Arc::new(AtomicBool::new(false));
|
||
|
|
let woken2 = woken.clone();
|
||
|
|
|
||
|
|
let rt = Builder::new_current_thread()
|
||
|
|
.enable_all()
|
||
|
|
.on_thread_park(move || {
|
||
|
|
if !woken2.swap(true, Ordering::SeqCst) {
|
||
|
|
let notify = notify2.clone();
|
||
|
|
tokio::spawn(async move {
|
||
|
|
notify.notify_one();
|
||
|
|
});
|
||
|
|
}
|
||
|
|
})
|
||
|
|
.build()
|
||
|
|
.unwrap();
|
||
|
|
|
||
|
|
rt.block_on(async {
|
||
|
|
// This will block until `notify` is notified.
|
||
|
|
// `before_park` should run when the runtime is about to park.
|
||
|
|
// It will spawn a task that notifies `notify`.
|
||
|
|
// The runtime should see the new task and NOT park.
|
||
|
|
// If it parks, it will deadlock.
|
||
|
|
notify.notified().await;
|
||
|
|
});
|
||
|
|
|
||
|
|
assert!(woken.load(Ordering::SeqCst));
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn wake_from_other_thread_block_on() {
|
||
|
|
let rt = Builder::new_current_thread().enable_all().build().unwrap();
|
||
|
|
let handle = rt.handle().clone();
|
||
|
|
let notify = Arc::new(Notify::new());
|
||
|
|
let notify2 = notify.clone();
|
||
|
|
|
||
|
|
let th = std::thread::spawn(move || {
|
||
|
|
// Give the main thread time to park
|
||
|
|
std::thread::sleep(std::time::Duration::from_millis(5));
|
||
|
|
handle.block_on(async move {
|
||
|
|
notify2.notify_one();
|
||
|
|
});
|
||
|
|
});
|
||
|
|
|
||
|
|
rt.block_on(async {
|
||
|
|
notify.notified().await;
|
||
|
|
});
|
||
|
|
|
||
|
|
th.join().unwrap();
|
||
|
|
}
|