karyon_core/async_util/
select.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

use pin_project_lite::pin_project;

/// Returns the result of the future that completes first, preferring future1
/// if both are ready.
///
/// # Examples
///
/// ```
/// use std::future;
///
/// use karyon_core::async_util::{select, Either};
///
///  async {
///     let fut1 = future::pending::<String>();
///     let fut2 = future::ready(0);
///     let res = select(fut1, fut2).await;
///     assert!(matches!(res, Either::Right(0)));
///     // ....
///  };
///
/// ```
///
pub fn select<T1, T2, F1, F2>(future1: F1, future2: F2) -> Select<F1, F2>
where
    F1: Future<Output = T1>,
    F2: Future<Output = T2>,
{
    Select { future1, future2 }
}

pin_project! {
    #[derive(Debug)]
    pub struct Select<F1, F2> {
        #[pin]
        future1: F1,
        #[pin]
        future2: F2,
    }
}

/// The return value from the [`select`] function, indicating which future
/// completed first.
#[derive(Debug)]
pub enum Either<T1, T2> {
    Left(T1),
    Right(T2),
}

// Implement the Future trait for the Select struct.
impl<T1, T2, F1, F2> Future for Select<F1, F2>
where
    F1: Future<Output = T1>,
    F2: Future<Output = T2>,
{
    type Output = Either<T1, T2>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();

        if let Poll::Ready(t) = this.future1.poll(cx) {
            return Poll::Ready(Either::Left(t));
        }

        if let Poll::Ready(t) = this.future2.poll(cx) {
            return Poll::Ready(Either::Right(t));
        }

        Poll::Pending
    }
}

#[cfg(test)]
mod tests {
    use std::future;

    use crate::{async_runtime::block_on, async_util::sleep};

    use super::{select, Either};

    #[test]
    fn test_async_select() {
        block_on(async move {
            let fut = select(sleep(std::time::Duration::MAX), future::ready(0 as u32)).await;
            assert!(matches!(fut, Either::Right(0)));

            let fut1 = future::pending::<String>();
            let fut2 = future::ready(0);
            let res = select(fut1, fut2).await;
            assert!(matches!(res, Either::Right(0)));

            let fut1 = future::ready(0);
            let fut2 = future::pending::<String>();
            let res = select(fut1, fut2).await;
            assert!(matches!(res, Either::Left(_)));
        });
    }
}