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::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::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_nonmove_block(x: u8) -> impl Future<Output = u8> {
  71     async move {
  72         let future = async {
  73             await!(wake_and_yield_once());
  74             x
  75         };
  76         await!(future)
  77     }
  78 }
  79
  80 fn async_closure(x: u8) -> impl Future<Output = u8> {
  81     (async move |x: u8| -> u8 {
  82         await!(wake_and_yield_once());
  83         x
  84     })(x)
  85 }
  86
  87 async fn async_fn(x: u8) -> u8 {
  88     await!(wake_and_yield_once());
  89     x
  90 }
  91
  92 async fn async_fn_with_borrow(x: &u8) -> u8 {
  93     await!(wake_and_yield_once());
  94     *x
  95 }
  96
  97 fn async_fn_with_internal_borrow(y: u8) -> impl Future<Output = u8> {
  98     async move {
  99         await!(async_fn_with_borrow(&y))
 100     }
 101 }
 102
 103 unsafe async fn unsafe_async_fn(x: u8) -> u8 {
 104     await!(wake_and_yield_once());
 105     x
 106 }
 107
 108 struct Foo;
 109
 110 trait Bar {
 111     fn foo() {}
 112 }
 113
 114 impl Foo {
 115     async fn async_method(x: u8) -> u8 {
 116         unsafe {
 117             await!(unsafe_async_fn(x))
 118         }
 119     }
 120 }
 121
 122 fn test_future_yields_once_then_returns<F, Fut>(f: F)
 123 where
 124     F: FnOnce(u8) -> Fut,
 125     Fut: Future<Output = u8>,
 126 {
 127     let mut fut = PinBox::new(f(9));
 128     let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
 129     let waker = local_waker_from_nonlocal(counter.clone());
 130     let spawner = &mut NoopSpawner;
 131     let cx = &mut Context::new(&waker, spawner);
 132
 133     assert_eq!(0, counter.wakes.load(atomic::Ordering::SeqCst));
 134     assert_eq!(Poll::Pending, fut.as_pin_mut().poll(cx));
 135     assert_eq!(1, counter.wakes.load(atomic::Ordering::SeqCst));
 136     assert_eq!(Poll::Ready(9), fut.as_pin_mut().poll(cx));
 137 }
 138
 139 fn main() {
 140     macro_rules! test {
 141         ($($fn_name:ident,)*) => { $(
 142             test_future_yields_once_then_returns($fn_name);
 143         )* }
 144     }
 145
 146     test! {
 147         async_block,
 148         async_nonmove_block,
 149         async_closure,
 150         async_fn,
 151         async_fn_with_internal_borrow,
 152     }
 153 }