feat(tower-fallback): support custom fallback policies (#10923)

* feat(tower-fallback): support custom fallback policies

Signed-off-by: weifanglab <weifanglab@outlook.com>

* Apply suggestions from code review

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

---------

Signed-off-by: weifanglab <weifanglab@outlook.com>
Co-authored-by: Conrado <conradoplg@gmail.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Conrado Gouvea <conrado@zfnd.org>
Co-authored-by: Alfredo Garcia <oxarbitrage@gmail.com>

Author: 4 people

tower-fallback/CHANGELOG.md

@@ -5,6 +5,13 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+
+### Added
+
+- Added `Fallback::new_with_policy()` for selecting fallback behavior based on
+  the first service's result.
+
 ## [0.2.42] - 2026-07-10
 
 ### Changed

tower-fallback/src/future.rs

@@ -11,19 +11,22 @@ use futures_core::ready;
 use pin_project::pin_project;
 use tower::Service;
 
-use crate::BoxedError;
+use crate::{BoxedError, FallbackPolicy, OnError};
 
 /// Future that completes either with the first service's successful response, or
 /// with the second service's response.
 #[pin_project]
-pub struct ResponseFuture<S1, S2, Request>
+pub struct ResponseFuture<S1, S2, Request, F = OnError>
 where
     S1: Service<Request>,
     S2: Service<Request, Response = <S1 as Service<Request>>::Response>,
+    F: FallbackPolicy<<S1 as Service<Request>>::Response>,
+    S1::Error: Into<BoxedError>,
     S2::Error: Into<BoxedError>,
 {
     #[pin]
     state: ResponseState<S1, S2, Request>,
+    policy: F,
 }
 
 #[pin_project(project_replace = __ResponseStateProjectionOwned, project = ResponseStateProj)]
@@ -51,23 +54,28 @@ where
     Tmp,
 }
 
-impl<S1, S2, Request> ResponseFuture<S1, S2, Request>
+impl<S1, S2, Request, F> ResponseFuture<S1, S2, Request, F>
 where
     S1: Service<Request>,
     S2: Service<Request, Response = <S1 as Service<Request>>::Response>,
+    F: FallbackPolicy<<S1 as Service<Request>>::Response>,
+    S1::Error: Into<BoxedError>,
     S2::Error: Into<BoxedError>,
 {
-    pub(crate) fn new(fut: S1::Future, req: Request, svc2: S2) -> Self {
+    pub(crate) fn new(fut: S1::Future, req: Request, svc2: S2, policy: F) -> Self {
         ResponseFuture {
             state: ResponseState::PollResponse1 { fut, req, svc2 },
+            policy,
         }
     }
 }
 
-impl<S1, S2, Request> Future for ResponseFuture<S1, S2, Request>
+impl<S1, S2, Request, F> Future for ResponseFuture<S1, S2, Request, F>
 where
     S1: Service<Request>,
     S2: Service<Request, Response = <S1 as Service<Request>>::Response>,
+    F: FallbackPolicy<<S1 as Service<Request>>::Response>,
+    S1::Error: Into<BoxedError>,
     S2::Error: Into<BoxedError>,
 {
     type Output = Result<<S1 as Service<Request>>::Response, BoxedError>;
@@ -83,19 +91,22 @@ where
         // only returns Pending when a future or service returns Pending.
         loop {
             match this.state.as_mut().project() {
-                ResponseStateProj::PollResponse1 { fut, .. } => match ready!(fut.poll(cx)) {
-                    Ok(rsp) => return Poll::Ready(Ok(rsp)),
-                    Err(_) => {
-                        tracing::debug!("got error from svc1, retrying on svc2");
+                ResponseStateProj::PollResponse1 { fut, .. } => {
+                    let result = ready!(fut.poll(cx)).map_err(Into::into);
+
+                    if this.policy.should_fallback(&result) {
+                        tracing::debug!("fallback policy selected svc2");
                         if let __ResponseStateProjectionOwned::PollResponse1 { req, svc2, .. } =
                             this.state.as_mut().project_replace(ResponseState::Tmp)
                         {
                             this.state.set(ResponseState::PollReady2 { req, svc2 });
                         } else {
                             unreachable!();
                         }
+                    } else {
+                        return Poll::Ready(result);
                     }
-                },
+                }
                 ResponseStateProj::PollReady2 { svc2, .. } => match ready!(svc2.poll_ready(cx)) {
                     Err(e) => return Poll::Ready(Err(e.into())),
                     Ok(()) => {
@@ -119,10 +130,12 @@ where
     }
 }
 
-impl<S1, S2, Request> Debug for ResponseFuture<S1, S2, Request>
+impl<S1, S2, Request, F> Debug for ResponseFuture<S1, S2, Request, F>
 where
     S1: Service<Request>,
     S2: Service<Request, Response = <S1 as Service<Request>>::Response>,
+    F: FallbackPolicy<<S1 as Service<Request>>::Response>,
+    S1::Error: Into<BoxedError>,
     Request: Debug,
     S1::Future: Debug,
     S2: Debug,

tower-fallback/src/lib.rs

@@ -1,5 +1,6 @@
 //! A service combinator that sends requests to a first service, then retries
-//! processing on a second fallback service if the first service errors.
+//! processing on a second fallback service if the first service errors, or if a
+//! custom fallback policy selects the fallback service.
 //!
 //! Fallback designs have [a number of downsides][aws-fallback] but may be useful
 //! in some cases. For instance, when using batch verification, the `Fallback`
@@ -13,7 +14,7 @@
 pub mod future;
 mod service;
 
-pub use self::service::Fallback;
+pub use self::service::{Fallback, FallbackPolicy, OnError};
 
 /// A boxed type-erased `std::error::Error` that can be sent between threads.
 pub type BoxedError = Box<dyn std::error::Error + Send + Sync + 'static>;

tower-fallback/src/service.rs

@@ -5,53 +5,102 @@ use tower::Service;
 use super::future::ResponseFuture;
 use crate::BoxedError;
 
-/// Provides fallback processing on a second service if the first service returned an error.
+/// Decides if a [`Fallback`] service should call its fallback service.
+pub trait FallbackPolicy<Response> {
+    /// Returns `true` if the fallback service should handle this request.
+    fn should_fallback(&self, result: &Result<Response, BoxedError>) -> bool;
+}
+
+impl<Response, F> FallbackPolicy<Response> for F
+where
+    F: Fn(&Result<Response, BoxedError>) -> bool,
+{
+    fn should_fallback(&self, result: &Result<Response, BoxedError>) -> bool {
+        self(result)
+    }
+}
+
+/// The default fallback policy.
+///
+/// Falls back whenever the first service returns an error.
+#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
+pub struct OnError;
+
+impl<Response> FallbackPolicy<Response> for OnError {
+    fn should_fallback(&self, result: &Result<Response, BoxedError>) -> bool {
+        result.is_err()
+    }
+}
+
+/// Provides fallback processing on a second service if its fallback policy selects it.
 #[derive(Debug)]
-pub struct Fallback<S1, S2>
+pub struct Fallback<S1, S2, F = OnError>
 where
     S2: Clone,
 {
     svc1: S1,
     svc2: S2,
+    policy: F,
 }
 
-impl<S1: Clone, S2: Clone> Clone for Fallback<S1, S2> {
+impl<S1: Clone, S2: Clone, F: Clone> Clone for Fallback<S1, S2, F> {
     fn clone(&self) -> Self {
         Self {
             svc1: self.svc1.clone(),
             svc2: self.svc2.clone(),
+            policy: self.policy.clone(),
         }
     }
 }
 
-impl<S1, S2: Clone> Fallback<S1, S2> {
+impl<S1, S2: Clone> Fallback<S1, S2, OnError> {
     /// Creates a new `Fallback` wrapping a pair of services.
     ///
     /// Requests are processed on `svc1`, and retried on `svc2` if `svc1` errored.
     pub fn new(svc1: S1, svc2: S2) -> Self {
-        Self { svc1, svc2 }
+        Self {
+            svc1,
+            svc2,
+            policy: OnError,
+        }
+    }
+}
+
+impl<S1, S2: Clone, F: Clone> Fallback<S1, S2, F> {
+    /// Creates a new `Fallback` wrapping a pair of services with a custom fallback policy.
+    ///
+    /// Requests are processed on `svc1`, and retried on `svc2` if `policy` returns `true`
+    /// for `svc1`'s result.
+    pub fn new_with_policy(svc1: S1, svc2: S2, policy: F) -> Self {
+        Self { svc1, svc2, policy }
     }
 }
 
-impl<S1, S2, Request> Service<Request> for Fallback<S1, S2>
+impl<S1, S2, F, Request> Service<Request> for Fallback<S1, S2, F>
 where
     S1: Service<Request>,
     S2: Service<Request, Response = <S1 as Service<Request>>::Response>,
+    F: FallbackPolicy<<S1 as Service<Request>>::Response> + Clone,
     S1::Error: Into<BoxedError>,
     S2::Error: Into<BoxedError>,
     S2: Clone,
     Request: Clone,
 {
     type Response = <S1 as Service<Request>>::Response;
     type Error = BoxedError;
-    type Future = ResponseFuture<S1, S2, Request>;
+    type Future = ResponseFuture<S1, S2, Request, F>;
 
     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         self.svc1.poll_ready(cx).map_err(Into::into)
     }
 
     fn call(&mut self, request: Request) -> Self::Future {
         let request2 = request.clone();
-        ResponseFuture::new(self.svc1.call(request), request2, self.svc2.clone())
+        ResponseFuture::new(
+            self.svc1.call(request),
+            request2,
+            self.svc2.clone(),
+            self.policy.clone(),
+        )
     }
 }

tower-fallback/tests/fallback.rs

@@ -1,7 +1,7 @@
 //! Tests for tower-fallback
 
 use tower::{service_fn, Service, ServiceExt};
-use tower_fallback::Fallback;
+use tower_fallback::{BoxedError, Fallback};
 
 #[tokio::test]
 async fn fallback() {
@@ -30,3 +30,35 @@ async fn fallback() {
     assert_eq!(svc.ready().await.unwrap().call(11).await.unwrap(), 111);
     assert!(svc.ready().await.unwrap().call(21).await.is_err());
 }
+
+#[tokio::test]
+async fn custom_fallback_policy() {
+    let _init_guard = zebra_test::init();
+
+    let svc1 = service_fn(|val: u64| async move {
+        if val < 10 {
+            Ok::<_, &'static str>(Some(val))
+        } else {
+            Ok(None)
+        }
+    });
+    let svc2 = service_fn(|val: u64| async move {
+        if val < 20 {
+            Ok(Some(100 + val))
+        } else {
+            Err("too big value on svc2")
+        }
+    });
+
+    let mut svc =
+        Fallback::new_with_policy(svc1, svc2, |result: &Result<Option<u64>, BoxedError>| {
+            matches!(result, Ok(None))
+        });
+
+    assert_eq!(svc.ready().await.unwrap().call(1).await.unwrap(), Some(1));
+    assert_eq!(
+        svc.ready().await.unwrap().call(11).await.unwrap(),
+        Some(111)
+    );
+    assert!(svc.ready().await.unwrap().call(21).await.is_err());
+}