src/test/ui/async-await/futures-api.rs

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