refactor: Better connection handling

src/server/conn.rs

@@ -0,0 +1,93 @@
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{self, Poll};
+
+use futures::io::{AsyncRead, AsyncWrite};
+use tower::Service;
+
+use super::pipeline::Connection;
+use super::Imap;
+use crate::proto::{Request, Response};
+
+#[pin_project::pin_project]
+pub struct Connecting<C, F, S> {
+    pub conn: Connection<C>,
+    #[pin]
+    pub state: ConnectingState<F, S>,
+    pub protocol: Imap,
+}
+
+#[pin_project::pin_project(project = ConnectingStateProj)]
+pub enum ConnectingState<F, S> {
+    Waiting {
+        #[pin]
+        service_fut: F,
+    },
+    Ready {
+        service: S,
+    },
+}
+
+impl<C, F, ME, S> Future for Connecting<C, F, S>
+where
+    C: AsyncRead + AsyncWrite + Unpin,
+    F: Future<Output = std::result::Result<S, ME>>,
+    ME: std::fmt::Display,
+    S: Service<Request, Response = Response>,
+    S::Error: std::fmt::Display,
+{
+    type Output = ();
+
+    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
+        use tokio_tower::pipeline::Server as PipelineServer;
+
+        let mut this = self.project();
+
+        loop {
+            let next = match this.state.as_mut().project() {
+                ConnectingStateProj::Waiting { service_fut } => {
+                    let service = match futures::ready!(service_fut.poll(cx)) {
+                        Ok(service) => service,
+                        Err(err) => {
+                            tracing::error!("Connection error: {}", err);
+                            return Poll::Ready(());
+                        }
+                    };
+
+                    // TODO: Properly handle server greeting
+                    {
+                        use imap_codec::types::response::{Response, Status};
+
+                        let status = match Status::ok(None, None, "Hello") {
+                            Ok(status) => status,
+                            Err(err) => {
+                                tracing::error!("Connection error: {}", err);
+                                return Poll::Ready(());
+                            }
+                        };
+                        let res = Response::Status(status);
+
+                        if let Err(err) = this.conn.send(res) {
+                            tracing::error!("Connection error: {}", err);
+                            return Poll::Ready(());
+                        };
+                    }
+
+                    ConnectingState::Ready { service }
+                }
+                ConnectingStateProj::Ready { service } => {
+                    let server = PipelineServer::new(this.conn, service);
+                    futures::pin_mut!(server);
+
+                    return server.poll(cx).map(|res| {
+                        if let Err(err) = res {
+                            tracing::debug!("Connection error: {}", err);
+                        }
+                    });
+                }
+            };
+
+            this.state.set(next);
+        }
+    }
+}

src/server/mod.rs

@@ -7,9 +7,11 @@ use imap_codec::types::response::Capability;
 
 use crate::proto::{Request, Response};
 use accept::Accept;
+use pipeline::Connection;
 pub use service::MakeServiceRef;
 
 pub mod accept;
+mod conn;
 mod pipeline;
 mod service;
 
@@ -70,7 +72,7 @@ where
                 let service_fut = this.make_service.make_service_ref(&conn);
 
                 tokio::task::spawn(conn::Connecting {
-                    conn,
+                    conn: Connection::new(conn),
                     state: conn::ConnectingState::Waiting { service_fut },
                     protocol: this.protocol.clone(),
                 });
@@ -115,93 +117,3 @@ impl<I: Accept> Builder<I> {
         }
     }
 }
-
-mod conn {
-    use std::future::Future;
-    use std::pin::Pin;
-    use std::task::{self, Poll};
-
-    use futures::io::{AsyncRead, AsyncWrite};
-    use tower::Service;
-
-    use super::Imap;
-    use crate::proto::{Request, Response};
-
-    #[pin_project::pin_project]
-    pub struct Connecting<C, F, S> {
-        pub conn: C,
-        #[pin]
-        pub state: ConnectingState<F, S>,
-        pub protocol: Imap,
-    }
-
-    #[pin_project::pin_project(project = ConnectingStateProj)]
-    pub enum ConnectingState<F, S> {
-        Waiting {
-            #[pin]
-            service_fut: F,
-        },
-        Ready {
-            service: S,
-        },
-    }
-
-    impl<C, F, ME, S> Future for Connecting<C, F, S>
-    where
-        C: AsyncRead + AsyncWrite + Unpin,
-        F: Future<Output = std::result::Result<S, ME>>,
-        ME: std::fmt::Display,
-        S: Service<Request, Response = Response>,
-        S::Error: std::fmt::Display,
-    {
-        type Output = ();
-
-        fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
-            use tokio_tower::pipeline::Server as PipelineServer;
-
-            use super::pipeline::Connection;
-
-            let mut this = self.project();
-
-            loop {
-                let next = match this.state.as_mut().project() {
-                    ConnectingStateProj::Waiting { service_fut } => {
-                        let service = match futures::ready!(service_fut.poll(cx)) {
-                            Ok(service) => service,
-                            Err(err) => {
-                                tracing::debug!("Connection error: {}", err);
-                                return Poll::Ready(());
-                            }
-                        };
-
-                        ConnectingState::Ready { service }
-                    }
-                    ConnectingStateProj::Ready { service } => {
-                        let mut conn = Connection::new(this.conn);
-
-                        // TODO: Properly handle server greeting
-                        {
-                            use imap_codec::types::response::{Response, Status};
-
-                            let status = Status::ok(None, None, "Hello").unwrap();
-                            let res = Response::Status(status);
-
-                            conn.send(res).unwrap();
-                        }
-
-                        let server = PipelineServer::new(conn, service);
-                        futures::pin_mut!(server);
-
-                        if let Err(err) = futures::ready!(server.poll(cx)) {
-                            tracing::debug!("Connection error: {}", err);
-                        }
-
-                        return Poll::Ready(());
-                    }
-                };
-
-                this.state.set(next);
-            }
-        }
-    }
-}

src/server/pipeline.rs

@@ -44,7 +44,13 @@ where
                 Ok(res) => res,
                 Err(e) if e.is_incomplete() => {
                     let mut buf = [0u8; 256];
+
                     let read = futures::ready!(this.conn.as_mut().poll_read(cx, &mut buf))?;
+                    tracing::trace!(read = read, "transport.poll_next");
+
+                    if read == 0 {
+                        return Poll::Ready(None);
+                    }
 
                     let data = &buf[..read];
                     this.read_buf.extend(data);
@@ -55,7 +61,7 @@ where
                     return Poll::Ready(Some(Err(format!("Error: {:?}", e).into())));
                 }
             };
-            tracing::debug!("Received: {:#?}", command);
+            tracing::debug!(command = ?command, "transport.recv");
 
             *this.read_buf = input.into();
 
@@ -74,6 +80,7 @@ where
 
         let mut this = self.project();
 
+        tracing::debug!(size = this.write_buf.len(), "transport.flush_buffer");
         while !this.write_buf.is_empty() {
             let written = futures::ready!(this.conn.as_mut().poll_write(cx, this.write_buf))?;
             this.write_buf.advance(written);
@@ -88,7 +95,7 @@ where
 
         let mut writer = BufMut::writer(&mut self.write_buf);
 
-        tracing::debug!("Sending: {:#?}", item);
+        tracing::debug!(item = ?item, "transport.send");
         item.encode(&mut writer)?;
 
         Ok(())