Make some functions less-generic to reduce binary bloat (#503)

* refactor: Extract FramedWrite::buffer to a less generic function

Should cut out another 23 KiB (since I see it duplicated)

* refactor: Extract some duplicated code to a function

* refactor: Extract part of flush into a less generic function

* refactor: Extract a less generic part of connection

* refactor: Factor out a less generic part of Connection::poll2

* refactor: Extract a non-generic part of handshake2

* refactor: Don't duplicate Streams code on Peer (-3.5%)

The `P: Peer` parameter is rarely used and there is already a mechanism
for using it dynamically.

* refactor: Make recv_frame less generic (-2.3%)

* Move out part of Connection::poll

* refactor: Extract parts of Connection

* refactor: Extract a non-generic part of reclaim_frame

* comments

Author: Marwes

src/client.rs

@@ -1124,6 +1124,20 @@ where
 
 // ===== impl Connection =====
 
+async fn bind_connection<T>(io: &mut T) -> Result<(), crate::Error>
+where
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    tracing::debug!("binding client connection");
+
+    let msg: &'static [u8] = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n";
+    io.write_all(msg).await.map_err(crate::Error::from_io)?;
+
+    tracing::debug!("client connection bound");
+
+    Ok(())
+}
+
 impl<T, B> Connection<T, B>
 where
     T: AsyncRead + AsyncWrite + Unpin,
@@ -1133,12 +1147,7 @@ where
         mut io: T,
         builder: Builder,
     ) -> Result<(SendRequest<B>, Connection<T, B>), crate::Error> {
-        tracing::debug!("binding client connection");
-
-        let msg: &'static [u8] = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n";
-        io.write_all(msg).await.map_err(crate::Error::from_io)?;
-
-        tracing::debug!("client connection bound");
+        bind_connection(&mut io).await?;
 
         // Create the codec
         let mut codec = Codec::new(io);

src/codec/framed_write.rs

@@ -23,6 +23,11 @@ pub struct FramedWrite<T, B> {
     /// Upstream `AsyncWrite`
     inner: T,
 
+    encoder: Encoder<B>,
+}
+
+#[derive(Debug)]
+struct Encoder<B> {
     /// HPACK encoder
     hpack: hpack::Encoder,
 
@@ -74,12 +79,14 @@ where
         let is_write_vectored = inner.is_write_vectored();
         FramedWrite {
             inner,
-            hpack: hpack::Encoder::default(),
-            buf: Cursor::new(BytesMut::with_capacity(DEFAULT_BUFFER_CAPACITY)),
-            next: None,
-            last_data_frame: None,
-            max_frame_size: frame::DEFAULT_MAX_FRAME_SIZE,
-            is_write_vectored,
+            encoder: Encoder {
+                hpack: hpack::Encoder::default(),
+                buf: Cursor::new(BytesMut::with_capacity(DEFAULT_BUFFER_CAPACITY)),
+                next: None,
+                last_data_frame: None,
+                max_frame_size: frame::DEFAULT_MAX_FRAME_SIZE,
+                is_write_vectored,
+            },
         }
     }
 
@@ -88,11 +95,11 @@ where
     /// Calling this function may result in the current contents of the buffer
     /// to be flushed to `T`.
     pub fn poll_ready(&mut self, cx: &mut Context) -> Poll<io::Result<()>> {
-        if !self.has_capacity() {
+        if !self.encoder.has_capacity() {
             // Try flushing
             ready!(self.flush(cx))?;
 
-            if !self.has_capacity() {
+            if !self.encoder.has_capacity() {
                 return Poll::Pending;
             }
         }
@@ -105,6 +112,128 @@ where
     /// `poll_ready` must be called first to ensure that a frame may be
     /// accepted.
     pub fn buffer(&mut self, item: Frame<B>) -> Result<(), UserError> {
+        self.encoder.buffer(item)
+    }
+
+    /// Flush buffered data to the wire
+    pub fn flush(&mut self, cx: &mut Context) -> Poll<io::Result<()>> {
+        let span = tracing::trace_span!("FramedWrite::flush");
+        let _e = span.enter();
+
+        loop {
+            while !self.encoder.is_empty() {
+                match self.encoder.next {
+                    Some(Next::Data(ref mut frame)) => {
+                        tracing::trace!(queued_data_frame = true);
+                        let mut buf = (&mut self.encoder.buf).chain(frame.payload_mut());
+                        ready!(write(
+                            &mut self.inner,
+                            self.encoder.is_write_vectored,
+                            &mut buf,
+                            cx,
+                        ))?
+                    }
+                    _ => {
+                        tracing::trace!(queued_data_frame = false);
+                        ready!(write(
+                            &mut self.inner,
+                            self.encoder.is_write_vectored,
+                            &mut self.encoder.buf,
+                            cx,
+                        ))?
+                    }
+                }
+            }
+
+            match self.encoder.unset_frame() {
+                ControlFlow::Continue => (),
+                ControlFlow::Break => break,
+            }
+        }
+
+        tracing::trace!("flushing buffer");
+        // Flush the upstream
+        ready!(Pin::new(&mut self.inner).poll_flush(cx))?;
+
+        Poll::Ready(Ok(()))
+    }
+
+    /// Close the codec
+    pub fn shutdown(&mut self, cx: &mut Context) -> Poll<io::Result<()>> {
+        ready!(self.flush(cx))?;
+        Pin::new(&mut self.inner).poll_shutdown(cx)
+    }
+}
+
+fn write<T, B>(
+    writer: &mut T,
+    is_write_vectored: bool,
+    buf: &mut B,
+    cx: &mut Context<'_>,
+) -> Poll<io::Result<()>>
+where
+    T: AsyncWrite + Unpin,
+    B: Buf,
+{
+    // TODO(eliza): when tokio-util 0.5.1 is released, this
+    // could just use `poll_write_buf`...
+    const MAX_IOVS: usize = 64;
+    let n = if is_write_vectored {
+        let mut bufs = [IoSlice::new(&[]); MAX_IOVS];
+        let cnt = buf.chunks_vectored(&mut bufs);
+        ready!(Pin::new(writer).poll_write_vectored(cx, &bufs[..cnt]))?
+    } else {
+        ready!(Pin::new(writer).poll_write(cx, buf.chunk()))?
+    };
+    buf.advance(n);
+    Ok(()).into()
+}
+
+#[must_use]
+enum ControlFlow {
+    Continue,
+    Break,
+}
+
+impl<B> Encoder<B>
+where
+    B: Buf,
+{
+    fn unset_frame(&mut self) -> ControlFlow {
+        // Clear internal buffer
+        self.buf.set_position(0);
+        self.buf.get_mut().clear();
+
+        // The data frame has been written, so unset it
+        match self.next.take() {
+            Some(Next::Data(frame)) => {
+                self.last_data_frame = Some(frame);
+                debug_assert!(self.is_empty());
+                ControlFlow::Break
+            }
+            Some(Next::Continuation(frame)) => {
+                // Buffer the continuation frame, then try to write again
+                let mut buf = limited_write_buf!(self);
+                if let Some(continuation) = frame.encode(&mut self.hpack, &mut buf) {
+                    // We previously had a CONTINUATION, and after encoding
+                    // it, we got *another* one? Let's just double check
+                    // that at least some progress is being made...
+                    if self.buf.get_ref().len() == frame::HEADER_LEN {
+                        // If *only* the CONTINUATION frame header was
+                        // written, and *no* header fields, we're stuck
+                        // in a loop...
+                        panic!("CONTINUATION frame write loop; header value too big to encode");
+                    }
+
+                    self.next = Some(Next::Continuation(continuation));
+                }
+                ControlFlow::Continue
+            }
+            None => ControlFlow::Break,
+        }
+    }
+
+    fn buffer(&mut self, item: Frame<B>) -> Result<(), UserError> {
         // Ensure that we have enough capacity to accept the write.
         assert!(self.has_capacity());
         let span = tracing::trace_span!("FramedWrite::buffer", frame = ?item);
@@ -185,93 +314,6 @@ where
         Ok(())
     }
 
-    /// Flush buffered data to the wire
-    pub fn flush(&mut self, cx: &mut Context) -> Poll<io::Result<()>> {
-        const MAX_IOVS: usize = 64;
-
-        let span = tracing::trace_span!("FramedWrite::flush");
-        let _e = span.enter();
-
-        loop {
-            while !self.is_empty() {
-                match self.next {
-                    Some(Next::Data(ref mut frame)) => {
-                        tracing::trace!(queued_data_frame = true);
-                        let mut buf = (&mut self.buf).chain(frame.payload_mut());
-                        // TODO(eliza): when tokio-util 0.5.1 is released, this
-                        // could just use `poll_write_buf`...
-                        let n = if self.is_write_vectored {
-                            let mut bufs = [IoSlice::new(&[]); MAX_IOVS];
-                            let cnt = buf.chunks_vectored(&mut bufs);
-                            ready!(Pin::new(&mut self.inner).poll_write_vectored(cx, &bufs[..cnt]))?
-                        } else {
-                            ready!(Pin::new(&mut self.inner).poll_write(cx, buf.chunk()))?
-                        };
-                        buf.advance(n);
-                    }
-                    _ => {
-                        tracing::trace!(queued_data_frame = false);
-                        let n = if self.is_write_vectored {
-                            let mut iovs = [IoSlice::new(&[]); MAX_IOVS];
-                            let cnt = self.buf.chunks_vectored(&mut iovs);
-                            ready!(
-                                Pin::new(&mut self.inner).poll_write_vectored(cx, &mut iovs[..cnt])
-                            )?
-                        } else {
-                            ready!(Pin::new(&mut self.inner).poll_write(cx, &mut self.buf.chunk()))?
-                        };
-                        self.buf.advance(n);
-                    }
-                }
-            }
-
-            // Clear internal buffer
-            self.buf.set_position(0);
-            self.buf.get_mut().clear();
-
-            // The data frame has been written, so unset it
-            match self.next.take() {
-                Some(Next::Data(frame)) => {
-                    self.last_data_frame = Some(frame);
-                    debug_assert!(self.is_empty());
-                    break;
-                }
-                Some(Next::Continuation(frame)) => {
-                    // Buffer the continuation frame, then try to write again
-                    let mut buf = limited_write_buf!(self);
-                    if let Some(continuation) = frame.encode(&mut self.hpack, &mut buf) {
-                        // We previously had a CONTINUATION, and after encoding
-                        // it, we got *another* one? Let's just double check
-                        // that at least some progress is being made...
-                        if self.buf.get_ref().len() == frame::HEADER_LEN {
-                            // If *only* the CONTINUATION frame header was
-                            // written, and *no* header fields, we're stuck
-                            // in a loop...
-                            panic!("CONTINUATION frame write loop; header value too big to encode");
-                        }
-
-                        self.next = Some(Next::Continuation(continuation));
-                    }
-                }
-                None => {
-                    break;
-                }
-            }
-        }
-
-        tracing::trace!("flushing buffer");
-        // Flush the upstream
-        ready!(Pin::new(&mut self.inner).poll_flush(cx))?;
-
-        Poll::Ready(Ok(()))
-    }
-
-    /// Close the codec
-    pub fn shutdown(&mut self, cx: &mut Context) -> Poll<io::Result<()>> {
-        ready!(self.flush(cx))?;
-        Pin::new(&mut self.inner).poll_shutdown(cx)
-    }
-
     fn has_capacity(&self) -> bool {
         self.next.is_none() && self.buf.get_ref().remaining_mut() >= MIN_BUFFER_CAPACITY
     }
@@ -284,26 +326,32 @@ where
     }
 }
 
+impl<B> Encoder<B> {
+    fn max_frame_size(&self) -> usize {
+        self.max_frame_size as usize
+    }
+}
+
 impl<T, B> FramedWrite<T, B> {
     /// Returns the max frame size that can be sent
     pub fn max_frame_size(&self) -> usize {
-        self.max_frame_size as usize
+        self.encoder.max_frame_size()
     }
 
     /// Set the peer's max frame size.
     pub fn set_max_frame_size(&mut self, val: usize) {
         assert!(val <= frame::MAX_MAX_FRAME_SIZE as usize);
-        self.max_frame_size = val as FrameSize;
+        self.encoder.max_frame_size = val as FrameSize;
     }
 
     /// Set the peer's header table size.
     pub fn set_header_table_size(&mut self, val: usize) {
-        self.hpack.update_max_size(val);
+        self.encoder.hpack.update_max_size(val);
     }
 
     /// Retrieve the last data frame that has been sent
     pub fn take_last_data_frame(&mut self) -> Option<frame::Data<B>> {
-        self.last_data_frame.take()
+        self.encoder.last_data_frame.take()
     }
 
     pub fn get_mut(&mut self) -> &mut T {