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::PinBox;
  16 use std::pin::PinMut;
  17 use std::future::Future;
  18 use std::sync::{
  19     Arc,
  20     atomic::{self, AtomicUsize},
  21 };
  22 use std::future::FutureObj;
  23 use std::task::{
  24     Context, Poll, Wake,
  25     Spawn, SpawnObjError,
  26     local_waker_from_nonlocal,
  27 };
  28
  29 struct Counter {
  30     wakes: AtomicUsize,
  31 }
  32
  33 impl Wake for Counter {
  34     fn wake(this: &Arc<Self>) {
  35         this.wakes.fetch_add(1, atomic::Ordering::SeqCst);
  36     }
  37 }
  38
  39 struct NoopSpawner;
  40 impl Spawn for NoopSpawner {
  41     fn spawn_obj(&mut self, _: FutureObj<'static, ()>) -> Result<(), SpawnObjError> {
  42         Ok(())
  43     }
  44 }
  45
  46 struct WakeOnceThenComplete(bool);
  47
  48 fn wake_and_yield_once() -> WakeOnceThenComplete { WakeOnceThenComplete(false) }
  49
  50 impl Future for WakeOnceThenComplete {
  51     type Output = ();
  52     fn poll(mut self: PinMut<Self>, cx: &mut Context) -> Poll<()> {
  53         if self.0 {
  54             Poll::Ready(())
  55         } else {
  56             cx.waker().wake();
  57             self.0 = true;
  58             Poll::Pending
  59         }
  60     }
  61 }
  62
  63 fn async_block(x: u8) -> impl Future<Output = u8> {
  64     async move {
  65         await!(wake_and_yield_once());
  66         x
  67     }
  68 }
  69
  70 fn async_block_with_borrow_named_lifetime<'a>(x: &'a u8) -> impl Future<Output = u8> + 'a {
  71     async move {
  72         await!(wake_and_yield_once());
  73         *x
  74     }
  75 }
  76
  77 fn async_nonmove_block(x: u8) -> impl Future<Output = u8> {
  78     async move {
  79         let future = async {
  80             await!(wake_and_yield_once());
  81             x
  82         };
  83         await!(future)
  84     }
  85 }
  86
  87 fn async_closure(x: u8) -> impl Future<Output = u8> {
  88     (async move |x: u8| -> u8 {
  89         await!(wake_and_yield_once());
  90         x
  91     })(x)
  92 }
  93
  94 async fn async_fn(x: u8) -> u8 {
  95     await!(wake_and_yield_once());
  96     x
  97 }
  98
  99 async fn async_fn_with_borrow(x: &u8) -> u8 {
 100     await!(wake_and_yield_once());
 101     *x
 102 }
 103
 104 async fn async_fn_with_borrow_named_lifetime<'a>(x: &'a u8) -> u8 {
 105     await!(wake_and_yield_once());
 106     *x
 107 }
 108
 109 fn async_fn_with_impl_future_named_lifetime<'a>(x: &'a u8) -> impl Future<Output = u8> + 'a {
 110     async move {
 111         await!(wake_and_yield_once());
 112         *x
 113     }
 114 }
 115
 116 async fn async_fn_with_named_lifetime_multiple_args<'a>(x: &'a u8, _y: &'a u8) -> u8 {
 117     await!(wake_and_yield_once());
 118     *x
 119 }
 120
 121 fn async_fn_with_internal_borrow(y: u8) -> impl Future<Output = u8> {
 122     async move {
 123         await!(async_fn_with_borrow(&y))
 124     }
 125 }
 126
 127 unsafe async fn unsafe_async_fn(x: u8) -> u8 {
 128     await!(wake_and_yield_once());
 129     x
 130 }
 131
 132 struct Foo;
 133
 134 trait Bar {
 135     fn foo() {}
 136 }
 137
 138 impl Foo {
 139     async fn async_method(x: u8) -> u8 {
 140         unsafe {
 141             await!(unsafe_async_fn(x))
 142         }
 143     }
 144 }
 145
 146 fn test_future_yields_once_then_returns<F, Fut>(f: F)
 147 where
 148     F: FnOnce(u8) -> Fut,
 149     Fut: Future<Output = u8>,
 150 {
 151     let mut fut = PinBox::new(f(9));
 152     let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
 153     let waker = local_waker_from_nonlocal(counter.clone());
 154     let spawner = &mut NoopSpawner;
 155     let cx = &mut Context::new(&waker, spawner);
 156
 157     assert_eq!(0, counter.wakes.load(atomic::Ordering::SeqCst));
 158     assert_eq!(Poll::Pending, fut.as_pin_mut().poll(cx));
 159     assert_eq!(1, counter.wakes.load(atomic::Ordering::SeqCst));
 160     assert_eq!(Poll::Ready(9), fut.as_pin_mut().poll(cx));
 161 }
 162
 163 fn main() {
 164     macro_rules! test {
 165         ($($fn_name:expr,)*) => { $(
 166             test_future_yields_once_then_returns($fn_name);
 167         )* }
 168     }
 169
 170     macro_rules! test_with_borrow {
 171         ($($fn_name:expr,)*) => { $(
 172             test_future_yields_once_then_returns(|x| {
 173                 async move {
 174                     await!($fn_name(&x))
 175                 }
 176             });
 177         )* }
 178     }
 179
 180     test! {
 181         async_block,
 182         async_nonmove_block,
 183         async_closure,
 184         async_fn,
 185         async_fn_with_internal_borrow,
 186         |x| {
 187             async move {
 188                 unsafe { await!(unsafe_async_fn(x)) }
 189             }
 190         },
 191     }
 192
 193     test_with_borrow! {
 194         async_block_with_borrow_named_lifetime,
 195         async_fn_with_borrow,
 196         async_fn_with_borrow_named_lifetime,
 197         async_fn_with_impl_future_named_lifetime,
 198         |x| {
 199             async move {
 200                 await!(async_fn_with_named_lifetime_multiple_args(x, x))
 201             }
 202         },
 203     }
 204 }