1 //! Functions for task management that don't belong inside the Runtime
2 //! trait.
3
4 use std::future::Future;
5 use std::pin::Pin;
6 use std::task::{Context, Poll};
7
8 /// Yield execution back to the runtime temporarily, so that other
9 /// tasks can run.
10 #[must_use = "yield_now returns a future that must be .awaited on."]
yield_now() -> YieldFuture11 pub fn yield_now() -> YieldFuture {
12 // TODO: There are functions similar to this in tokio and
13 // async_std and futures_lite. It would be lovely if futures had
14 // one too. If it does, we should probably use it.
15 YieldFuture { first_time: true }
16 }
17
18 /// A future returned by [`yield_now()`].
19 ///
20 /// It returns `Poll::Pending` once, and `Poll::Ready` thereafter.
21 #[derive(Debug)]
22 #[must_use = "Futures do nothing unless .awaited on."]
23 pub struct YieldFuture {
24 /// True if this future has not yet been polled.
25 first_time: bool,
26 }
27
28 impl Future for YieldFuture {
29 type Output = ();
poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()>30 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
31 if self.first_time {
32 self.first_time = false;
33 cx.waker().wake_by_ref();
34 Poll::Pending
35 } else {
36 Poll::Ready(())
37 }
38 }
39 }
40
41 #[cfg(all(test, any(feature = "tokio", feature = "async-std")))]
42 mod test {
43 use super::yield_now;
44 use crate::test_with_all_runtimes;
45
46 use std::sync::atomic::{AtomicBool, Ordering};
47
48 #[test]
test_yield() -> std::io::Result<()>49 fn test_yield() -> std::io::Result<()> {
50 test_with_all_runtimes!(|_| async {
51 let b = AtomicBool::new(false);
52 use Ordering::SeqCst;
53
54 // Both tasks here run in a loop, trying to set 'b' to their
55 // favorite value, and returning once they've done it 10 times.
56 //
57 // Without 'yield_now', one task is likely to monopolize
58 // the scheduler.
59 futures::join!(
60 async {
61 let mut n = 0_usize;
62 while n < 10 {
63 if b.compare_exchange(false, true, SeqCst, SeqCst).is_ok() {
64 n += 1;
65 }
66 yield_now().await;
67 }
68 },
69 async {
70 let mut n = 0_usize;
71 while n < 10 {
72 if b.compare_exchange(true, false, SeqCst, SeqCst).is_ok() {
73 n += 1;
74 }
75 yield_now().await;
76 }
77 }
78 );
79 std::io::Result::Ok(())
80 })?;
81 Ok(())
82 }
83 }
84