diff --git a/zebra-network/src/lib.rs b/zebra-network/src/lib.rs index 3670b12b..afd13cd1 100644 --- a/zebra-network/src/lib.rs +++ b/zebra-network/src/lib.rs @@ -52,6 +52,7 @@ mod config; mod constants; mod meta_addr; mod network; +mod policies; mod peer; mod peer_set; mod protocol; @@ -62,6 +63,7 @@ pub use crate::{ config::Config, peer_set::init, protocol::internal::{Request, Response}, + policies::{RetryErrors, RetryLimit}, }; /// Types used in the definition of [`Request`] and [`Response`] messages. diff --git a/zebra-network/src/policies.rs b/zebra-network/src/policies.rs new file mode 100644 index 00000000..61a8f237 --- /dev/null +++ b/zebra-network/src/policies.rs @@ -0,0 +1,61 @@ +use tower::retry::Policy; +use futures::future; + +/// A very basic retry policy with a limited number of retry attempts. +/// +/// XXX Remove this when https://github.com/tower-rs/tower/pull/414 lands. +#[derive(Clone, Debug)] +pub struct RetryLimit { + remaining_tries: usize, +} + +impl RetryLimit { + /// Create a policy with the given number of retry attempts. + pub fn new(retry_attempts: usize) -> Self { + RetryLimit { + remaining_tries: retry_attempts, + } + } +} + +impl Policy for RetryLimit { + type Future = future::Ready; + fn retry(&self, _: &Req, result: Result<&Res, &E>) -> Option { + if result.is_err() { + if self.remaining_tries > 0 { + Some(future::ready(RetryLimit { + remaining_tries: self.remaining_tries - 1, + })) + } else { + None + } + } else { + None + } + } + + fn clone_request(&self, req: &Req) -> Option { + Some(req.clone()) + } +} + +/// A very basic retry policy that always retries failed requests. +/// +/// XXX remove this when https://github.com/tower-rs/tower/pull/414 lands. +#[derive(Clone, Debug)] +pub struct RetryErrors; + +impl Policy for RetryErrors { + type Future = future::Ready; + fn retry(&self, _: &Req, result: Result<&Res, &E>) -> Option { + if result.is_err() { + Some(future::ready(RetryErrors)) + } else { + None + } + } + + fn clone_request(&self, req: &Req) -> Option { + Some(req.clone()) + } +}