vendor/tokio/src/sync/tests/loom_atomic_waker.rs

   1 use crate::sync::task::AtomicWaker;
   2
   3 use futures::future::poll_fn;
   4 use loom::future::block_on;
   5 use loom::sync::atomic::AtomicUsize;
   6 use loom::thread;
   7 use std::sync::atomic::Ordering::Relaxed;
   8 use std::sync::Arc;
   9 use std::task::Poll::{Pending, Ready};
  10
  11 struct Chan {
  12     num: AtomicUsize,
  13     task: AtomicWaker,
  14 }
  15
  16 #[test]
  17 fn basic_notification() {
  18     const NUM_NOTIFY: usize = 2;
  19
  20     loom::model(|| {
  21         let chan = Arc::new(Chan {
  22             num: AtomicUsize::new(0),
  23             task: AtomicWaker::new(),
  24         });
  25
  26         for _ in 0..NUM_NOTIFY {
  27             let chan = chan.clone();
  28
  29             thread::spawn(move || {
  30                 chan.num.fetch_add(1, Relaxed);
  31                 chan.task.wake();
  32             });
  33         }
  34
  35         block_on(poll_fn(move |cx| {
  36             chan.task.register_by_ref(cx.waker());
  37
  38             if NUM_NOTIFY == chan.num.load(Relaxed) {
  39                 return Ready(());
  40             }
  41
  42             Pending
  43         }));
  44     });
  45 }