diff --git a/quinn-h3/src/client.rs b/quinn-h3/src/client.rs index 93a692b52..54705f20a 100644 --- a/quinn-h3/src/client.rs +++ b/quinn-h3/src/client.rs @@ -91,6 +91,7 @@ use crate::{ proto::{headers::Header, settings::Settings, ErrorCode}, Error, SendData, ZeroRttAccepted, }; +use futures_util::future; /// Configure and build a new HTTP/3 client /// @@ -663,19 +664,39 @@ impl RecvResponse { } } - /// Cancel a HTTP/3 response reception + /// Cancel an HTTP/3 response reception /// /// Server will receive a request error with `REQUEST_CANCELLED` code. Any call on any /// object related with this request will fail. - pub fn cancel(&mut self) { - if let Some(recv) = self.recv.as_mut() { - self.conn - .h3 - .lock() - .unwrap() - .cancel_request(self.stream_id.unwrap()); - recv.reset(ErrorCode::REQUEST_CANCELLED); + pub async fn cancel(&mut self) { + let stream_id = match self.state { + RecvResponseState::Finished => None, + RecvResponseState::Opening(ref mut o) => { + future::poll_fn(|cx| { + Poll::Ready(match o.poll_unpin(cx) { + Poll::Ready(Ok((mut r, i))) => { + let _ = r.stop(ErrorCode::REQUEST_CANCELLED.into()); + Some(i) + } + _ => None, + }) + }) + .await + } + RecvResponseState::Receiving => { + self.recv + .take() + .unwrap() + .reset(ErrorCode::REQUEST_CANCELLED); + self.stream_id.take() + } + }; + + if let Some(id) = stream_id { + self.conn.h3.lock().unwrap().cancel_request(id); } + + self.state = RecvResponseState::Finished; } } diff --git a/quinn-h3/src/tests/mod.rs b/quinn-h3/src/tests/mod.rs index 2ca948fe1..8370878a2 100644 --- a/quinn-h3/src/tests/mod.rs +++ b/quinn-h3/src/tests/mod.rs @@ -119,13 +119,12 @@ async fn client_cancel_response() { .await .expect("accept"); let recv_req = incoming_req.next().await.expect("wait request"); - delay_for(Duration::from_millis(25)).await; let (_, mut sender) = recv_req.await.expect("recv_req"); sender .send_response( Response::builder() .status(StatusCode::OK) - .body(Body::from(())) + .body(Body::from("a".repeat(1024 * 1024 * 100).as_ref())) .unwrap(), ) .await @@ -133,10 +132,9 @@ async fn client_cancel_response() { }); let conn = helper.make_connection().await; - delay_for(Duration::from_millis(50)).await; let (req, mut resp) = conn.send_request(get("/")); req.await.unwrap(); - resp.cancel(); + resp.cancel().await; assert_matches!( timeout_join(server_handle).await,