src/test/run-pass/async-await.rs

   1 // Copyright 2018 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 // compile-flags: --edition=2018
  12
  13 #![feature(arbitrary_self_types, async_await, await_macro, futures_api, pin)]
  14
  15 use std::boxed::PinBox;
  16 use std::mem::PinMut;
  17 use std::future::Future;
  18 use std::sync::{
  19     Arc,
  20     atomic::{self, AtomicUsize},
  21 };
  22 use std::task::{
  23     Context, Poll, Wake,
  24     Executor, TaskObj, SpawnObjError,
  25     local_waker_from_nonlocal,
  26 };
  27
  28 struct Counter {
  29     wakes: AtomicUsize,
  30 }
  31
  32 impl Wake for Counter {
  33     fn wake(this: &Arc<Self>) {
  34         this.wakes.fetch_add(1, atomic::Ordering::SeqCst);
  35     }
  36 }
  37
  38 struct NoopExecutor;
  39 impl Executor for NoopExecutor {
  40     fn spawn_obj(&mut self, _: TaskObj) -> Result<(), SpawnObjError> {
  41         Ok(())
  42     }
  43 }
  44
  45 struct WakeOnceThenComplete(bool);
  46
  47 fn wake_and_yield_once() -> WakeOnceThenComplete { WakeOnceThenComplete(false) }
  48
  49 impl Future for WakeOnceThenComplete {
  50     type Output = ();
  51     fn poll(mut self: PinMut<Self>, cx: &mut Context) -> Poll<()> {
  52         if self.0 {
  53             Poll::Ready(())
  54         } else {
  55             cx.waker().wake();
  56             self.0 = true;
  57             Poll::Pending
  58         }
  59     }
  60 }
  61
  62 fn async_block(x: u8) -> impl Future<Output = u8> {
  63     async move {
  64         await!(wake_and_yield_once());
  65         x
  66     }
  67 }
  68
  69 fn async_nonmove_block(x: u8) -> impl Future<Output = u8> {
  70     async move {
  71         let future = async {
  72             await!(wake_and_yield_once());
  73             x
  74         };
  75         await!(future)
  76     }
  77 }
  78
  79 fn async_closure(x: u8) -> impl Future<Output = u8> {
  80     (async move |x: u8| -> u8 {
  81         await!(wake_and_yield_once());
  82         x
  83     })(x)
  84 }
  85
  86 async fn async_fn(x: u8) -> u8 {
  87     await!(wake_and_yield_once());
  88     x
  89 }
  90
  91 async fn async_fn_with_borrow(x: &u8) -> u8 {
  92     await!(wake_and_yield_once());
  93     *x
  94 }
  95
  96 fn async_fn_with_internal_borrow(y: u8) -> impl Future<Output = u8> {
  97     async move {
  98         await!(async_fn_with_borrow(&y))
  99     }
 100 }
 101
 102 unsafe async fn unsafe_async_fn(x: u8) -> u8 {
 103     await!(wake_and_yield_once());
 104     x
 105 }
 106
 107 struct Foo {
 108     async fn async_method(x: u8) -> u8 {
 109         unsafe {
 110             await!(unsafe_async_fn())
 111         }
 112     }
 113 }
 114
 115 fn test_future_yields_once_then_returns<F, Fut>(f: F)
 116 where
 117     F: FnOnce(u8) -> Fut,
 118     Fut: Future<Output = u8>,
 119 {
 120     let mut fut = PinBox::new(f(9));
 121     let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
 122     let waker = local_waker_from_nonlocal(counter.clone());
 123     let executor = &mut NoopExecutor;
 124     let cx = &mut Context::new(&waker, executor);
 125
 126     assert_eq!(0, counter.wakes.load(atomic::Ordering::SeqCst));
 127     assert_eq!(Poll::Pending, fut.as_pin_mut().poll(cx));
 128     assert_eq!(1, counter.wakes.load(atomic::Ordering::SeqCst));
 129     assert_eq!(Poll::Ready(9), fut.as_pin_mut().poll(cx));
 130 }
 131
 132 fn main() {
 133     macro_rules! test {
 134         ($($fn_name:ident,)*) => { $(
 135             test_future_yields_once_then_returns($fn_name);
 136         )* }
 137     }
 138
 139     test! {
 140         async_block,
 141         async_nonmove_block,
 142         async_closure,
 143         async_fn,
 144         async_fn_with_internal_borrow,
 145     }
 146 }