diff --git a/tower/src/retry/future.rs b/tower/src/retry/future.rs index 7e2218440..ff8fb6eb4 100644 --- a/tower/src/retry/future.rs +++ b/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), }); diff --git a/tower/src/retry/mod.rs b/tower/src/retry/mod.rs index 1bb5e29ed..3f38a66f9 100644 --- a/tower/src/retry/mod.rs +++ b/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) } } diff --git a/tower/src/retry/policy.rs b/tower/src/retry/policy.rs index 57b80a71c..44dd178e4 100644 --- a/tower/src/retry/policy.rs +++ b/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 = ()>>, +) { +}