Skip to content

Commit

Permalink
Add basic retry policies to zebra-network.
Browse files Browse the repository at this point in the history
This should be removed when tower-rs/tower#414 lands
but is good enough for our purposes for now.
  • Loading branch information
hdevalence committed Feb 11, 2020
1 parent 5191b9d commit acccdda
Show file tree
Hide file tree
Showing 2 changed files with 63 additions and 0 deletions.
2 changes: 2 additions & 0 deletions zebra-network/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ mod config;
mod constants;
mod meta_addr;
mod network;
mod policies;
mod peer;
mod peer_set;
mod protocol;
Expand All @@ -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.
Expand Down
61 changes: 61 additions & 0 deletions zebra-network/src/policies.rs
Original file line number Diff line number Diff line change
@@ -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<Req: Clone, Res, E> Policy<Req, Res, E> for RetryLimit {
type Future = future::Ready<Self>;
fn retry(&self, _: &Req, result: Result<&Res, &E>) -> Option<Self::Future> {
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<Req> {
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<Req: Clone, Res, E> Policy<Req, Res, E> for RetryErrors {
type Future = future::Ready<Self>;
fn retry(&self, _: &Req, result: Result<&Res, &E>) -> Option<Self::Future> {
if result.is_err() {
Some(future::ready(RetryErrors))
} else {
None
}
}

fn clone_request(&self, req: &Req) -> Option<Req> {
Some(req.clone())
}
}

0 comments on commit acccdda

Please sign in to comment.