From 03773e48f2057f1d71c5952d69269382a7cf3afb Mon Sep 17 00:00:00 2001 From: Brent Echols Date: Sat, 15 Aug 2026 14:17:15 -0700 Subject: [PATCH] perf: avoid relocking response stream references --- src/client.rs | 23 ++++++++-------- src/proto/streams/store.rs | 6 +++++ src/proto/streams/streams.rs | 21 +++++++++------ src/server.rs | 8 ------ src/share.rs | 8 ------ tests/h2-tests/tests/client_request.rs | 10 +++---- .../h2-tests/tests/informational_responses.rs | 9 ++++++- tests/h2-tests/tests/push_promise.rs | 27 ++++++++++--------- 8 files changed, 58 insertions(+), 54 deletions(-) diff --git a/src/client.rs b/src/client.rs index f5ce1406..aca107a7 100644 --- a/src/client.rs +++ b/src/client.rs @@ -237,6 +237,7 @@ pub struct Connection { #[must_use = "futures do nothing unless polled"] pub struct ResponseFuture { inner: proto::OpaqueStreamRef, + body: Option, push_promise_consumed: bool, } @@ -517,7 +518,7 @@ where self.inner .send_request(request, end_of_stream, self.pending.as_ref()) .map_err(Into::into) - .map(|(stream, is_full)| { + .map(|(stream, response, body, is_full)| { if stream.is_pending_open() && is_full { // Only prevent sending another request when the request queue // is not full. @@ -525,7 +526,8 @@ where } let response = ResponseFuture { - inner: stream.clone_to_opaque(), + inner: response, + body: Some(body), push_promise_consumed: false, }; @@ -1470,7 +1472,11 @@ impl Future for ResponseFuture { fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let (parts, _) = ready!(self.inner.poll_response(cx))?.into_parts(); - let body = RecvStream::new(FlowControl::new(self.inner.clone())); + let body = RecvStream::new(FlowControl::new( + self.body + .take() + .expect("ResponseFuture polled after completion"), + )); Poll::Ready(Ok(Response::from_parts(parts, body))) } @@ -1478,10 +1484,6 @@ impl Future for ResponseFuture { impl ResponseFuture { /// Returns the stream ID of the response stream. - /// - /// # Panics - /// - /// If the lock on the stream store has been poisoned. pub fn stream_id(&self) -> crate::StreamId { crate::StreamId::from_internal(self.inner.stream_id()) } @@ -1532,10 +1534,11 @@ impl PushPromises { cx: &mut Context<'_>, ) -> Poll>> { match self.inner.poll_pushed(cx) { - Poll::Ready(Some(Ok((request, response)))) => { + Poll::Ready(Some(Ok((request, response, body)))) => { let response = PushedResponseFuture { inner: ResponseFuture { inner: response, + body: Some(body), push_promise_consumed: false, }, }; @@ -1589,10 +1592,6 @@ impl Future for PushedResponseFuture { impl PushedResponseFuture { /// Returns the stream ID of the response stream. - /// - /// # Panics - /// - /// If the lock on the stream store has been poisoned. pub fn stream_id(&self) -> crate::StreamId { self.inner.stream_id() } diff --git a/src/proto/streams/store.rs b/src/proto/streams/store.rs index b2685ed3..0bced776 100644 --- a/src/proto/streams/store.rs +++ b/src/proto/streams/store.rs @@ -29,6 +29,12 @@ pub(crate) struct Key { stream_id: StreamId, } +impl Key { + pub(crate) fn stream_id(self) -> StreamId { + self.stream_id + } +} + // We can never have more than `StreamId::MAX` streams in the store, // so we can save a smaller index (u32 vs usize). #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/src/proto/streams/streams.rs b/src/proto/streams/streams.rs index ba326d0d..1a781402 100644 --- a/src/proto/streams/streams.rs +++ b/src/proto/streams/streams.rs @@ -264,7 +264,7 @@ where mut request: Request<()>, end_of_stream: bool, pending: Option<&OpaqueStreamRef>, - ) -> Result<(StreamRef, bool), SendError> { + ) -> Result<(StreamRef, OpaqueStreamRef, OpaqueStreamRef, bool), SendError> { use super::stream::ContentLength; use http::Method; @@ -344,14 +344,18 @@ where // TODO: ideally, OpaqueStreamRefs::new would do this, but we're holding // the lock, so it can't. - me.refs += 1; + me.refs += 3; let is_full = me.counts.next_send_stream_will_reach_capacity(); + let response = OpaqueStreamRef::new(self.inner.clone(), &mut stream); + let body = OpaqueStreamRef::new(self.inner.clone(), &mut stream); Ok(( StreamRef { opaque: OpaqueStreamRef::new(self.inner.clone(), &mut stream), send_buffer: self.send_buffer.clone(), }, + response, + body, is_full, )) } @@ -1472,7 +1476,7 @@ impl OpaqueStreamRef { pub fn poll_pushed( &mut self, cx: &Context, - ) -> Poll, OpaqueStreamRef), proto::Error>>> { + ) -> Poll, OpaqueStreamRef, OpaqueStreamRef), proto::Error>>> { let mut me = self.inner.lock().unwrap(); let me = &mut *me; @@ -1481,10 +1485,11 @@ impl OpaqueStreamRef { .recv .poll_pushed(cx, &mut stream) .map_ok(|(h, key)| { - me.refs += 1; - let opaque_ref = - OpaqueStreamRef::new(self.inner.clone(), &mut me.store.resolve(key)); - (h, opaque_ref) + me.refs += 2; + let stream = &mut me.store.resolve(key); + let response = OpaqueStreamRef::new(self.inner.clone(), stream); + let body = OpaqueStreamRef::new(self.inner.clone(), stream); + (h, response, body) }) } @@ -1557,7 +1562,7 @@ impl OpaqueStreamRef { } pub fn stream_id(&self) -> StreamId { - self.inner.lock().unwrap().store[self.key].id + self.key.stream_id() } } diff --git a/src/server.rs b/src/server.rs index da6f259b..1b0a1a75 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1297,10 +1297,6 @@ impl SendResponse { } /// Returns the stream ID of the response stream. - /// - /// # Panics - /// - /// If the lock on the stream store has been poisoned. pub fn stream_id(&self) -> crate::StreamId { crate::StreamId::from_internal(self.inner.stream_id()) } @@ -1369,10 +1365,6 @@ impl SendPushedResponse { } /// Returns the stream ID of the response stream. - /// - /// # Panics - /// - /// If the lock on the stream store has been poisoned. pub fn stream_id(&self) -> crate::StreamId { self.inner.stream_id() } diff --git a/src/share.rs b/src/share.rs index 10198734..7c1894d3 100644 --- a/src/share.rs +++ b/src/share.rs @@ -373,10 +373,6 @@ impl SendStream { } /// Returns the stream ID of this `SendStream`. - /// - /// # Panics - /// - /// If the lock on the stream store has been poisoned. pub fn stream_id(&self) -> StreamId { StreamId::from_internal(self.inner.stream_id()) } @@ -451,10 +447,6 @@ impl RecvStream { } /// Returns the stream ID of this stream. - /// - /// # Panics - /// - /// If the lock on the stream store has been poisoned. pub fn stream_id(&self) -> StreamId { self.inner.stream_id() } diff --git a/tests/h2-tests/tests/client_request.rs b/tests/h2-tests/tests/client_request.rs index 04f8ddda..4600514f 100644 --- a/tests/h2-tests/tests/client_request.rs +++ b/tests/h2-tests/tests/client_request.rs @@ -1,4 +1,4 @@ -use futures::future::{ready, Either}; +use futures::future::{poll_fn, ready, Either}; use futures::stream::FuturesUnordered; use futures::StreamExt; use h2_support::prelude::*; @@ -47,12 +47,12 @@ async fn client_other_thread() { .uri("https://http2.akamai.com/") .body(()) .unwrap(); - let _res = client - .send_request(request, true) - .unwrap() - .0 + let mut response = client.send_request(request, true).unwrap().0; + let stream_id = response.stream_id(); + let _res = poll_fn(|cx| Pin::new(&mut response).poll(cx)) .await .expect("request"); + assert_eq!(response.stream_id(), stream_id); }); h2.await.expect("h2"); }; diff --git a/tests/h2-tests/tests/informational_responses.rs b/tests/h2-tests/tests/informational_responses.rs index 53363239..50b847d9 100644 --- a/tests/h2-tests/tests/informational_responses.rs +++ b/tests/h2-tests/tests/informational_responses.rs @@ -3,6 +3,7 @@ use futures::{future::poll_fn, StreamExt}; use h2_support::prelude::*; use http::{Response, StatusCode}; +use std::{pin::Pin, task::Poll}; #[tokio::test] async fn send_100_continue() { @@ -297,8 +298,14 @@ async fn client_poll_informational_responses_none() { sync_sender.send(()).unwrap(); // Get the final response - let response = response_future.await.expect("response error"); + let response = poll_fn(|cx| Pin::new(&mut response_future).poll(cx)) + .await + .expect("response error"); assert_eq!(response.status(), StatusCode::OK); + assert!(matches!( + poll_fn(|cx| Poll::Ready(response_future.poll_informational(cx))).await, + Poll::Pending + )); let (_hdr, mut recv_stream) = response.into_parts(); let data = recv_stream.data().await.unwrap().unwrap(); assert_eq!("request body", data); diff --git a/tests/h2-tests/tests/push_promise.rs b/tests/h2-tests/tests/push_promise.rs index 59596ebf..a304ee9e 100644 --- a/tests/h2-tests/tests/push_promise.rs +++ b/tests/h2-tests/tests/push_promise.rs @@ -1,5 +1,6 @@ -use futures::{StreamExt, TryStreamExt}; +use futures::{future::poll_fn, StreamExt, TryStreamExt}; use h2_support::prelude::*; +use std::pin::Pin; #[tokio::test] async fn recv_push_works() { @@ -32,27 +33,29 @@ async fn recv_push_works() { .body(()) .unwrap(); let (mut resp, _) = client.send_request(request, true).unwrap(); - let pushed = resp.push_promises(); - let check_resp_status = async move { - let resp = resp.await.unwrap(); - assert_eq!(resp.status(), StatusCode::NOT_FOUND); - }; - let check_pushed_response = async move { + let check_responses = async move { + let response = poll_fn(|cx| Pin::new(&mut resp).poll(cx)).await.unwrap(); + assert_eq!(response.status(), StatusCode::NOT_FOUND); + + let pushed = resp.push_promises(); let p = pushed.and_then(|headers| async move { - let (request, response) = headers.into_parts(); + let (request, mut response) = headers.into_parts(); assert_eq!(request.into_parts().0.method, Method::GET); - let resp = response.await.unwrap(); + let stream_id = response.stream_id(); + let resp = poll_fn(|cx| Pin::new(&mut response).poll(cx)) + .await + .unwrap(); + assert_eq!(response.stream_id(), stream_id); assert_eq!(resp.status(), StatusCode::OK); let b = util::concat(resp.into_body()).await.unwrap(); assert_eq!(b, "promised_data"); Ok(()) }); let ps: Vec<_> = p.collect().await; - assert_eq!(1, ps.len()) + assert_eq!(1, ps.len()); }; - h2.drive(join(check_resp_status, check_pushed_response)) - .await; + h2.drive(check_responses).await; }; join(mock, h2).await;