Refactor TcpStream::connect into resolving loop and TcpStream::connect_to (#119)

src/net/tcp/stream.rs

@@ -78,18 +78,28 @@ impl TcpStream {
     /// # Ok(()) }) }
     /// ```
     pub async fn connect<A: ToSocketAddrs>(addrs: A) -> io::Result<TcpStream> {
-        enum State {
-            Waiting(TcpStream),
-            Error(io::Error),
-            Done,
-        }
-
         let mut last_err = None;
 
         for addr in addrs.to_socket_addrs()? {
-            let mut state = {
-                match mio::net::TcpStream::connect(&addr) {
-                    Ok(mio_stream) => {
+            let res = Self::connect_to(addr).await;
+
+            match res {
+                Ok(stream) => return Ok(stream),
+                Err(err) => last_err = Some(err),
+            }
+        }
+
+        Err(last_err.unwrap_or_else(|| {
+            io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "could not resolve to any addresses",
+            )
+        }))
+    }
+
+    /// Creates a new TCP stream connected to the specified address.
+    async fn connect_to(addr: SocketAddr) -> io::Result<TcpStream> {
+        let stream = mio::net::TcpStream::connect(&addr).map(|mio_stream| {
             #[cfg(unix)]
             let stream = TcpStream {
                 raw_fd: mio_stream.as_raw_fd(),
@@ -102,13 +112,19 @@ impl TcpStream {
                 io_handle: IoHandle::new(mio_stream),
             };
 
-                        State::Waiting(stream)
+            stream
+        });
+
+        enum State {
+            Waiting(TcpStream),
+            Error(io::Error),
+            Done,
         }
+        let mut state = match stream {
+            Ok(stream) => State::Waiting(stream),
             Err(err) => State::Error(err),
-                }
         };
-
-            let res = future::poll_fn(|cx| {
+        future::poll_fn(|cx| {
             match mem::replace(&mut state, State::Done) {
                 State::Waiting(stream) => {
                     // Once we've connected, wait for the stream to be writable as that's when
@@ -129,23 +145,10 @@ impl TcpStream {
                     Poll::Ready(Ok(stream))
                 }
                 State::Error(err) => Poll::Ready(Err(err)),
-                    State::Done => panic!("`TcpStream::connect()` future polled after completion"),
+                State::Done => panic!("`TcpStream::connect_to()` future polled after completion"),
             }
         })
-            .await;
-
-            match res {
-                Ok(stream) => return Ok(stream),
-                Err(err) => last_err = Some(err),
-            }
-        }
-
-        Err(last_err.unwrap_or_else(|| {
-            io::Error::new(
-                io::ErrorKind::InvalidInput,
-                "could not resolve to any addresses",
-            )
-        }))
+        .await
     }
 
     /// Returns the local address that this stream is connected to.