src/test/run-pass/futures-api.rs

   1 // aux-build:arc_wake.rs
   2
   3 extern crate arc_wake;
   4
   5 use std::future::Future;
   6 use std::pin::Pin;
   7 use std::sync::{
   8     Arc,
   9     atomic::{self, AtomicUsize},
  10 };
  11 use std::task::{
  12     Context, Poll,
  13 };
  14 use arc_wake::ArcWake;
  15
  16 struct Counter {
  17     wakes: AtomicUsize,
  18 }
  19
  20 impl ArcWake for Counter {
  21     fn wake(self: Arc<Self>) {
  22         Self::wake_by_ref(&self)
  23     }
  24     fn wake_by_ref(arc_self: &Arc<Self>) {
  25         arc_self.wakes.fetch_add(1, atomic::Ordering::SeqCst);
  26     }
  27 }
  28
  29 struct MyFuture;
  30
  31 impl Future for MyFuture {
  32     type Output = ();
  33     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
  34         // Wake twice
  35         let waker = cx.waker();
  36         waker.wake_by_ref();
  37         waker.wake_by_ref();
  38         Poll::Ready(())
  39     }
  40 }
  41
  42 fn test_waker() {
  43     let counter = Arc::new(Counter {
  44         wakes: AtomicUsize::new(0),
  45     });
  46     let waker = ArcWake::into_waker(counter.clone());
  47     assert_eq!(2, Arc::strong_count(&counter));
  48     {
  49         let mut context = Context::from_waker(&waker);
  50         assert_eq!(Poll::Ready(()), Pin::new(&mut MyFuture).poll(&mut context));
  51         assert_eq!(2, counter.wakes.load(atomic::Ordering::SeqCst));
  52     }
  53     drop(waker);
  54     assert_eq!(1, Arc::strong_count(&counter));
  55 }
  56
  57 fn main() {
  58     test_waker();
  59 }