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

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");
+        bind_connection(&mut io).await?;
-
-        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");
 
         // 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(),
+            encoder: Encoder {
-            buf: Cursor::new(BytesMut::with_capacity(DEFAULT_BUFFER_CAPACITY)),
+                hpack: hpack::Encoder::default(),
-            next: None,
+                buf: Cursor::new(BytesMut::with_capacity(DEFAULT_BUFFER_CAPACITY)),
-            last_data_frame: None,
+                next: None,
-            max_frame_size: frame::DEFAULT_MAX_FRAME_SIZE,
+                last_data_frame: None,
-            is_write_vectored,
+                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 {

src/proto/connection.rs

@@ -17,6 +17,19 @@ use tokio::io::{AsyncRead, AsyncWrite};
 /// An H2 connection
 #[derive(Debug)]
 pub(crate) struct Connection<T, P, B: Buf = Bytes>
+where
+    P: Peer,
+{
+    /// Read / write frame values
+    codec: Codec<T, Prioritized<B>>,
+
+    inner: ConnectionInner<P, B>,
+}
+
+// Extracted part of `Connection` which does not depend on `T`. Reduces the amount of duplicated
+// method instantiations.
+#[derive(Debug)]
+struct ConnectionInner<P, B: Buf = Bytes>
 where
     P: Peer,
 {
@@ -29,9 +42,6 @@ where
     /// graceful shutdown.
     error: Option<Reason>,
 
-    /// Read / write frame values
-    codec: Codec<T, Prioritized<B>>,
-
     /// Pending GOAWAY frames to write.
     go_away: GoAway,
 
@@ -51,6 +61,18 @@ where
     _phantom: PhantomData<P>,
 }
 
+struct DynConnection<'a, B: Buf = Bytes> {
+    state: &'a mut State,
+
+    go_away: &'a mut GoAway,
+
+    streams: DynStreams<'a, B>,
+
+    error: &'a mut Option<Reason>,
+
+    ping_pong: &'a mut PingPong,
+}
+
 #[derive(Debug, Clone)]
 pub(crate) struct Config {
     pub next_stream_id: StreamId,
@@ -79,51 +101,56 @@ where
     B: Buf,
 {
     pub fn new(codec: Codec<T, Prioritized<B>>, config: Config) -> Connection<T, P, B> {
-        let streams = Streams::new(streams::Config {
+        fn streams_config(config: &Config) -> streams::Config {
-            local_init_window_sz: config
+            streams::Config {
-                .settings
+                local_init_window_sz: config
-                .initial_window_size()
+                    .settings
-                .unwrap_or(DEFAULT_INITIAL_WINDOW_SIZE),
+                    .initial_window_size()
-            initial_max_send_streams: config.initial_max_send_streams,
+                    .unwrap_or(DEFAULT_INITIAL_WINDOW_SIZE),
-            local_next_stream_id: config.next_stream_id,
+                initial_max_send_streams: config.initial_max_send_streams,
-            local_push_enabled: config.settings.is_push_enabled().unwrap_or(true),
+                local_next_stream_id: config.next_stream_id,
-            local_reset_duration: config.reset_stream_duration,
+                local_push_enabled: config.settings.is_push_enabled().unwrap_or(true),
-            local_reset_max: config.reset_stream_max,
+                local_reset_duration: config.reset_stream_duration,
-            remote_init_window_sz: DEFAULT_INITIAL_WINDOW_SIZE,
+                local_reset_max: config.reset_stream_max,
-            remote_max_initiated: config
+                remote_init_window_sz: DEFAULT_INITIAL_WINDOW_SIZE,
-                .settings
+                remote_max_initiated: config
-                .max_concurrent_streams()
+                    .settings
-                .map(|max| max as usize),
+                    .max_concurrent_streams()
-        });
+                    .map(|max| max as usize),
+            }
+        }
+        let streams = Streams::new(streams_config(&config));
         Connection {
-            state: State::Open,
-            error: None,
             codec,
-            go_away: GoAway::new(),
+            inner: ConnectionInner {
-            ping_pong: PingPong::new(),
+                state: State::Open,
-            settings: Settings::new(config.settings),
+                error: None,
-            streams,
+                go_away: GoAway::new(),
-            span: tracing::debug_span!("Connection", peer = %P::NAME),
+                ping_pong: PingPong::new(),
-            _phantom: PhantomData,
+                settings: Settings::new(config.settings),
+                streams,
+                span: tracing::debug_span!("Connection", peer = %P::NAME),
+                _phantom: PhantomData,
+            },
         }
     }
 
     /// connection flow control
     pub(crate) fn set_target_window_size(&mut self, size: WindowSize) {
-        self.streams.set_target_connection_window_size(size);
+        self.inner.streams.set_target_connection_window_size(size);
     }
 
     /// Send a new SETTINGS frame with an updated initial window size.
     pub(crate) fn set_initial_window_size(&mut self, size: WindowSize) -> Result<(), UserError> {
         let mut settings = frame::Settings::default();
         settings.set_initial_window_size(Some(size));
-        self.settings.send_settings(settings)
+        self.inner.settings.send_settings(settings)
     }
 
     /// Returns the maximum number of concurrent streams that may be initiated
     /// by this peer.
     pub(crate) fn max_send_streams(&self) -> usize {
-        self.streams.max_send_streams()
+        self.inner.streams.max_send_streams()
     }
 
     /// Returns `Ready` when the connection is ready to receive a frame.
@@ -131,16 +158,17 @@ where
     /// Returns `RecvError` as this may raise errors that are caused by delayed
     /// processing of received frames.
     fn poll_ready(&mut self, cx: &mut Context) -> Poll<Result<(), RecvError>> {
-        let _e = self.span.enter();
+        let _e = self.inner.span.enter();
         let span = tracing::trace_span!("poll_ready");
         let _e = span.enter();
         // The order of these calls don't really matter too much
-        ready!(self.ping_pong.send_pending_pong(cx, &mut self.codec))?;
+        ready!(self.inner.ping_pong.send_pending_pong(cx, &mut self.codec))?;
-        ready!(self.ping_pong.send_pending_ping(cx, &mut self.codec))?;
+        ready!(self.inner.ping_pong.send_pending_ping(cx, &mut self.codec))?;
         ready!(self
+            .inner
             .settings
-            .poll_send(cx, &mut self.codec, &mut self.streams))?;
+            .poll_send(cx, &mut self.codec, &mut self.inner.streams))?;
-        ready!(self.streams.send_pending_refusal(cx, &mut self.codec))?;
+        ready!(self.inner.streams.send_pending_refusal(cx, &mut self.codec))?;
 
         Poll::Ready(Ok(()))
     }
@@ -150,32 +178,15 @@ where
     /// This will return `Some(reason)` if the connection should be closed
     /// afterwards. If this is a graceful shutdown, this returns `None`.
     fn poll_go_away(&mut self, cx: &mut Context) -> Poll<Option<io::Result<Reason>>> {
-        self.go_away.send_pending_go_away(cx, &mut self.codec)
+        self.inner.go_away.send_pending_go_away(cx, &mut self.codec)
-    }
-
-    fn go_away(&mut self, id: StreamId, e: Reason) {
-        let frame = frame::GoAway::new(id, e);
-        self.streams.send_go_away(id);
-        self.go_away.go_away(frame);
-    }
-
-    fn go_away_now(&mut self, e: Reason) {
-        let last_processed_id = self.streams.last_processed_id();
-        let frame = frame::GoAway::new(last_processed_id, e);
-        self.go_away.go_away_now(frame);
     }
 
     pub fn go_away_from_user(&mut self, e: Reason) {
-        let last_processed_id = self.streams.last_processed_id();
+        self.inner.as_dyn().go_away_from_user(e)
-        let frame = frame::GoAway::new(last_processed_id, e);
-        self.go_away.go_away_from_user(frame);
-
-        // Notify all streams of reason we're abruptly closing.
-        self.streams.recv_err(&proto::Error::Proto(e));
     }
 
     fn take_error(&mut self, ours: Reason) -> Poll<Result<(), proto::Error>> {
-        let reason = if let Some(theirs) = self.error.take() {
+        let reason = if let Some(theirs) = self.inner.error.take() {
             match (ours, theirs) {
                 // If either side reported an error, return that
                 // to the user.
@@ -202,13 +213,13 @@ where
     pub fn maybe_close_connection_if_no_streams(&mut self) {
         // If we poll() and realize that there are no streams or references
         // then we can close the connection by transitioning to GOAWAY
-        if !self.streams.has_streams_or_other_references() {
+        if !self.inner.streams.has_streams_or_other_references() {
-            self.go_away_now(Reason::NO_ERROR);
+            self.inner.as_dyn().go_away_now(Reason::NO_ERROR);
         }
     }
 
     pub(crate) fn take_user_pings(&mut self) -> Option<UserPings> {
-        self.ping_pong.take_user_pings()
+        self.inner.ping_pong.take_user_pings()
     }
 
     /// Advances the internal state of the connection.
@@ -217,79 +228,39 @@ where
         // order to placate the borrow checker — `self` is mutably borrowed by
         // `poll2`, which means that we can't borrow `self.span` to enter it.
         // The clone is just an atomic ref bump.
-        let span = self.span.clone();
+        let span = self.inner.span.clone();
         let _e = span.enter();
         let span = tracing::trace_span!("poll");
         let _e = span.enter();
-        use crate::codec::RecvError::*;
 
         loop {
-            tracing::trace!(connection.state = ?self.state);
+            tracing::trace!(connection.state = ?self.inner.state);
             // TODO: probably clean up this glob of code
-            match self.state {
+            match self.inner.state {
                 // When open, continue to poll a frame
                 State::Open => {
-                    match self.poll2(cx) {
+                    let result = match self.poll2(cx) {
-                        // The connection has shutdown normally
+                        Poll::Ready(result) => result,
-                        Poll::Ready(Ok(())) => self.state = State::Closing(Reason::NO_ERROR),
                         // The connection is not ready to make progress
                         Poll::Pending => {
                             // Ensure all window updates have been sent.
                             //
                             // This will also handle flushing `self.codec`
-                            ready!(self.streams.poll_complete(cx, &mut self.codec))?;
+                            ready!(self.inner.streams.poll_complete(cx, &mut self.codec))?;
 
-                            if (self.error.is_some() || self.go_away.should_close_on_idle())
+                            if (self.inner.error.is_some()
-                                && !self.streams.has_streams()
+                                || self.inner.go_away.should_close_on_idle())
+                                && !self.inner.streams.has_streams()
                             {
-                                self.go_away_now(Reason::NO_ERROR);
+                                self.inner.as_dyn().go_away_now(Reason::NO_ERROR);
                                 continue;
                             }
 
                             return Poll::Pending;
                         }
-                        // Attempting to read a frame resulted in a connection level
+                    };
-                        // error. This is handled by setting a GOAWAY frame followed by
-                        // terminating the connection.
-                        Poll::Ready(Err(Connection(e))) => {
-                            tracing::debug!(error = ?e, "Connection::poll; connection error");
 
-                            // We may have already sent a GOAWAY for this error,
+                    self.inner.as_dyn().handle_poll2_result(result)?
-                            // if so, don't send another, just flush and close up.
-                            if let Some(reason) = self.go_away.going_away_reason() {
-                                if reason == e {
-                                    tracing::trace!("    -> already going away");
-                                    self.state = State::Closing(e);
-                                    continue;
-                                }
-                            }
-
-                            // Reset all active streams
-                            self.streams.recv_err(&e.into());
-                            self.go_away_now(e);
-                        }
-                        // Attempting to read a frame resulted in a stream level error.
-                        // This is handled by resetting the frame then trying to read
-                        // another frame.
-                        Poll::Ready(Err(Stream { id, reason })) => {
-                            tracing::trace!(?id, ?reason, "stream error");
-                            self.streams.send_reset(id, reason);
-                        }
-                        // Attempting to read a frame resulted in an I/O error. All
-                        // active streams must be reset.
-                        //
-                        // TODO: Are I/O errors recoverable?
-                        Poll::Ready(Err(Io(e))) => {
-                            tracing::debug!(error = ?e, "Connection::poll; IO error");
-                            let e = e.into();
-
-                            // Reset all active streams
-                            self.streams.recv_err(&e);
-
-                            // Return the error
-                            return Poll::Ready(Err(e));
-                        }
-                    }
                 }
                 State::Closing(reason) => {
                     tracing::trace!("connection closing after flush");
@@ -297,7 +268,7 @@ where
                     ready!(self.codec.shutdown(cx))?;
 
                     // Transition the state to error
-                    self.state = State::Closed(reason);
+                    self.inner.state = State::Closed(reason);
                 }
                 State::Closed(reason) => return self.take_error(reason),
             }
@@ -305,8 +276,6 @@ where
     }
 
     fn poll2(&mut self, cx: &mut Context) -> Poll<Result<(), RecvError>> {
-        use crate::frame::Frame::*;
-
         // This happens outside of the loop to prevent needing to do a clock
         // check and then comparison of the queue possibly multiple times a
         // second (and thus, the clock wouldn't have changed enough to matter).
@@ -319,8 +288,8 @@ where
             // - poll_go_away may buffer a graceful shutdown GOAWAY frame
             // - If it has, we've also added a PING to be sent in poll_ready
             if let Some(reason) = ready!(self.poll_go_away(cx)?) {
-                if self.go_away.should_close_now() {
+                if self.inner.go_away.should_close_now() {
-                    if self.go_away.is_user_initiated() {
+                    if self.inner.go_away.is_user_initiated() {
                         // A user initiated abrupt shutdown shouldn't return
                         // the same error back to the user.
                         return Poll::Ready(Ok(()));
@@ -337,61 +306,20 @@ where
             }
             ready!(self.poll_ready(cx))?;
 
-            match ready!(Pin::new(&mut self.codec).poll_next(cx)?) {
+            match self
-                Some(Headers(frame)) => {
+                .inner
-                    tracing::trace!(?frame, "recv HEADERS");
+                .as_dyn()
-                    self.streams.recv_headers(frame)?;
+                .recv_frame(ready!(Pin::new(&mut self.codec).poll_next(cx)?))?
+            {
+                ReceivedFrame::Settings(frame) => {
+                    self.inner.settings.recv_settings(
+                        frame,
+                        &mut self.codec,
+                        &mut self.inner.streams,
+                    )?;
                 }
-                Some(Data(frame)) => {
+                ReceivedFrame::Continue => (),
-                    tracing::trace!(?frame, "recv DATA");
+                ReceivedFrame::Done => {
-                    self.streams.recv_data(frame)?;
-                }
-                Some(Reset(frame)) => {
-                    tracing::trace!(?frame, "recv RST_STREAM");
-                    self.streams.recv_reset(frame)?;
-                }
-                Some(PushPromise(frame)) => {
-                    tracing::trace!(?frame, "recv PUSH_PROMISE");
-                    self.streams.recv_push_promise(frame)?;
-                }
-                Some(Settings(frame)) => {
-                    tracing::trace!(?frame, "recv SETTINGS");
-                    self.settings
-                        .recv_settings(frame, &mut self.codec, &mut self.streams)?;
-                }
-                Some(GoAway(frame)) => {
-                    tracing::trace!(?frame, "recv GOAWAY");
-                    // This should prevent starting new streams,
-                    // but should allow continuing to process current streams
-                    // until they are all EOS. Once they are, State should
-                    // transition to GoAway.
-                    self.streams.recv_go_away(&frame)?;
-                    self.error = Some(frame.reason());
-                }
-                Some(Ping(frame)) => {
-                    tracing::trace!(?frame, "recv PING");
-                    let status = self.ping_pong.recv_ping(frame);
-                    if status.is_shutdown() {
-                        assert!(
-                            self.go_away.is_going_away(),
-                            "received unexpected shutdown ping"
-                        );
-
-                        let last_processed_id = self.streams.last_processed_id();
-                        self.go_away(last_processed_id, Reason::NO_ERROR);
-                    }
-                }
-                Some(WindowUpdate(frame)) => {
-                    tracing::trace!(?frame, "recv WINDOW_UPDATE");
-                    self.streams.recv_window_update(frame)?;
-                }
-                Some(Priority(frame)) => {
-                    tracing::trace!(?frame, "recv PRIORITY");
-                    // TODO: handle
-                }
-                None => {
-                    tracing::trace!("codec closed");
-                    self.streams.recv_eof(false).expect("mutex poisoned");
                     return Poll::Ready(Ok(()));
                 }
             }
@@ -399,17 +327,190 @@ where
     }
 
     fn clear_expired_reset_streams(&mut self) {
-        self.streams.clear_expired_reset_streams();
+        self.inner.streams.clear_expired_reset_streams();
     }
 }
 
+impl<P, B> ConnectionInner<P, B>
+where
+    P: Peer,
+    B: Buf,
+{
+    fn as_dyn(&mut self) -> DynConnection<'_, B> {
+        let ConnectionInner {
+            state,
+            go_away,
+            streams,
+            error,
+            ping_pong,
+            ..
+        } = self;
+        let streams = streams.as_dyn();
+        DynConnection {
+            state,
+            go_away,
+            streams,
+            error,
+            ping_pong,
+        }
+    }
+}
+
+impl<B> DynConnection<'_, B>
+where
+    B: Buf,
+{
+    fn go_away(&mut self, id: StreamId, e: Reason) {
+        let frame = frame::GoAway::new(id, e);
+        self.streams.send_go_away(id);
+        self.go_away.go_away(frame);
+    }
+
+    fn go_away_now(&mut self, e: Reason) {
+        let last_processed_id = self.streams.last_processed_id();
+        let frame = frame::GoAway::new(last_processed_id, e);
+        self.go_away.go_away_now(frame);
+    }
+
+    fn go_away_from_user(&mut self, e: Reason) {
+        let last_processed_id = self.streams.last_processed_id();
+        let frame = frame::GoAway::new(last_processed_id, e);
+        self.go_away.go_away_from_user(frame);
+
+        // Notify all streams of reason we're abruptly closing.
+        self.streams.recv_err(&proto::Error::Proto(e));
+    }
+
+    fn handle_poll2_result(&mut self, result: Result<(), RecvError>) -> Result<(), Error> {
+        use crate::codec::RecvError::*;
+        match result {
+            // The connection has shutdown normally
+            Ok(()) => {
+                *self.state = State::Closing(Reason::NO_ERROR);
+                Ok(())
+            }
+            // Attempting to read a frame resulted in a connection level
+            // error. This is handled by setting a GOAWAY frame followed by
+            // terminating the connection.
+            Err(Connection(e)) => {
+                tracing::debug!(error = ?e, "Connection::poll; connection error");
+
+                // We may have already sent a GOAWAY for this error,
+                // if so, don't send another, just flush and close up.
+                if let Some(reason) = self.go_away.going_away_reason() {
+                    if reason == e {
+                        tracing::trace!("    -> already going away");
+                        *self.state = State::Closing(e);
+                        return Ok(());
+                    }
+                }
+
+                // Reset all active streams
+                self.streams.recv_err(&e.into());
+                self.go_away_now(e);
+                Ok(())
+            }
+            // Attempting to read a frame resulted in a stream level error.
+            // This is handled by resetting the frame then trying to read
+            // another frame.
+            Err(Stream { id, reason }) => {
+                tracing::trace!(?id, ?reason, "stream error");
+                self.streams.send_reset(id, reason);
+                Ok(())
+            }
+            // Attempting to read a frame resulted in an I/O error. All
+            // active streams must be reset.
+            //
+            // TODO: Are I/O errors recoverable?
+            Err(Io(e)) => {
+                tracing::debug!(error = ?e, "Connection::poll; IO error");
+                let e = e.into();
+
+                // Reset all active streams
+                self.streams.recv_err(&e);
+
+                // Return the error
+                Err(e)
+            }
+        }
+    }
+
+    fn recv_frame(&mut self, frame: Option<Frame>) -> Result<ReceivedFrame, RecvError> {
+        use crate::frame::Frame::*;
+        match frame {
+            Some(Headers(frame)) => {
+                tracing::trace!(?frame, "recv HEADERS");
+                self.streams.recv_headers(frame)?;
+            }
+            Some(Data(frame)) => {
+                tracing::trace!(?frame, "recv DATA");
+                self.streams.recv_data(frame)?;
+            }
+            Some(Reset(frame)) => {
+                tracing::trace!(?frame, "recv RST_STREAM");
+                self.streams.recv_reset(frame)?;
+            }
+            Some(PushPromise(frame)) => {
+                tracing::trace!(?frame, "recv PUSH_PROMISE");
+                self.streams.recv_push_promise(frame)?;
+            }
+            Some(Settings(frame)) => {
+                tracing::trace!(?frame, "recv SETTINGS");
+                return Ok(ReceivedFrame::Settings(frame));
+            }
+            Some(GoAway(frame)) => {
+                tracing::trace!(?frame, "recv GOAWAY");
+                // This should prevent starting new streams,
+                // but should allow continuing to process current streams
+                // until they are all EOS. Once they are, State should
+                // transition to GoAway.
+                self.streams.recv_go_away(&frame)?;
+                *self.error = Some(frame.reason());
+            }
+            Some(Ping(frame)) => {
+                tracing::trace!(?frame, "recv PING");
+                let status = self.ping_pong.recv_ping(frame);
+                if status.is_shutdown() {
+                    assert!(
+                        self.go_away.is_going_away(),
+                        "received unexpected shutdown ping"
+                    );
+
+                    let last_processed_id = self.streams.last_processed_id();
+                    self.go_away(last_processed_id, Reason::NO_ERROR);
+                }
+            }
+            Some(WindowUpdate(frame)) => {
+                tracing::trace!(?frame, "recv WINDOW_UPDATE");
+                self.streams.recv_window_update(frame)?;
+            }
+            Some(Priority(frame)) => {
+                tracing::trace!(?frame, "recv PRIORITY");
+                // TODO: handle
+            }
+            None => {
+                tracing::trace!("codec closed");
+                self.streams.recv_eof(false).expect("mutex poisoned");
+                return Ok(ReceivedFrame::Done);
+            }
+        }
+        Ok(ReceivedFrame::Continue)
+    }
+}
+
+enum ReceivedFrame {
+    Settings(frame::Settings),
+    Continue,
+    Done,
+}
+
 impl<T, B> Connection<T, client::Peer, B>
 where
     T: AsyncRead + AsyncWrite,
     B: Buf,
 {
     pub(crate) fn streams(&self) -> &Streams<B, client::Peer> {
-        &self.streams
+        &self.inner.streams
     }
 }
 
@@ -419,12 +520,12 @@ where
     B: Buf,
 {
     pub fn next_incoming(&mut self) -> Option<StreamRef<B>> {
-        self.streams.next_incoming()
+        self.inner.streams.next_incoming()
     }
 
     // Graceful shutdown only makes sense for server peers.
     pub fn go_away_gracefully(&mut self) {
-        if self.go_away.is_going_away() {
+        if self.inner.go_away.is_going_away() {
             // No reason to start a new one.
             return;
         }
@@ -440,11 +541,11 @@ where
         // > send another GOAWAY frame with an updated last stream identifier.
         // > This ensures that a connection can be cleanly shut down without
         // > losing requests.
-        self.go_away(StreamId::MAX, Reason::NO_ERROR);
+        self.inner.as_dyn().go_away(StreamId::MAX, Reason::NO_ERROR);
 
         // We take the advice of waiting 1 RTT literally, and wait
         // for a pong before proceeding.
-        self.ping_pong.ping_shutdown();
+        self.inner.ping_pong.ping_shutdown();
     }
 }
 
@@ -455,6 +556,6 @@ where
 {
     fn drop(&mut self) {
         // Ignore errors as this indicates that the mutex is poisoned.
-        let _ = self.streams.recv_eof(true);
+        let _ = self.inner.streams.recv_eof(true);
     }
 }

src/proto/mod.rs

@@ -10,7 +10,7 @@ pub(crate) use self::connection::{Config, Connection};
 pub(crate) use self::error::Error;
 pub(crate) use self::peer::{Dyn as DynPeer, Peer};
 pub(crate) use self::ping_pong::UserPings;
-pub(crate) use self::streams::{OpaqueStreamRef, StreamRef, Streams};
+pub(crate) use self::streams::{DynStreams, OpaqueStreamRef, StreamRef, Streams};
 pub(crate) use self::streams::{Open, PollReset, Prioritized};
 
 use crate::codec::Codec;

src/proto/streams/mod.rs

@@ -12,7 +12,7 @@ mod streams;
 pub(crate) use self::prioritize::Prioritized;
 pub(crate) use self::recv::Open;
 pub(crate) use self::send::PollReset;
-pub(crate) use self::streams::{OpaqueStreamRef, StreamRef, Streams};
+pub(crate) use self::streams::{DynStreams, OpaqueStreamRef, StreamRef, Streams};
 
 use self::buffer::Buffer;
 use self::counts::Counts;

src/proto/streams/prioritize.rs

@@ -545,43 +545,57 @@ impl Prioritize {
 
         // First check if there are any data chunks to take back
         if let Some(frame) = dst.take_last_data_frame() {
-            tracing::trace!(
+            self.reclaim_frame_inner(buffer, store, frame)
-                ?frame,
+        } else {
-                sz = frame.payload().inner.get_ref().remaining(),
+            false
-                "reclaimed"
+        }
-            );
+    }
 
-            let mut eos = false;
+    fn reclaim_frame_inner<B>(
-            let key = frame.payload().stream;
+        &mut self,
+        buffer: &mut Buffer<Frame<B>>,
+        store: &mut Store,
+        frame: frame::Data<Prioritized<B>>,
+    ) -> bool
+    where
+        B: Buf,
+    {
+        tracing::trace!(
+            ?frame,
+            sz = frame.payload().inner.get_ref().remaining(),
+            "reclaimed"
+        );
 
-            match mem::replace(&mut self.in_flight_data_frame, InFlightData::Nothing) {
+        let mut eos = false;
-                InFlightData::Nothing => panic!("wasn't expecting a frame to reclaim"),
+        let key = frame.payload().stream;
-                InFlightData::Drop => {
+
-                    tracing::trace!("not reclaiming frame for cancelled stream");
+        match mem::replace(&mut self.in_flight_data_frame, InFlightData::Nothing) {
-                    return false;
+            InFlightData::Nothing => panic!("wasn't expecting a frame to reclaim"),
-                }
+            InFlightData::Drop => {
-                InFlightData::DataFrame(k) => {
+                tracing::trace!("not reclaiming frame for cancelled stream");
-                    debug_assert_eq!(k, key);
+                return false;
-                }
+            }
+            InFlightData::DataFrame(k) => {
+                debug_assert_eq!(k, key);
+            }
+        }
+
+        let mut frame = frame.map(|prioritized| {
+            // TODO: Ensure fully written
+            eos = prioritized.end_of_stream;
+            prioritized.inner.into_inner()
+        });
+
+        if frame.payload().has_remaining() {
+            let mut stream = store.resolve(key);
+
+            if eos {
+                frame.set_end_stream(true);
             }
 
-            let mut frame = frame.map(|prioritized| {
+            self.push_back_frame(frame.into(), buffer, &mut stream);
-                // TODO: Ensure fully written
-                eos = prioritized.end_of_stream;
-                prioritized.inner.into_inner()
-            });
 
-            if frame.payload().has_remaining() {
+            return true;
-                let mut stream = store.resolve(key);
-
-                if eos {
-                    frame.set_end_stream(true);
-                }
-
-                self.push_back_frame(frame.into(), buffer, &mut stream);
-
-                return true;
-            }
         }
 
         false