perf: avoid relocking response stream references

Author: brentechols

src/client.rs

@@ -237,6 +237,7 @@ pub struct Connection<T, B: Buf = Bytes> {
 #[must_use = "futures do nothing unless polled"]
 pub struct ResponseFuture {
     inner: proto::OpaqueStreamRef,
+    body: Option<proto::OpaqueStreamRef>,
     push_promise_consumed: bool,
 }
 
@@ -517,15 +518,16 @@ 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.
                     self.pending = Some(stream.clone_to_opaque());
                 }
 
                 let response = ResponseFuture {
-                    inner: stream.clone_to_opaque(),
+                    inner: response,
+                    body: Some(body),
                     push_promise_consumed: false,
                 };
 
@@ -1470,18 +1472,18 @@ impl Future for ResponseFuture {
 
     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         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)))
     }
 }
 
 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<Option<Result<PushPromise, crate::Error>>> {
         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()
     }

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)]

src/proto/streams/streams.rs

@@ -264,7 +264,7 @@ where
         mut request: Request<()>,
         end_of_stream: bool,
         pending: Option<&OpaqueStreamRef>,
-    ) -> Result<(StreamRef<B>, bool), SendError> {
+    ) -> Result<(StreamRef<B>, 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<Option<Result<(Request<()>, OpaqueStreamRef), proto::Error>>> {
+    ) -> Poll<Option<Result<(Request<()>, 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()
     }
 }
 

src/server.rs

@@ -1297,10 +1297,6 @@ impl<B: Buf> SendResponse<B> {
     }
 
     /// 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<B: Buf> SendPushedResponse<B> {
     }
 
     /// 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()
     }

src/share.rs

@@ -373,10 +373,6 @@ impl<B: Buf> SendStream<B> {
     }
 
     /// 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()
     }

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");
     };

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);

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;