improve flexiblity of retry policy

tower/src/retry/future.rs

@@ -65,9 +65,9 @@ where
         loop {
             match this.state.as_mut().project() {
                 StateProj::Called(future) => {
-                    let result = ready!(future.poll(cx));
-                    if let Some(ref req) = this.request {
-                        match this.retry.policy.retry(req, result.as_ref()) {
+                    let mut result = ready!(future.poll(cx));
+                    if let Some(ref mut req) = this.request {
+                        match this.retry.policy.retry(req, &mut result) {
                             Some(checking) => {
                                 this.state.set(State::Checking(checking));
                             }

tower/src/retry/policy.rs

@@ -16,7 +16,7 @@ use std::future::Future;
 /// impl<E> Policy<Req, Res, E> for Attempts {
 ///     type Future = future::Ready<Self>;
 ///
-///     fn retry(&self, req: &Req, result: Result<&Res, &E>) -> Option<Self::Future> {
+///     fn retry(&self, req: &mut Req, result: &mut Result<Res, E>) -> Option<Self::Future> {
 ///         match result {
 ///             Ok(_) => {
 ///                 // Treat all `Response`s as success,
@@ -58,7 +58,7 @@ pub trait Policy<Req, Res, E>: Sized {
     ///
     /// [`Service::Response`]: crate::Service::Response
     /// [`Service::Error`]: crate::Service::Error
-    fn retry(&self, req: &Req, result: Result<&Res, &E>) -> Option<Self::Future>;
+    fn retry(&self, req: &mut Req, result: &mut Result<Res, E>) -> Option<Self::Future>;
 
     /// Tries to clone a request before being passed to the inner service.
     ///

tower/tests/builder.rs

@@ -45,7 +45,7 @@ where
 {
     type Future = Ready<Self>;
 
-    fn retry(&self, _req: &Req, _result: Result<&Res, &E>) -> Option<Self::Future> {
+    fn retry(&self, _req: &mut Req, _result: &mut Result<Res, E>) -> Option<Self::Future> {
         None
     }
 

tower/tests/retry/main.rs

@@ -90,6 +90,30 @@ async fn success_with_cannot_clone() {
     assert_ready_ok!(fut.poll(), "world");
 }
 
+#[tokio::test(flavor = "current_thread")]
+async fn retry_mutating_policy() {
+    let _t = support::trace_init();
+
+    // Even though the request couldn't be cloned, if the first request succeeds,
+    // it should succeed overall.
+    let (mut service, mut handle) = new_service(MutatingPolicy { remaining: 2 });
+
+    assert_ready_ok!(service.poll_ready());
+    let mut fut = task::spawn(service.call("hello"));
+
+    assert_request_eq!(handle, "hello").send_response("world");
+    assert_pending!(fut.poll());
+    // the policy alters the request. in real life, this might be setting
+    // a header
+    assert_request_eq!(handle, "retrying").send_response("world");
+
+    assert_pending!(fut.poll());
+
+    assert_request_eq!(handle, "retrying").send_response("world");
+
+    assert_ready_err!(fut.poll(), "out of retries");
+}
+
 type Req = &'static str;
 type Res = &'static str;
 type InnerError = &'static str;
@@ -102,7 +126,7 @@ struct RetryErrors;
 
 impl Policy<Req, Res, Error> for RetryErrors {
     type Future = future::Ready<Self>;
-    fn retry(&self, _: &Req, result: Result<&Res, &Error>) -> Option<Self::Future> {
+    fn retry(&self, _: &mut Req, result: &mut Result<Res, Error>) -> Option<Self::Future> {
         if result.is_err() {
             Some(future::ready(RetryErrors))
         } else {
@@ -120,7 +144,7 @@ struct Limit(usize);
 
 impl Policy<Req, Res, Error> for Limit {
     type Future = future::Ready<Self>;
-    fn retry(&self, _: &Req, result: Result<&Res, &Error>) -> Option<Self::Future> {
+    fn retry(&self, _: &mut Req, result: &mut Result<Res, Error>) -> Option<Self::Future> {
         if result.is_err() && self.0 > 0 {
             Some(future::ready(Limit(self.0 - 1)))
         } else {
@@ -138,8 +162,8 @@ struct UnlessErr(InnerError);
 
 impl Policy<Req, Res, Error> for UnlessErr {
     type Future = future::Ready<Self>;
-    fn retry(&self, _: &Req, result: Result<&Res, &Error>) -> Option<Self::Future> {
-        result.err().and_then(|err| {
+    fn retry(&self, _: &mut Req, result: &mut Result<Res, Error>) -> Option<Self::Future> {
+        result.as_ref().err().and_then(|err| {
             if err.to_string() != self.0 {
                 Some(future::ready(self.clone()))
             } else {
@@ -158,7 +182,7 @@ struct CannotClone;
 
 impl Policy<Req, Res, Error> for CannotClone {
     type Future = future::Ready<Self>;
-    fn retry(&self, _: &Req, _: Result<&Res, &Error>) -> Option<Self::Future> {
+    fn retry(&self, _: &mut Req, _: &mut Result<Res, Error>) -> Option<Self::Future> {
         unreachable!("retry cannot be called since request isn't cloned");
     }
 
@@ -167,6 +191,34 @@ impl Policy<Req, Res, Error> for CannotClone {
     }
 }
 
+#[derive(Clone)]
+struct MutatingPolicy {
+    remaining: usize,
+}
+
+impl Policy<Req, Res, Error> for MutatingPolicy
+where
+    Error: From<&'static str>,
+{
+    type Future = future::Ready<Self>;
+
+    fn retry(&self, req: &mut Req, result: &mut Result<Res, Error>) -> Option<Self::Future> {
+        if self.remaining == 0 {
+            *result = Err("out of retries".into());
+            None
+        } else {
+            *req = "retrying";
+            Some(future::ready(MutatingPolicy {
+                remaining: self.remaining - 1,
+            }))
+        }
+    }
+
+    fn clone_request(&self, req: &Req) -> Option<Req> {
+        Some(*req)
+    }
+}
+
 fn new_service<P: Policy<Req, Res, Error> + Clone>(
     policy: P,
 ) -> (mock::Spawn<tower::retry::Retry<P, Mock>>, Handle) {