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 // edition:2018
  12
  13 #![feature(arbitrary_self_types, async_await, await_macro, futures_api, pin)]
  14
  15 use std::pin::Pin;
  16 use std::future::Future;
  17 use std::sync::{
  18     Arc,
  19     atomic::{self, AtomicUsize},
  20 };
  21 use std::future::FutureObj;
  22 use std::task::{
  23     Context, Poll, Wake,
  24     Spawn, 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 NoopSpawner;
  39 impl Spawn for NoopSpawner {
  40     fn spawn_obj(&mut self, _: FutureObj<'static, ()>) -> 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: Pin<&mut 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_block_with_borrow_named_lifetime<'a>(x: &'a u8) -> impl Future<Output = u8> + 'a {
  70     async move {
  71         await!(wake_and_yield_once());
  72         *x
  73     }
  74 }
  75
  76 fn async_nonmove_block(x: u8) -> impl Future<Output = u8> {
  77     async move {
  78         let future = async {
  79             await!(wake_and_yield_once());
  80             x
  81         };
  82         await!(future)
  83     }
  84 }
  85
  86 fn async_closure(x: u8) -> impl Future<Output = u8> {
  87     (async move |x: u8| -> u8 {
  88         await!(wake_and_yield_once());
  89         x
  90     })(x)
  91 }
  92
  93 async fn async_fn(x: u8) -> u8 {
  94     await!(wake_and_yield_once());
  95     x
  96 }
  97
  98 async fn async_fn_with_borrow(x: &u8) -> u8 {
  99     await!(wake_and_yield_once());
 100     *x
 101 }
 102
 103 async fn async_fn_with_borrow_named_lifetime<'a>(x: &'a u8) -> u8 {
 104     await!(wake_and_yield_once());
 105     *x
 106 }
 107
 108 fn async_fn_with_impl_future_named_lifetime<'a>(x: &'a u8) -> impl Future<Output = u8> + 'a {
 109     async move {
 110         await!(wake_and_yield_once());
 111         *x
 112     }
 113 }
 114
 115 async fn async_fn_with_named_lifetime_multiple_args<'a>(x: &'a u8, _y: &'a u8) -> u8 {
 116     await!(wake_and_yield_once());
 117     *x
 118 }
 119
 120 fn async_fn_with_internal_borrow(y: u8) -> impl Future<Output = u8> {
 121     async move {
 122         await!(async_fn_with_borrow(&y))
 123     }
 124 }
 125
 126 unsafe async fn unsafe_async_fn(x: u8) -> u8 {
 127     await!(wake_and_yield_once());
 128     x
 129 }
 130
 131 struct Foo;
 132
 133 trait Bar {
 134     fn foo() {}
 135 }
 136
 137 impl Foo {
 138     async fn async_method(x: u8) -> u8 {
 139         unsafe {
 140             await!(unsafe_async_fn(x))
 141         }
 142     }
 143 }
 144
 145 fn test_future_yields_once_then_returns<F, Fut>(f: F)
 146 where
 147     F: FnOnce(u8) -> Fut,
 148     Fut: Future<Output = u8>,
 149 {
 150     let mut fut = Box::pinned(f(9));
 151     let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
 152     let waker = local_waker_from_nonlocal(counter.clone());
 153     let spawner = &mut NoopSpawner;
 154     let cx = &mut Context::new(&waker, spawner);
 155
 156     assert_eq!(0, counter.wakes.load(atomic::Ordering::SeqCst));
 157     assert_eq!(Poll::Pending, fut.as_mut().poll(cx));
 158     assert_eq!(1, counter.wakes.load(atomic::Ordering::SeqCst));
 159     assert_eq!(Poll::Ready(9), fut.as_mut().poll(cx));
 160 }
 161
 162 fn main() {
 163     macro_rules! test {
 164         ($($fn_name:expr,)*) => { $(
 165             test_future_yields_once_then_returns($fn_name);
 166         )* }
 167     }
 168
 169     macro_rules! test_with_borrow {
 170         ($($fn_name:expr,)*) => { $(
 171             test_future_yields_once_then_returns(|x| {
 172                 async move {
 173                     await!($fn_name(&x))
 174                 }
 175             });
 176         )* }
 177     }
 178
 179     test! {
 180         async_block,
 181         async_nonmove_block,
 182         async_closure,
 183         async_fn,
 184         async_fn_with_internal_borrow,
 185         |x| {
 186             async move {
 187                 unsafe { await!(unsafe_async_fn(x)) }
 188             }
 189         },
 190     }
 191
 192     test_with_borrow! {
 193         async_block_with_borrow_named_lifetime,
 194         async_fn_with_borrow,
 195         async_fn_with_borrow_named_lifetime,
 196         async_fn_with_impl_future_named_lifetime,
 197         |x| {
 198             async move {
 199                 await!(async_fn_with_named_lifetime_multiple_args(x, x))
 200             }
 201         },
 202     }
 203 }