diff --git a/actix-service/src/and_then.rs b/actix-service/src/and_then.rs index 1c7d52b4..690d0cb4 100644 --- a/actix-service/src/and_then.rs +++ b/actix-service/src/and_then.rs @@ -51,7 +51,7 @@ where fn call(&mut self, req: A::Request) -> Self::Future { AndThenServiceResponse { - state: State::A(self.0.get_mut().0.call(req), self.0.clone()), + state: State::A(self.0.get_mut().0.call(req), Some(self.0.clone())), } } } @@ -72,8 +72,9 @@ where A: Service, B: Service, { - A(#[pin] A::Future, Cell<(A, B)>), + A(#[pin] A::Future, Option>), B(#[pin] B::Future), + Empty, } impl Future for AndThenServiceResponse @@ -91,13 +92,19 @@ where match this.state.as_mut().project() { State::A(fut, b) => match fut.poll(cx)? { Poll::Ready(res) => { + let mut b = b.take().unwrap(); + this.state.set(State::Empty); // drop fut A let fut = b.get_mut().1.call(res); this.state.set(State::B(fut)); self.poll(cx) } Poll::Pending => Poll::Pending, }, - State::B(fut) => fut.poll(cx), + State::B(fut) => fut.poll(cx).map(|r| { + this.state.set(State::Empty); + r + }), + State::Empty => panic!("future must not be polled after it returned `Poll::Ready`"), } } } diff --git a/actix-service/src/and_then_apply_fn.rs b/actix-service/src/and_then_apply_fn.rs index 3ae7e59f..2ac7474d 100644 --- a/actix-service/src/and_then_apply_fn.rs +++ b/actix-service/src/and_then_apply_fn.rs @@ -79,7 +79,7 @@ where fn call(&mut self, req: A::Request) -> Self::Future { AndThenApplyFnFuture { - state: State::A(self.a.call(req), self.b.clone()), + state: State::A(self.a.call(req), Some(self.b.clone())), } } } @@ -108,8 +108,9 @@ where Err: From, Err: From, { - A(#[pin] A::Future, Cell<(B, F)>), + A(#[pin] A::Future, Option>), B(#[pin] Fut), + Empty, } impl Future for AndThenApplyFnFuture @@ -130,6 +131,8 @@ where match this.state.as_mut().project() { State::A(fut, b) => match fut.poll(cx)? { Poll::Ready(res) => { + let mut b = b.take().unwrap(); + this.state.set(State::Empty); let b = b.get_mut(); let fut = (&mut b.1)(res, &mut b.0); this.state.set(State::B(fut)); @@ -137,7 +140,11 @@ where } Poll::Pending => Poll::Pending, }, - State::B(fut) => fut.poll(cx), + State::B(fut) => fut.poll(cx).map(|r| { + this.state.set(State::Empty); + r + }), + State::Empty => panic!("future must not be polled after it returned `Poll::Ready`"), } } } diff --git a/actix-service/src/then.rs b/actix-service/src/then.rs index 9741ddfa..eaf2da5e 100644 --- a/actix-service/src/then.rs +++ b/actix-service/src/then.rs @@ -58,7 +58,7 @@ where fn call(&mut self, req: A::Request) -> Self::Future { ThenServiceResponse { - state: ThenServiceResponseState::A(self.a.call(req), self.b.clone()), + state: State::A(self.a.call(req), Some(self.b.clone())), } } } @@ -70,17 +70,18 @@ where B: Service>, { #[pin] - state: ThenServiceResponseState, + state: State, } #[pin_project::pin_project] -enum ThenServiceResponseState +enum State where A: Service, B: Service>, { - A(#[pin] A::Future, Cell), + A(#[pin] A::Future, Option>), B(#[pin] B::Future), + Empty } impl Future for ThenServiceResponse @@ -96,15 +97,21 @@ where #[project] match this.state.as_mut().project() { - ThenServiceResponseState::A(fut, b) => match fut.poll(cx) { + State::A(fut, b) => match fut.poll(cx) { Poll::Ready(res) => { + let mut b = b.take().unwrap(); + this.state.set(State::Empty); // drop fut A let fut = b.get_mut().call(res); - this.state.set(ThenServiceResponseState::B(fut)); + this.state.set(State::B(fut)); self.poll(cx) } Poll::Pending => Poll::Pending, }, - ThenServiceResponseState::B(fut) => fut.poll(cx), + State::B(fut) => fut.poll(cx).map(|r| { + this.state.set(State::Empty); + r + }), + State::Empty => panic!("future must not be polled after it returned `Poll::Ready`") } } }