4 // aux-build:arc_wake.rs
6 #![feature(async_await, async_closure)]
11 use std::future::Future;
14 atomic::{self, AtomicUsize},
16 use std::task::{Context, Poll};
17 use arc_wake::ArcWake;
23 impl ArcWake for Counter {
24 fn wake(self: Arc<Self>) {
25 Self::wake_by_ref(&self)
27 fn wake_by_ref(arc_self: &Arc<Self>) {
28 arc_self.wakes.fetch_add(1, atomic::Ordering::SeqCst);
32 struct WakeOnceThenComplete(bool);
34 fn wake_and_yield_once() -> WakeOnceThenComplete { WakeOnceThenComplete(false) }
36 impl Future for WakeOnceThenComplete {
38 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
42 cx.waker().wake_by_ref();
49 fn async_closure(x: u8) -> impl Future<Output = u8> {
50 (async move |x: u8| -> u8 {
51 wake_and_yield_once().await;
56 fn async_closure_in_unsafe_block(x: u8) -> impl Future<Output = u8> {
58 async move |x: u8| unsafe_fn(unsafe_async_fn(x).await)
62 async unsafe fn unsafe_async_fn(x: u8) -> u8 {
63 wake_and_yield_once().await;
67 unsafe fn unsafe_fn(x: u8) -> u8 {
71 fn test_future_yields_once_then_returns<F, Fut>(f: F)
74 Fut: Future<Output = u8>,
76 let mut fut = Box::pin(f(9));
77 let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
78 let waker = ArcWake::into_waker(counter.clone());
79 let mut cx = Context::from_waker(&waker);
80 assert_eq!(0, counter.wakes.load(atomic::Ordering::SeqCst));
81 assert_eq!(Poll::Pending, fut.as_mut().poll(&mut cx));
82 assert_eq!(1, counter.wakes.load(atomic::Ordering::SeqCst));
83 assert_eq!(Poll::Ready(9), fut.as_mut().poll(&mut cx));
88 ($($fn_name:expr,)*) => { $(
89 test_future_yields_once_then_returns($fn_name);
95 async_closure_in_unsafe_block,