Refactor Connection to a synchronous state machine

src/connection.rs

@@ -87,6 +87,7 @@
 //   command processing instead of having to scan the whole collection of
 //   `Stream`s on each loop iteration, which is not great.
 
+mod closing;
 mod control;
 mod stream;
 
@@ -109,6 +110,7 @@ use std::collections::VecDeque;
 use std::task::Context;
 use std::{fmt, sync::Arc, task::Poll};
 
+use crate::connection::closing::Closing;
 pub use control::Control;
 pub use stream::{Packet, State, Stream};
 
@@ -219,7 +221,7 @@ enum ConnectionState<T> {
     Active(Active<T>),
     /// Our user requested to shutdown the connection, we are working on it.
     Closing {
-        inner: BoxFuture<'static, Result<()>>,
+        inner: Closing<T>,
         reply: oneshot::Sender<()>,
     },
     /// An error occurred and we are cleaning up our resources.
@@ -256,7 +258,7 @@ impl<T: AsyncRead + AsyncWrite + Unpin + Send + 'static> ConnectionState<T> {
                         }
                         Poll::Ready(Some(ControlCommand::CloseConnection(reply))) => {
                             *self = ConnectionState::Closing {
-                                inner: active.close().boxed(),
+                                inner: active.close(),
                                 reply,
                             };
                             continue;
@@ -449,59 +451,13 @@ impl<T: AsyncRead + AsyncWrite + Unpin> Active<T> {
     }
 
     /// Gracefully close the connection to the remote.
-    async fn close(mut self) -> Result<()> {
-        // Prevent more control commands being sent.
-        self.control_receiver.close();
-
-        // Drain all remaining control commands.
-        while let Some(cmd) = self.control_receiver.next().await {
-            match cmd {
-                ControlCommand::OpenStream(reply) => {
-                    let _ = reply.send(Err(ConnectionError::Closed));
-                }
-                ControlCommand::CloseConnection(reply) => {
-                    let _ = reply.send(());
-                }
-            }
-        }
-
-        // Prevent any more work being sent from streams.
-        self.stream_receiver.close();
-
-        // Drain all remaining stream commands.
-        while let Some(cmd) = self.stream_receiver.next().await {
-            match cmd {
-                StreamCommand::SendFrame(frame) => {
-                    self.on_send_frame(frame);
-                }
-                StreamCommand::CloseStream { id, ack } => {
-                    self.on_close_stream(id, ack);
-                }
-            }
-        }
-
-        // Send `GoAway` frame.
-        log::debug!("{}: sending term", self.id);
-        self.pending_frames.push_back(Frame::term().into());
-
-        // Drain all remaining frames.
-        while let Some(pending_frame) = self.pending_frames.pop_front() {
-            self.socket
-                .feed(pending_frame)
-                .await
-                .or(Err(ConnectionError::Closed))?;
-        }
-
-        self.socket.flush().await.or(Err(ConnectionError::Closed))?;
-
-        // Close the socket.
-        self.socket
-            .get_mut()
-            .close()
-            .await
-            .or(Err(ConnectionError::Closed))?;
-
-        Ok(())
+    fn close(self) -> Closing<T> {
+        Closing::new(
+            self.control_receiver,
+            self.stream_receiver,
+            self.pending_frames,
+            self.socket,
+        )
     }
 
     /// Cleanup all our resources.

src/connection/closing.rs

@@ -0,0 +1,117 @@
+use crate::connection::Result;
+use crate::connection::{ControlCommand, StreamCommand};
+use crate::frame::Frame;
+use crate::{frame, ConnectionError};
+use futures::channel::mpsc;
+use futures::stream::Fuse;
+use futures::{ready, AsyncRead, AsyncWrite, SinkExt, StreamExt};
+use std::collections::VecDeque;
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+/// A [`Future`] that gracefully closes the yamux connection.
+#[must_use]
+pub struct Closing<T> {
+    state: State,
+    control_receiver: mpsc::Receiver<ControlCommand>,
+    stream_receiver: mpsc::Receiver<StreamCommand>,
+    pending_frames: VecDeque<Frame<()>>,
+    socket: Fuse<frame::Io<T>>,
+}
+
+impl<T> Closing<T>
+where
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    pub(crate) fn new(
+        control_receiver: mpsc::Receiver<ControlCommand>,
+        stream_receiver: mpsc::Receiver<StreamCommand>,
+        pending_frames: VecDeque<Frame<()>>,
+        socket: Fuse<frame::Io<T>>,
+    ) -> Self {
+        Self {
+            state: State::ClosingControlReceiver,
+            control_receiver,
+            stream_receiver,
+            pending_frames,
+            socket,
+        }
+    }
+}
+
+impl<T> Future for Closing<T>
+where
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    type Output = Result<()>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let mut this = self.get_mut();
+
+        loop {
+            match this.state {
+                State::ClosingControlReceiver => {
+                    this.control_receiver.close();
+                    this.state = State::DrainingControlReceiver;
+                }
+                State::DrainingControlReceiver => {
+                    match ready!(this.control_receiver.poll_next_unpin(cx)) {
+                        Some(ControlCommand::OpenStream(reply)) => {
+                            let _ = reply.send(Err(ConnectionError::Closed));
+                        }
+                        Some(ControlCommand::CloseConnection(reply)) => {
+                            let _ = reply.send(());
+                        }
+                        None => this.state = State::ClosingStreamReceiver,
+                    }
+                }
+                State::ClosingStreamReceiver => {
+                    this.stream_receiver.close();
+                    this.state = State::DrainingStreamReceiver;
+                }
+
+                State::DrainingStreamReceiver => {
+                    this.stream_receiver.close();
+
+                    match ready!(this.stream_receiver.poll_next_unpin(cx)) {
+                        Some(StreamCommand::SendFrame(frame)) => {
+                            this.pending_frames.push_back(frame.into())
+                        }
+                        Some(StreamCommand::CloseStream { id, ack }) => this
+                            .pending_frames
+                            .push_back(Frame::close_stream(id, ack).into()),
+                        None => this.state = State::SendingTermFrame,
+                    }
+                }
+                State::SendingTermFrame => {
+                    this.pending_frames.push_back(Frame::term().into());
+                    this.state = State::FlushingPendingFrames;
+                }
+                State::FlushingPendingFrames => {
+                    ready!(this.socket.poll_ready_unpin(cx))?;
+
+                    match this.pending_frames.pop_front() {
+                        Some(frame) => this.socket.start_send_unpin(frame)?,
+                        None => this.state = State::ClosingSocket,
+                    }
+                }
+                State::ClosingSocket => {
+                    ready!(this.socket.poll_close_unpin(cx))?;
+
+                    return Poll::Ready(Ok(()));
+                }
+            }
+        }
+    }
+}
+
+pub enum State {
+    ClosingControlReceiver,
+    DrainingControlReceiver,
+    ClosingStreamReceiver,
+    DrainingStreamReceiver,
+    SendingTermFrame,
+    FlushingPendingFrames,
+    ClosingSocket,
+}