for thramp

retry_future.rs

extern crate futures;
// XXX I haven't actually tested this, but I think it should work.

use futures::{Async, Future, Poll};

pub struct RetryFuture<F, N> {
    current: F,
    new_future: N,
    retries: usize,
    max_retries: Option<usize>,
}

#[derive(Clone, Debug)]
pub struct RetryError<E> {
    error: E,
    after_retries: usize,
}

impl<F, N> RetryFuture<F, N>
where
    F: Future,
    N: FnMut() -> F,
{
    pub fn new(mut new_future: N) -> Self {
        Self {
            current: new_future(),
            new_future,
            retries: 0,
            max_retries: None,
        }
    }

    pub fn with_max_retries(self, max_retries: usize) -> Self {
        Self {
            max_retries: Some(max_retries),
            ..self
        }
    }
}

impl<F, N> Future for RetryFuture<F, N>
where
    F: Future,
    N: FnMut() -> F,
{
    type Item = F::Item;
    type Error = RetryError<F::Error>;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            match self.current.poll() {
                // The future finished successfully, return the result.
                Ok(Async::Ready(item)) => return Ok(Async::Ready(item)),

                // The future isn't ready yet, so this future isn't ready
                // either.
                Ok(Async::NotReady) => return Ok(Async::NotReady),

                Err(error) => {
                    // Retry budget exceeded.
                    if self.max_retries.as_ref() == Some(&self.retries) {
                        let error = RetryError {
                            error,
                            after_retries: self.retries,
                        };
                        return Err(error);
                    } else {
                        // The current future failed, but we have retries
                        // remaining. Increment the number of retries, and
                        // replace the current future with a new one.
                        self.retries += 1;
                        self.current = (self.new_future)();
                        // We will poll the new current future again on the next
                        // iteration of the loop.
                    };
                }
            }
        }
    }
}