zebra_network/
policies.rs

1use std::pin::Pin;
2
3use futures::{Future, FutureExt};
4use tower::retry::Policy;
5
6/// A very basic retry policy with a limited number of retry attempts.
7///
8/// TODO: Remove this when <https://github.com/tower-rs/tower/pull/414> lands.
9#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
10pub struct RetryLimit {
11    remaining_tries: usize,
12}
13
14impl RetryLimit {
15    /// Create a policy with the given number of retry attempts.
16    pub fn new(retry_attempts: usize) -> Self {
17        RetryLimit {
18            remaining_tries: retry_attempts,
19        }
20    }
21}
22
23impl<Req: Clone + std::fmt::Debug, Res, E: std::fmt::Debug> Policy<Req, Res, E> for RetryLimit {
24    type Future = Pin<Box<dyn Future<Output = Self> + Send + Sync + 'static>>;
25
26    fn retry(&self, req: &Req, result: Result<&Res, &E>) -> Option<Self::Future> {
27        if let Err(e) = result {
28            if self.remaining_tries > 0 {
29                tracing::debug!(?req, ?e, remaining_tries = self.remaining_tries, "retrying");
30
31                let remaining_tries = self.remaining_tries - 1;
32                let retry_outcome = RetryLimit { remaining_tries };
33
34                Some(
35                    // Let other tasks run, so we're more likely to choose a different peer,
36                    // and so that any notfound inv entries win the race to the PeerSet.
37                    //
38                    // # Security
39                    //
40                    // We want to choose different peers for retries, so we have a better chance of getting each block.
41                    // This is implemented by the connection state machine sending synthetic `notfound`s to the
42                    // `InventoryRegistry`, as well as forwarding actual `notfound`s from peers.
43                    Box::pin(tokio::task::yield_now().map(move |()| retry_outcome)),
44                )
45            } else {
46                None
47            }
48        } else {
49            None
50        }
51    }
52
53    fn clone_request(&self, req: &Req) -> Option<Req> {
54        Some(req.clone())
55    }
56}