zebra_network/peer_set/
unready_service.rs

1//! Services that are busy or newly created.
2//!
3//! The [`UnreadyService`] implementation is adapted from the one in [tower::Balance][tower-balance].
4//!
5//! [tower-balance]: https://github.com/tower-rs/tower/tree/master/tower/src/balance
6
7use std::{
8    future::Future,
9    marker::PhantomData,
10    pin::Pin,
11    task::{Context, Poll},
12};
13
14use futures::{channel::oneshot, ready};
15use tokio::time::Sleep;
16use tower::Service;
17
18use crate::peer_set::set::CancelClientWork;
19
20#[cfg(test)]
21mod tests;
22
23/// A Future that becomes satisfied when an `S`-typed service is ready.
24///
25/// May fail due to cancellation or continuous unreadiness for `REQUEST_TIMEOUT`.
26/// This bound is separate from the response timeout, which starts at request dispatch.
27#[pin_project]
28#[derive(Debug)]
29pub(super) struct UnreadyService<K, S, Req> {
30    /// The key used to lookup `service`.
31    pub(super) key: Option<K>,
32
33    /// A oneshot used to cancel the request the `service` is currently working on, if any.
34    #[pin]
35    pub(super) cancel: oneshot::Receiver<CancelClientWork>,
36
37    /// Limits continuous unreadiness, starting when the service enters the unready set.
38    #[pin]
39    pub(super) readiness_timeout: Sleep,
40
41    /// The `service` that is busy (or newly created).
42    pub(super) service: Option<S>,
43
44    /// Dropping `service` might drop a request.
45    /// This [`PhantomData`] tells the Rust compiler to do a drop check for `Req`.
46    pub(super) _req: PhantomData<Req>,
47}
48
49#[derive(Debug, Eq, PartialEq)]
50pub(super) enum Error<E> {
51    Inner(E),
52    Canceled,
53    CancelHandleDropped(oneshot::Canceled),
54    TimedOut,
55}
56
57impl<K, S: Service<Req>, Req> Future for UnreadyService<K, S, Req> {
58    type Output = Result<(K, S), (K, Error<S::Error>)>;
59
60    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
61        let this = self.project();
62
63        if let Poll::Ready(oneshot_result) = this.cancel.poll(cx) {
64            let key = this.key.take().expect("polled after ready");
65
66            // # Correctness
67            //
68            // Return an error if the service is explicitly canceled,
69            // or its cancel handle is dropped, implicitly cancelling it.
70            match oneshot_result {
71                Ok(CancelClientWork) => return Poll::Ready(Err((key, Error::Canceled))),
72                Err(canceled_error) => {
73                    return Poll::Ready(Err((key, Error::CancelHandleDropped(canceled_error))))
74                }
75            }
76        }
77
78        // # Correctness
79        //
80        // The current task must be scheduled for wakeup every time we return
81        // `Poll::Pending`.
82        //
83        // Check readiness before the deadline so a service ready at the deadline is retained.
84        // If it is still unready, always poll the timer, even amid repeated service wakeups.
85        // Both the timer and `cancel.poll` also schedule this task for wakeup.
86        let res = match this
87            .service
88            .as_mut()
89            .expect("polled after ready")
90            .poll_ready(cx)
91        {
92            Poll::Ready(res) => res,
93            Poll::Pending => {
94                ready!(this.readiness_timeout.poll(cx));
95                let key = this.key.take().expect("polled after ready");
96                return Poll::Ready(Err((key, Error::TimedOut)));
97            }
98        };
99
100        let key = this.key.take().expect("polled after ready");
101        let svc = this.service.take().expect("polled after ready");
102
103        match res {
104            Ok(()) => Poll::Ready(Ok((key, svc))),
105            Err(e) => Poll::Ready(Err((key, Error::Inner(e)))),
106        }
107    }
108}
zebra_network/peer_set/unready_service.rs

zebra_network/peer_set/
unready_service.rs
