From ef9fe7d78bf0030546693b95907c6382ea4df5e5 Mon Sep 17 00:00:00 2001
From: dignifiedquire <me@dignifiedquire.com>
Date: Mon, 11 Nov 2019 13:49:52 +0100
Subject: [PATCH] feat(net): implement more methods for TcpStream

- connect_timeout
- take_error

Ref #499
---
 src/net/tcp/stream.rs | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/src/net/tcp/stream.rs b/src/net/tcp/stream.rs
index 13a1752f2..37e4e0a33 100644
--- a/src/net/tcp/stream.rs
+++ b/src/net/tcp/stream.rs
@@ -1,6 +1,7 @@
 use std::io::{IoSlice, IoSliceMut, Read as _, Write as _};
 use std::net::SocketAddr;
 use std::pin::Pin;
+use std::time::Duration;
 
 use crate::future;
 use crate::io::{self, Read, Write};
@@ -96,6 +97,18 @@ impl TcpStream {
         }))
     }
 
+    /// Opens a TCP connection to a remote host with a timeout.
+    ///
+    /// Unlike `connect`, `connect_timeout` takes a single `SocketAddr` since
+    /// timeout must be applied to individual addresses.
+    ///
+    /// It is an error to pass a zero `Duration` to this function.
+    pub async fn connect_timeout(addr: &SocketAddr, timeout: Duration) -> io::Result<TcpStream> {
+        let stream = io::timeout(timeout, async move { TcpStream::connect(addr).await }).await?;
+
+        Ok(stream)
+    }
+
     /// Returns the local address that this stream is connected to.
     ///
     /// ## Examples
@@ -255,6 +268,29 @@ impl TcpStream {
         self.watcher.get_ref().set_nodelay(nodelay)
     }
 
+    /// Gets the value of the `SO_ERROR` option on this socket.
+    ///
+    /// This will retrieve the stored error in the underlying socket, clearing
+    /// the field in the process. This can be useful for checking errors between
+    /// calls.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # fn main() -> std::io::Result<()> { async_std::task::block_on(async {
+    /// #
+    /// use async_std::net::TcpStream;
+    ///
+    /// let stream = TcpStream::connect("127.0.0.1:8080").await
+    ///                        .expect("Couldn't connect to the server...");
+    /// stream.take_error().expect("No error was expected...");
+    ///
+    /// # Ok(()) }) }
+    /// ```
+    pub fn take_error(&self) -> io::Result<Option<io::Error>> {
+        self.watcher.get_ref().take_error()
+    }
+
     /// Shuts down the read, write, or both halves of this connection.
     ///
     /// This method will cause all pending and future I/O on the specified portions to return