summaryrefslogtreecommitdiffstats
path: root/third_party/rust/tokio/src/sync/tests/loom_atomic_waker.rs
blob: f8bae65d130be3d28493ecc7f7d43a2ddf5847ed (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
use crate::sync::task::AtomicWaker;

use futures::future::poll_fn;
use loom::future::block_on;
use loom::sync::atomic::AtomicUsize;
use loom::thread;
use std::sync::atomic::Ordering::Relaxed;
use std::sync::Arc;
use std::task::Poll::{Pending, Ready};

struct Chan {
    num: AtomicUsize,
    task: AtomicWaker,
}

#[test]
fn basic_notification() {
    const NUM_NOTIFY: usize = 2;

    loom::model(|| {
        let chan = Arc::new(Chan {
            num: AtomicUsize::new(0),
            task: AtomicWaker::new(),
        });

        for _ in 0..NUM_NOTIFY {
            let chan = chan.clone();

            thread::spawn(move || {
                chan.num.fetch_add(1, Relaxed);
                chan.task.wake();
            });
        }

        block_on(poll_fn(move |cx| {
            chan.task.register_by_ref(cx.waker());

            if NUM_NOTIFY == chan.num.load(Relaxed) {
                return Ready(());
            }

            Pending
        }));
    });
}

#[test]
fn test_panicky_waker() {
    use std::panic;
    use std::ptr;
    use std::task::{RawWaker, RawWakerVTable, Waker};

    static PANICKING_VTABLE: RawWakerVTable =
        RawWakerVTable::new(|_| panic!("clone"), |_| (), |_| (), |_| ());

    let panicking = unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &PANICKING_VTABLE)) };

    // If you're working with this test (and I sure hope you never have to!),
    // uncomment the following section because there will be a lot of panics
    // which would otherwise log.
    //
    // We can't however leaved it uncommented, because it's global.
    // panic::set_hook(Box::new(|_| ()));

    const NUM_NOTIFY: usize = 2;

    loom::model(move || {
        let chan = Arc::new(Chan {
            num: AtomicUsize::new(0),
            task: AtomicWaker::new(),
        });

        for _ in 0..NUM_NOTIFY {
            let chan = chan.clone();

            thread::spawn(move || {
                chan.num.fetch_add(1, Relaxed);
                chan.task.wake();
            });
        }

        // Note: this panic should have no effect on the overall state of the
        // waker and it should proceed as normal.
        //
        // A thread above might race to flag a wakeup, and a WAKING state will
        // be preserved if this expected panic races with that so the below
        // procedure should be allowed to continue uninterrupted.
        let _ = panic::catch_unwind(|| chan.task.register_by_ref(&panicking));

        block_on(poll_fn(move |cx| {
            chan.task.register_by_ref(cx.waker());

            if NUM_NOTIFY == chan.num.load(Relaxed) {
                return Ready(());
            }

            Pending
        }));
    });
}