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."]
11
16845
pub fn yield_now() -> YieldFuture {
12
16845
    // TODO: There are functions similar to this in tokio and
13
16845
    // async_std and futures_lite.  It would be lovely if futures had
14
16845
    // one too.  If it does, we should probably use it.
15
16845
    YieldFuture { first_time: true }
16
16845
}
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 = ();
30
33580
    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
31
33580
        if self.first_time {
32
16860
            self.first_time = false;
33
16860
            cx.waker().wake_by_ref();
34
16860
            Poll::Pending
35
        } else {
36
16720
            Poll::Ready(())
37
        }
38
33580
    }
39
}
40

            
41
#[cfg(all(
42
    test,
43
    any(feature = "native-tls", feature = "rustls"),
44
    any(feature = "tokio", feature = "async-std")
45
))]
46
mod test {
47
    use super::yield_now;
48
    use crate::test_with_all_runtimes;
49

            
50
    use std::sync::atomic::{AtomicBool, Ordering};
51

            
52
    #[test]
53
    fn test_yield() {
54
        test_with_all_runtimes!(|_| async {
55
            let b = AtomicBool::new(false);
56
            use Ordering::SeqCst;
57

            
58
            // Both tasks here run in a loop, trying to set 'b' to their
59
            // favorite value, and returning once they've done it 10 times.
60
            //
61
            // Without 'yield_now', one task is likely to monopolize
62
            // the scheduler.
63
            futures::join!(
64
                async {
65
                    let mut n = 0_usize;
66
                    while n < 10 {
67
                        if b.compare_exchange(false, true, SeqCst, SeqCst).is_ok() {
68
                            n += 1;
69
                        }
70
                        yield_now().await;
71
                    }
72
                },
73
                async {
74
                    let mut n = 0_usize;
75
                    while n < 10 {
76
                        if b.compare_exchange(true, false, SeqCst, SeqCst).is_ok() {
77
                            n += 1;
78
                        }
79
                        yield_now().await;
80
                    }
81
                }
82
            );
83
            std::io::Result::Ok(())
84
        });
85
    }
86
}