Retry middleware improvements for non-cloneable requests

tower/src/retry/future.rs

@@ -16,7 +16,7 @@ pin_project! {
         P: Policy<Request, S::Response, S::Error>,
         S: Service<Request>,
     {
-        request: Option<Request>,
+        request: P::CloneableRequest,
         #[pin]
         retry: Retry<P, S>,
         #[pin]
@@ -49,7 +49,7 @@ where
     S: Service<Request>,
 {
     pub(crate) fn new(
-        request: Option<Request>,
+        request: P::CloneableRequest,
         retry: Retry<P, S>,
         future: S::Future,
     ) -> ResponseFuture<P, S, Request> {
@@ -75,16 +75,11 @@ where
             match this.state.as_mut().project() {
                 StateProj::Called { future } => {
                     let mut result = ready!(future.poll(cx));
-                    if let Some(req) = &mut this.request {
-                        match this.retry.policy.retry(req, &mut result) {
-                            Some(waiting) => {
-                                this.state.set(State::Waiting { waiting });
-                            }
-                            None => return Poll::Ready(result),
+                    match this.retry.policy.retry(this.request, &mut result) {
+                        Some(waiting) => {
+                            this.state.set(State::Waiting { waiting });
                         }
-                    } else {
-                        // request wasn't cloned, so no way to retry it
-                        return Poll::Ready(result);
+                        None => return Poll::Ready(result),
                     }
                 }
                 StateProj::Waiting { waiting } => {
@@ -105,11 +100,7 @@ where
                     // in Ready to make it Unpin so that we can get &mut Ready as needed to call
                     // poll_ready on it.
                     ready!(this.retry.as_mut().project().service.poll_ready(cx))?;
-                    let req = this
-                        .request
-                        .take()
-                        .expect("retrying requires cloned request");
-                    *this.request = this.retry.policy.clone_request(&req);
+                    let req = this.retry.policy.clone_request(this.request);
                     this.state.set(State::Called {
                         future: this.retry.as_mut().project().service.call(req),
                     });

tower/src/retry/mod.rs

@@ -86,9 +86,10 @@ where
     }
 
     fn call(&mut self, request: Request) -> Self::Future {
-        let cloned = self.policy.clone_request(&request);
-        let future = self.service.call(request);
+        let cloneable = self.policy.create_cloneable_request(request);
+        let req = self.policy.clone_request(&cloneable);
+        let future = self.service.call(req);
 
-        ResponseFuture::new(cloned, self.clone(), future)
+        ResponseFuture::new(cloneable, self.clone(), future)
     }
 }

tower/src/retry/policy.rs

@@ -47,6 +47,10 @@ pub trait Policy<Req, Res, E> {
     /// The [`Future`] type returned by [`Policy::retry`].
     type Future: Future<Output = ()>;
 
+    /// A type that is able to store request object, that can be
+    /// cloned back to original request.
+    type CloneableRequest;
+
     /// Check the policy if a certain request should be retried.
     ///
     /// This method is passed a reference to the original request, and either
@@ -80,15 +84,25 @@ pub trait Policy<Req, Res, E> {
     ///
     /// [`Service::Response`]: crate::Service::Response
     /// [`Service::Error`]: crate::Service::Error
-    fn retry(&mut self, req: &mut Req, result: &mut Result<Res, E>) -> Option<Self::Future>;
+    fn retry(
+        &mut self,
+        req: &mut Self::CloneableRequest,
+        result: &mut Result<Res, E>,
+    ) -> Option<Self::Future>;
 
-    /// Tries to clone a request before being passed to the inner service.
-    ///
-    /// If the request cannot be cloned, return [`None`]. Moreover, the retry
-    /// function will not be called if the [`None`] is returned.
-    fn clone_request(&mut self, req: &Req) -> Option<Req>;
+    /// Consume initial request and returns `CloneableRequest` which
+    /// will be used to recreate original request objects for each retry.
+    /// This is essential in cases where original request cannot be cloned,
+    /// but can only be consumed.
+    fn create_cloneable_request(&mut self, req: Req) -> Self::CloneableRequest;
+
+    /// Recreates original request object for each retry.
+    fn clone_request(&mut self, req: &Self::CloneableRequest) -> Req;
 }
 
 // Ensure `Policy` is object safe
 #[cfg(test)]
-fn _obj_safe(_: Box<dyn Policy<(), (), (), Future = futures::future::Ready<()>>>) {}
+fn _obj_safe(
+    _: Box<dyn Policy<(), (), (), Future = futures::future::Ready<()>, CloneableRequest = ()>>,
+) {
+}