perf: avoid relocking response stream references

Author: brentechols

src/client.rs

@@ -236,7 +236,8 @@ pub struct Connection<T, B: Buf = Bytes> {
 #[derive(Debug)]
 #[must_use = "futures do nothing unless polled"]
 pub struct ResponseFuture {
-    inner: proto::OpaqueStreamRef,
+    inner: Option<proto::OpaqueStreamRef>,
+    stream_id: crate::StreamId,
     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, 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(),
+                    stream_id: crate::StreamId::from_internal(response.stream_id()),
+                    inner: Some(response),
                     push_promise_consumed: false,
                 };
 
@@ -1469,21 +1471,21 @@ impl Future for ResponseFuture {
     type Output = Result<Response<RecvStream>, crate::Error>;
 
     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 inner = self
+            .inner
+            .as_mut()
+            .expect("ResponseFuture polled after completion");
+        let (parts, _) = ready!(inner.poll_response(cx))?.into_parts();
+        let body = RecvStream::new(FlowControl::new(self.inner.take().unwrap()));
 
         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())
+        self.stream_id
     }
 
     /// Polls for informational responses (1xx status codes).
@@ -1498,7 +1500,11 @@ impl ResponseFuture {
         &mut self,
         cx: &mut Context<'_>,
     ) -> Poll<Option<Result<Response<()>, crate::Error>>> {
-        self.inner.poll_informational(cx).map_err(Into::into)
+        self.inner
+            .as_mut()
+            .expect("ResponseFuture polled after completion")
+            .poll_informational(cx)
+            .map_err(Into::into)
     }
 
     /// Returns a stream of PushPromises
@@ -1513,7 +1519,7 @@ impl ResponseFuture {
         }
         self.push_promise_consumed = true;
         PushPromises {
-            inner: self.inner.clone(),
+            inner: self.inner.as_ref().unwrap().clone(),
         }
     }
 }
@@ -1535,7 +1541,8 @@ impl PushPromises {
             Poll::Ready(Some(Ok((request, response)))) => {
                 let response = PushedResponseFuture {
                     inner: ResponseFuture {
-                        inner: response,
+                        stream_id: crate::StreamId::from_internal(response.stream_id()),
+                        inner: Some(response),
                         push_promise_consumed: false,
                     },
                 };
@@ -1589,10 +1596,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, bool), SendError> {
         use super::stream::ContentLength;
         use http::Method;
 
@@ -344,14 +344,16 @@ where
 
         // TODO: ideally, OpaqueStreamRefs::new would do this, but we're holding
         // the lock, so it can't.
-        me.refs += 1;
+        me.refs += 2;
 
         let is_full = me.counts.next_send_stream_will_reach_capacity();
+        let response = OpaqueStreamRef::new(self.inner.clone(), &mut stream);
         Ok((
             StreamRef {
                 opaque: OpaqueStreamRef::new(self.inner.clone(), &mut stream),
                 send_buffer: self.send_buffer.clone(),
             },
+            response,
             is_full,
         ))
     }
@@ -1557,7 +1559,7 @@ impl OpaqueStreamRef {
     }
 
     pub fn stream_id(&self) -> StreamId {
-        self.inner.lock().unwrap().store[self.key].id
+        self.key.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");
     };