Skip to main content

ntex_util/future/
select.rs

1use std::{future::Future, pin::Pin, task::Context, task::Poll};
2
3use crate::future::Either;
4
5/// Waits for either one of two differently-typed futures to complete.
6pub async fn select<A, B>(fut_a: A, fut_b: B) -> Either<A::Output, B::Output>
7where
8    A: Future,
9    B: Future,
10{
11    Select { fut_a, fut_b }.await
12}
13
14pin_project_lite::pin_project! {
15    pub(crate) struct Select<A, B> {
16        #[pin]
17        fut_a: A,
18        #[pin]
19        fut_b: B,
20    }
21}
22
23impl<A, B> Future for Select<A, B>
24where
25    A: Future,
26    B: Future,
27{
28    type Output = Either<A::Output, B::Output>;
29
30    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
31        let this = self.project();
32
33        if let Poll::Ready(item) = this.fut_a.poll(cx) {
34            return Poll::Ready(Either::Left(item));
35        }
36
37        if let Poll::Ready(item) = this.fut_b.poll(cx) {
38            return Poll::Ready(Either::Right(item));
39        }
40
41        Poll::Pending
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use futures_util::future::pending;
48    use std::future::ready;
49
50    use super::*;
51    use crate::time;
52
53    #[ntex::test]
54    async fn select_tests() {
55        let res = select(ready(Ok::<_, ()>("test")), pending::<()>()).await;
56        assert_eq!(res, Either::Left(Ok("test")));
57
58        let res = select(pending::<()>(), time::sleep(time::Millis(50))).await;
59        assert_eq!(res, Either::Right(()));
60    }
61}