Quelle  rt_common_before_park.rs

#![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();
}