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
49    use super::*;
50    use crate::{future::Ready, time};
51
52    #[ntex_macros::rt_test2]
53    async fn select_tests() {
54        let res = select(Ready::<_, ()>::Ok("test"), pending::<()>()).await;
55        assert_eq!(res, Either::Left(Ok("test")));
56
57        let res = select(pending::<()>(), time::sleep(time::Millis(50))).await;
58        assert_eq!(res, Either::Right(()));
59    }
60}