saluki_io/net/util/middleware/

retry_circuit_breaker.rs

use std::{
    fmt,
    future::Future,
    pin::Pin,
    sync::{Arc, Mutex},
    task::{ready, Context, Poll},
};

use futures::FutureExt as _;
use pin_project_lite::pin_project;
use tokio_util::sync::{CancellationToken, WaitForCancellationFutureOwned};
use tower::{retry::Policy, Layer, Service};
use tracing::debug;

/// An error from [`RetryCircuitBreaker`].
#[derive(Debug)]
pub enum Error<E, R> {
    /// The inner service returned a readiness error or a final request error that will not be retried.
    Service(E),

    /// The inner service was called and completed, and `Policy::retry` returned `Some(backoff)`.
    Retry(R),

    /// The circuit breaker was already in backoff at call time, so the inner service was not called.
    Open(R),

    /// The cancellation token was cancelled while the inner service was handling the request, so the request was
    /// abandoned before it completed.
    Cancelled(R),
}

impl<E, R> std::error::Error for Error<E, R>
where
    E: std::error::Error + 'static,
    R: fmt::Debug,
{
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::Service(e) => Some(e),
            Self::Retry(_) | Self::Open(_) | Self::Cancelled(_) => None,
        }
    }
}

impl<E, R> fmt::Display for Error<E, R>
where
    E: fmt::Display,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Service(e) => write!(f, "service error: {}", e),
            Self::Retry(_) => write!(f, "request should be retried"),
            Self::Open(_) => write!(f, "circuit breaker open"),
            Self::Cancelled(_) => write!(f, "request cancelled"),
        }
    }
}

impl<E, R> PartialEq for Error<E, R>
where
    E: PartialEq,
    R: PartialEq,
{
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Self::Service(a), Self::Service(b)) => a == b,
            (Self::Retry(a), Self::Retry(b)) => a == b,
            (Self::Open(a), Self::Open(b)) => a == b,
            (Self::Cancelled(a), Self::Cancelled(b)) => a == b,
            _ => false,
        }
    }
}

pin_project! {
    /// Response future for [`RetryCircuitBreaker`].
    pub struct ResponseFuture<P, F, Request> {
        state: Arc<Mutex<State<P>>>,
        #[pin]
        inner: Option<F>,
        req: Option<Request>,
        #[pin]
        cancelled: Option<WaitForCancellationFutureOwned>,
    }
}

impl<P, F, T, E, Request> Future for ResponseFuture<P, F, Request>
where
    P: Policy<Request, T, E>,
    P::Future: Send + 'static,
    F: Future<Output = Result<T, E>>,
{
    type Output = Result<T, Error<E, Request>>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // Our response future exists in two states: the circuit breaker was either closed or open when we created it.
        //
        // When the circuit breaker is open while creating the response future, there's no actual response future to
        // call in this case. We simply store the original request and pass it back while indicating to the caller that
        // the circuit breaker is open. Simple.
        //
        // When the circuit breaker is closed while creating the response future, this means we can proceed, and we
        // generate a legitimate response future to poll. However, the retry policy may return `None` when trying to
        // clone the request, which indicates the request actually isn't eligible to be retried at all. Thus, when we
        // don't have an original request here, we just return the inner service's response as-is. When we _do_ have the
        // original request, we utilize the retry policy to determine if it can be retried, and if so, potentially
        // update our circuit breaker state based on what the retry policy tells us.
        //
        // When a cancellation token is configured and we have the original request, we also stop waiting for the
        // inner response once the token is cancelled, and pass the original request back instead. The response is
        // polled first, so a response that is already available wins over a cancellation.

        let this = self.project();
        if let Some(inner) = this.inner.as_pin_mut() {
            let mut result = match inner.poll(cx) {
                Poll::Ready(result) => result,
                Poll::Pending => {
                    if let Some(cancelled) = this.cancelled.as_pin_mut() {
                        ready!(cancelled.poll(cx));

                        debug!("request cancelled while in flight, returning error");
                        return Poll::Ready(Err(Error::Cancelled(
                            this.req.take().expect("response future polled after completion"),
                        )));
                    }
                    return Poll::Pending;
                }
            };

            let mut state = this.state.lock().unwrap();
            match this.req.take() {
                Some(mut req) => match state.policy.retry(&mut req, &mut result) {
                    Some(backoff) => {
                        // The policy has indicated that the request should be retried, so we need to open the circuit
                        // breaker by setting the backoff future to use. Another request's retry decision may have
                        // already beat us to the punch, though, so don't overwrite it if it's already set.
                        if state.backoff.is_none() {
                            debug!("no existing backoff future present, setting delay backoff");
                            state.backoff = Some(backoff.boxed());
                        }

                        Poll::Ready(Err(Error::Retry(req)))
                    }
                    None => {
                        debug!("request completed, no retry indicated");
                        Poll::Ready(result.map_err(Error::Service))
                    }
                },
                None => {
                    debug!("request completed, but request not cloneable so returning response as-is");
                    Poll::Ready(result.map_err(Error::Service))
                }
            }
        } else {
            debug!("circuit breaker open prior to call, returning error");
            Poll::Ready(Err(Error::Open(
                this.req.take().expect("response future polled after completion"),
            )))
        }
    }
}

struct State<P> {
    policy: P,
    backoff: Option<Pin<Box<dyn Future<Output = ()> + Send>>>,
}

impl<P> std::fmt::Debug for State<P>
where
    P: std::fmt::Debug,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let backoff = if self.backoff.is_some() { "set" } else { "unset" };
        f.debug_struct("State")
            .field("policy", &self.policy)
            .field("backoff", &backoff)
            .finish()
    }
}

/// Wraps a service in a [circuit breaker][circuit_breaker] and signals when a request must be retried at a later time.
///
/// [circuit_breaker]: https://en.wikipedia.org/wiki/Circuit_breaker_design_pattern
pub struct RetryCircuitBreakerLayer<P> {
    policy: P,
    cancellation: Option<CancellationToken>,
}

impl<P> RetryCircuitBreakerLayer<P> {
    /// Creates a new [`RetryCircuitBreakerLayer`] with the given policy.
    pub const fn new(policy: P) -> Self {
        Self {
            policy,
            cancellation: None,
        }
    }

    /// Sets the cancellation token for in-flight requests.
    ///
    /// See [`RetryCircuitBreaker::with_cancellation`] for details.
    pub fn with_cancellation(mut self, token: CancellationToken) -> Self {
        self.cancellation = Some(token);
        self
    }
}

impl<P, S> Layer<S> for RetryCircuitBreakerLayer<P>
where
    P: Clone,
{
    type Service = RetryCircuitBreaker<S, P>;

    fn layer(&self, inner: S) -> Self::Service {
        let service = RetryCircuitBreaker::new(inner, self.policy.clone());
        match &self.cancellation {
            Some(token) => service.with_cancellation(token.clone()),
            None => service,
        }
    }
}

/// Wraps a service in a [circuit breaker][circuit_breaker] and signals when a request must be retried at a later time.
///
/// This circuit breaker implementation is specific to retrying requests. In many cases, a request can fail in two
/// ways: unrecoverable errors, which shouldn't be retried, and recoverable errors, which should be retried after a
/// some period of time. When a request can be retried, it may not be advantageous to wait for the given request to
/// be retried successfully, as the request should perhaps be stored in a queue and retried at a later time,
/// potentially to avoid applying backpressure to the client.
///
/// [`RetryCircuitBreaker`] provides this capability by separating the logic of determining whether or not a request
/// should be retried from actually performing the retry itself. When a request leads to an unrecoverable error,
/// that error is immediately passed back to the caller without affecting the circuit breaker state. However, when a
/// recoverable error is encountered, the circuit breaker will signal to the caller that the request should be
/// retried, and update its internal state to open the circuit breaker for a configurable period of time. Further
/// requests to the circuit breaker will be rejected with an error (indicating the open state) until that period of
/// time has passed.
///
/// [circuit_breaker]: https://en.wikipedia.org/wiki/Circuit_breaker_design_pattern
#[derive(Debug)]
pub struct RetryCircuitBreaker<S, P> {
    inner: S,
    state: Arc<Mutex<State<P>>>,
    cancellation: Option<CancellationToken>,
}

impl<S, P> RetryCircuitBreaker<S, P> {
    /// Creates a new [`RetryCircuitBreaker`].
    pub fn new(inner: S, policy: P) -> Self {
        Self {
            inner,
            state: Arc::new(Mutex::new(State { policy, backoff: None })),
            cancellation: None,
        }
    }

    /// Sets the cancellation token for in-flight requests.
    ///
    /// When `token` is cancelled, every request still waiting on the inner service is abandoned, and its response
    /// future resolves to [`Error::Cancelled`] with the copy of the request made by the retry policy. The circuit
    /// breaker state is not affected. A request that the retry policy could not clone has no copy to hand back, so it
    /// is not abandoned and runs to completion.
    pub fn with_cancellation(mut self, token: CancellationToken) -> Self {
        self.cancellation = Some(token);
        self
    }
}

impl<S, P, Request> Service<Request> for RetryCircuitBreaker<S, P>
where
    S: Service<Request>,
    P: Policy<Request, S::Response, S::Error>,
    P::Future: Send + 'static,
{
    type Response = S::Response;
    type Error = Error<S::Error, Request>;
    type Future = ResponseFuture<P, S::Future, Request>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        {
            // Check if we're currently in a backoff state.
            let mut state = self.state.lock().unwrap();
            if let Some(backoff) = state.backoff.as_mut() {
                ready!(backoff.as_mut().poll(cx));

                debug!("circuit breaker backoff complete");

                // The backoff future has completed, so we can reset the circuit breaker state.
                state.backoff = None;
            }
        }

        // Check the readiness of the inner service.
        self.inner.poll_ready(cx).map_err(Error::Service)
    }

    fn call(&mut self, req: Request) -> Self::Future {
        let response_state = Arc::clone(&self.state);

        let mut state = self.state.lock().unwrap();
        if state.backoff.is_some() {
            ResponseFuture {
                state: response_state,
                inner: None,
                req: Some(req),
                cancelled: None,
            }
        } else {
            // The circuit breaker is closed, so we can proceed with the request.
            let cloned_req = state.policy.clone_request(&req);
            let inner = self.inner.call(req);
            let cancelled = match (&self.cancellation, &cloned_req) {
                (Some(token), Some(_)) => Some(token.clone().cancelled_owned()),
                _ => None,
            };

            ResponseFuture {
                state: response_state,
                inner: Some(inner),
                req: cloned_req,
                cancelled,
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::{
        future::{ready, Ready},
        time::Duration,
    };

    use tokio::time::Sleep;
    use tokio_test::{assert_pending, assert_ready_ok};
    use tower::{retry::Policy, ServiceExt as _};

    use super::*;

    const BACKOFF_DUR: Duration = Duration::from_secs(1);

    #[derive(Clone, Debug, Eq, PartialEq)]
    enum BasicRequest {
        Ok(String),
        Err(String),
    }

    impl BasicRequest {
        fn success<S: AsRef<str>>(value: S) -> Self {
            Self::Ok(value.as_ref().to_string())
        }

        fn failure<S: AsRef<str>>(value: S) -> Self {
            Self::Err(value.as_ref().to_string())
        }

        fn as_service_response(&self) -> Result<String, Error<String, Self>> {
            match self {
                Self::Ok(value) => Ok(value.clone()),
                Self::Err(value) => Err(Error::Service(value.clone())),
            }
        }

        fn as_retry_response(&self) -> Result<String, Error<String, Self>> {
            Err(Error::Retry(self.clone()))
        }

        fn as_open_response(&self) -> Result<String, Error<String, Self>> {
            Err(Error::Open(self.clone()))
        }

        fn as_cancelled_response(&self) -> Result<String, Error<String, Self>> {
            Err(Error::Cancelled(self.clone()))
        }
    }

    impl PartialEq<Result<String, String>> for BasicRequest {
        fn eq(&self, other: &Result<String, String>) -> bool {
            match self {
                Self::Ok(value) => other.as_ref() == Ok(value),
                Self::Err(value) => other.as_ref() == Err(value),
            }
        }
    }

    #[derive(Debug)]
    struct LoopbackService;

    impl Service<BasicRequest> for LoopbackService {
        type Response = String;
        type Error = String;
        type Future = Ready<Result<String, String>>;

        fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
            Poll::Ready(Ok(()))
        }

        fn call(&mut self, req: BasicRequest) -> Self::Future {
            let res = match req {
                BasicRequest::Ok(value) => Ok(value),
                BasicRequest::Err(value) => Err(value),
            };
            ready(res)
        }
    }

    /// Like [`LoopbackService`], but each response takes [`RESPONSE_DELAY`] to arrive.
    #[derive(Debug)]
    struct DelayedLoopbackService;

    const RESPONSE_DELAY: Duration = Duration::from_secs(5);

    impl Service<BasicRequest> for DelayedLoopbackService {
        type Response = String;
        type Error = String;
        type Future = Pin<Box<dyn Future<Output = Result<String, String>> + Send>>;

        fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
            Poll::Ready(Ok(()))
        }

        fn call(&mut self, req: BasicRequest) -> Self::Future {
            let response = LoopbackService.call(req);
            Box::pin(async move {
                tokio::time::sleep(RESPONSE_DELAY).await;
                response.await
            })
        }
    }

    #[derive(Clone, Debug)]
    struct CloneableTestRetryPolicy;

    impl<Req, T, E> Policy<Req, T, E> for CloneableTestRetryPolicy
    where
        Req: Clone,
    {
        type Future = Sleep;

        fn retry(&mut self, _: &mut Req, res: &mut Result<T, E>) -> Option<Self::Future> {
            match res {
                Ok(_) => None,
                Err(_) => Some(tokio::time::sleep(BACKOFF_DUR)),
            }
        }

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

    #[derive(Debug)]
    struct NonCloneableTestRetryPolicy;

    impl<Req, T, E> Policy<Req, T, E> for NonCloneableTestRetryPolicy {
        type Future = Sleep;

        fn retry(&mut self, _: &mut Req, res: &mut Result<T, E>) -> Option<Self::Future> {
            match res {
                Ok(_) => None,
                Err(_) => Some(tokio::time::sleep(BACKOFF_DUR)),
            }
        }

        fn clone_request(&mut self, _: &Req) -> Option<Req> {
            None
        }
    }

    #[tokio::test(start_paused = true)]
    async fn basic() {
        let good_req = BasicRequest::success("good");
        let bad_req = BasicRequest::failure("bad");

        let mut circuit_breaker = RetryCircuitBreaker::new(LoopbackService, CloneableTestRetryPolicy);

        // First request should succeed.
        //
        // We should see that it called through to the inner service.
        let svc = circuit_breaker.ready().await.expect("should never fail to be ready");
        let fut = svc.call(good_req.clone());
        let result = fut.await;
        assert_eq!(result, good_req.as_service_response());

        // Second request should fail and should be retried.
        //
        // We should see that it called through to the inner service
        let svc = circuit_breaker.ready().await.expect("should never fail to be ready");
        let fut = svc.call(bad_req.clone());
        let result = fut.await;
        assert_eq!(result, bad_req.as_retry_response());

        // When trying to make our third request, we should have to wait for the backoff duration before the service
        // indicates that it's ready for another call.
        let mut svc_fut = tokio_test::task::spawn(circuit_breaker.ready());
        assert_pending!(svc_fut.poll());

        // Advance time past the backoff duration, which should make our service ready.
        tokio::time::advance(BACKOFF_DUR + Duration::from_millis(1)).await;
        assert!(svc_fut.is_woken());
        let svc = assert_ready_ok!(svc_fut.poll());

        let fut = svc.call(good_req.clone());
        let result = fut.await;
        assert_eq!(result, good_req.as_service_response());

        // Fourth request should succeed unimpeded since the breaker is closed again.
        let svc = circuit_breaker.ready().await.expect("should never fail to be ready");
        let fut = svc.call(good_req.clone());
        let result = fut.await;
        assert_eq!(result, good_req.as_service_response());
    }

    #[tokio::test]
    async fn retry_policy_no_clone() {
        let good_req = BasicRequest::success("good");
        let bad_req = BasicRequest::failure("bad");

        // First request should succeed.
        //
        // We should see that it called through to the inner service.
        let mut circuit_breaker = RetryCircuitBreaker::new(LoopbackService, NonCloneableTestRetryPolicy);
        let svc = circuit_breaker.ready().await.expect("should never fail to be ready");
        let fut = svc.call(good_req.clone());
        let result = fut.await;
        assert_eq!(result, good_req.as_service_response());

        // Second request should fail and should be a service error, because without being able to clone the request, it
        // can't be retried anyways.
        let mut circuit_breaker = RetryCircuitBreaker::new(LoopbackService, NonCloneableTestRetryPolicy);
        let svc = circuit_breaker.ready().await.expect("should never fail to be ready");
        let fut = svc.call(bad_req.clone());
        let result = fut.await;
        assert_eq!(result, bad_req.as_service_response());
    }

    #[tokio::test(start_paused = true)]
    async fn concurrent_calls_can_advance() {
        let good_req = BasicRequest::success("good");
        let bad_req = BasicRequest::failure("bad");

        let mut circuit_breaker = RetryCircuitBreaker::new(LoopbackService, CloneableTestRetryPolicy);

        // First request should succeed. This is just a warmup.
        let svc = circuit_breaker.ready().await.expect("should never fail to be ready");
        let fut = svc.call(good_req.clone());
        let result = fut.await;
        assert_eq!(result, good_req.as_service_response());

        // Now we'll create two calls -- one that should fail and one that succeed -- but won't poll them until both are
        // created. This simulates two concurrent calls happening, and what we want to show is that the circuit breaker
        // should only mark itself as open to _new_ calls after it the state changes to open, and should not affect
        // running requests.
        let svc = circuit_breaker.ready().await.expect("should never fail to be ready");
        let bad_fut = svc.call(bad_req.clone());

        let svc = circuit_breaker.ready().await.expect("should never fail to be ready");
        let good_fut = svc.call(good_req.clone());

        let bad_result = bad_fut.await;
        assert_eq!(bad_result, bad_req.as_retry_response());

        let good_result = good_fut.await;
        assert_eq!(good_result, good_req.as_service_response());

        // Now we'll go to make a fourth request, and we'll manually check the readiness of the service to ensure that
        // we're now in a backoff state.
        let mut svc_fut = tokio_test::task::spawn(circuit_breaker.ready());
        assert_pending!(svc_fut.poll());
    }

    #[tokio::test(start_paused = true)]
    async fn breaker_open_between_ready_and_call() {
        let good_req = BasicRequest::success("good");
        let bad_req = BasicRequest::failure("bad");

        let mut circuit_breaker = RetryCircuitBreaker::new(LoopbackService, CloneableTestRetryPolicy);

        // We'll create two calls -- one that should fail and one that succeed -- but order their creation / polling such that the
        // bad request updates the breaker state to be open _before_ we create the good request. This is to exercise
        // that even though the service may report itself as ready, an in-flight request that completes and ultimately
        // changes the breaker state to open should cause subsequent calls to immediately fail.
        let svc = circuit_breaker.ready().await.expect("should never fail to be ready");
        let bad_fut = svc.call(bad_req.clone());

        // We're just making sure here that the service is ready to accept another call, but we're not creating that
        // call yet.
        let svc = circuit_breaker.ready().await.expect("should never fail to be ready");

        // Run our bad request first and ensure it reports that the completed request should be retried.
        let bad_result = bad_fut.await;
        assert_eq!(bad_result, bad_req.as_retry_response());

        // Now _create_ the good request and ensure that it fails with an open error without reaching the inner service.
        let good_fut = svc.call(good_req.clone());
        let good_result = good_fut.await;
        assert_eq!(good_result, good_req.as_open_response());
    }

    #[tokio::test(start_paused = true)]
    async fn cancellation_hands_back_in_flight_request() {
        let req = BasicRequest::success("in-flight");
        let token = CancellationToken::new();

        let mut circuit_breaker =
            RetryCircuitBreaker::new(DelayedLoopbackService, CloneableTestRetryPolicy).with_cancellation(token.clone());
        let svc = circuit_breaker.ready().await.expect("should never fail to be ready");
        let mut fut = tokio_test::task::spawn(svc.call(req.clone()));
        assert_pending!(fut.poll());

        token.cancel();
        assert!(fut.is_woken());
        assert_eq!(fut.await, req.as_cancelled_response());

        // A cancelled request doesn't open the breaker.
        let mut svc_fut = tokio_test::task::spawn(circuit_breaker.ready());
        assert_ready_ok!(svc_fut.poll());
    }

    #[tokio::test]
    async fn response_wins_over_cancellation() {
        let req = BasicRequest::success("good");
        let token = CancellationToken::new();
        token.cancel();

        // The response is available on the first poll, as is the cancellation. The response must win, so that a
        // request the remote end has answered isn't handed back to be sent again.
        let mut circuit_breaker =
            RetryCircuitBreaker::new(LoopbackService, CloneableTestRetryPolicy).with_cancellation(token);
        let svc = circuit_breaker.ready().await.expect("should never fail to be ready");
        assert_eq!(svc.call(req.clone()).await, req.as_service_response());
    }

    #[tokio::test(start_paused = true)]
    async fn cancellation_does_not_abandon_uncloned_request() {
        let req = BasicRequest::success("uncloned");
        let token = CancellationToken::new();

        let mut circuit_breaker = RetryCircuitBreaker::new(DelayedLoopbackService, NonCloneableTestRetryPolicy)
            .with_cancellation(token.clone());
        let svc = circuit_breaker.ready().await.expect("should never fail to be ready");
        let mut fut = tokio_test::task::spawn(svc.call(req.clone()));
        assert_pending!(fut.poll());

        // Without a copy of the request, there is nothing to hand back, so the request runs to completion.
        token.cancel();
        assert_pending!(fut.poll());

        tokio::time::advance(RESPONSE_DELAY).await;
        assert_eq!(fut.await, req.as_service_response());
    }

    #[tokio::test(start_paused = true)]
    async fn layer_applies_cancellation() {
        let req = BasicRequest::success("in-flight");
        let token = CancellationToken::new();

        let mut svc = RetryCircuitBreakerLayer::new(CloneableTestRetryPolicy)
            .with_cancellation(token.clone())
            .layer(DelayedLoopbackService);
        let svc = svc.ready().await.expect("should never fail to be ready");
        let fut = svc.call(req.clone());

        token.cancel();
        assert_eq!(fut.await, req.as_cancelled_response());
    }
}