async_utils/
futures.rs

1// Copyright 2020 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5//! Provides utilities for working with futures.
6
7use std::pin::Pin;
8
9use futures::future::FusedFuture;
10use futures::{task, Future};
11
12/// A future that yields to the executor only once.
13///
14/// This future returns [`Poll::Pending`] the first time it's polled after
15/// waking the context waker. This effectively yields the currently running task
16/// to the executor, but puts it back in the executor's ready task queue.
17///
18/// Example:
19/// ```
20/// loop {
21///   let read = read_big_thing().await;
22///
23///   while let Some(x) = read.next() {
24///     process_one_thing(x);
25///     YieldToExecutorOnce::new().await;
26///   }
27/// }
28/// ```
29#[derive(Default)]
30pub struct YieldToExecutorOnce(YieldToExecutorOnceInner);
31
32#[derive(Default)]
33enum YieldToExecutorOnceInner {
34    #[default]
35    NotPolled,
36    Ready,
37    Terminated,
38}
39
40impl YieldToExecutorOnce {
41    /// Creates a new `YieldToExecutorOnce`.
42    pub fn new() -> Self {
43        Self::default()
44    }
45}
46
47impl Future for YieldToExecutorOnce {
48    type Output = ();
49
50    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
51        let Self(inner) = self.get_mut();
52        match *inner {
53            YieldToExecutorOnceInner::NotPolled => {
54                *inner = YieldToExecutorOnceInner::Ready;
55                // Wake the executor before returning pending. We only want to yield
56                // once.
57                cx.waker().wake_by_ref();
58                task::Poll::Pending
59            }
60            YieldToExecutorOnceInner::Ready => {
61                *inner = YieldToExecutorOnceInner::Terminated;
62                task::Poll::Ready(())
63            }
64            YieldToExecutorOnceInner::Terminated => {
65                panic!("polled future after completion");
66            }
67        }
68    }
69}
70
71impl FusedFuture for YieldToExecutorOnce {
72    fn is_terminated(&self) -> bool {
73        let Self(inner) = self;
74        match inner {
75            YieldToExecutorOnceInner::Ready | YieldToExecutorOnceInner::NotPolled => false,
76            YieldToExecutorOnceInner::Terminated => true,
77        }
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    #[test]
84    fn yield_to_executor_once() {
85        use futures::future::FusedFuture as _;
86        use futures::FutureExt as _;
87
88        let (waker, count) = futures_test::task::new_count_waker();
89        let mut context = std::task::Context::from_waker(&waker);
90        let mut fut = super::YieldToExecutorOnce::new();
91
92        assert!(!fut.is_terminated());
93        assert_eq!(count, 0);
94        assert_eq!(fut.poll_unpin(&mut context), std::task::Poll::Pending);
95        assert!(!fut.is_terminated());
96        assert_eq!(count, 1);
97        assert_eq!(fut.poll_unpin(&mut context), std::task::Poll::Ready(()));
98        assert!(fut.is_terminated());
99        // The waker is never hit again.
100        assert_eq!(count, 1);
101    }
102}
async_utils/futures.rs

async_utils/
futures.rs
