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.
6///
7/// `fut_a` is polled first, so it wins when both futures are ready during the
8/// same poll. The other future is dropped.
9pub async fn select<A, B>(fut_a: A, fut_b: B) -> Either<A::Output, B::Output>
10where
11    A: Future,
12    B: Future,
13{
14    Select { fut_a, fut_b }.await
15}
16
17pin_project_lite::pin_project! {
18    pub(crate) struct Select<A, B> {
19        #[pin]
20        fut_a: A,
21        #[pin]
22        fut_b: B,
23    }
24}
25
26impl<A, B> Future for Select<A, B>
27where
28    A: Future,
29    B: Future,
30{
31    type Output = Either<A::Output, B::Output>;
32
33    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
34        let this = self.project();
35
36        if let Poll::Ready(item) = this.fut_a.poll(cx) {
37            return Poll::Ready(Either::Left(item));
38        }
39
40        if let Poll::Ready(item) = this.fut_b.poll(cx) {
41            return Poll::Ready(Either::Right(item));
42        }
43
44        Poll::Pending
45    }
46}
47
48#[cfg(test)]
49mod tests {
50    use futures_util::future::pending;
51    use std::future::ready;
52
53    use super::*;
54    use crate::time;
55
56    #[ntex::test]
57    async fn select_tests() {
58        let res = select(ready(Ok::<_, ()>("test")), pending::<()>()).await;
59        assert_eq!(res, Either::Left(Ok("test")));
60
61        let res = select(pending::<()>(), time::sleep(time::Millis(50))).await;
62        assert_eq!(res, Either::Right(()));
63    }
64}