RUST-360 Guard connection establishment by connectTimeoutMS (#743)

This fixes a regression that was introduced in #721.

Author: patrickfreed

src/cmap/establish/mod.rs

@@ -1,5 +1,7 @@
 pub(super) mod handshake;
 
+use std::time::Duration;
+
 use self::handshake::{Handshaker, HandshakerOptions};
 use super::{
     conn::{ConnectionGeneration, LoadBalancedGeneration, PendingConnection},
@@ -13,7 +15,7 @@ use crate::{
     },
     error::{Error as MongoError, ErrorKind, Result},
     hello::HelloReply,
-    runtime::{AsyncStream, HttpClient, TlsConfig},
+    runtime::{self, stream::DEFAULT_CONNECT_TIMEOUT, AsyncStream, HttpClient, TlsConfig},
     sdam::HandshakePhase,
 };
 
@@ -26,11 +28,14 @@ pub(crate) struct ConnectionEstablisher {
 
     /// Cached configuration needed to create TLS connections, if needed.
     tls_config: Option<TlsConfig>,
+
+    connect_timeout: Duration,
 }
 
 pub(crate) struct EstablisherOptions {
     handshake_options: HandshakerOptions,
     tls_options: Option<TlsOptions>,
+    connect_timeout: Option<Duration>,
 }
 
 impl EstablisherOptions {
@@ -44,6 +49,7 @@ impl EstablisherOptions {
                 load_balanced: opts.load_balanced.unwrap_or(false),
             },
             tls_options: opts.tls_options(),
+            connect_timeout: opts.connect_timeout,
         }
     }
 }
@@ -59,14 +65,25 @@ impl ConnectionEstablisher {
             None
         };
 
+        let connect_timeout = match options.connect_timeout {
+            Some(d) if d.is_zero() => Duration::MAX,
+            Some(d) => d,
+            None => DEFAULT_CONNECT_TIMEOUT,
+        };
+
         Ok(Self {
             handshaker,
             tls_config,
+            connect_timeout,
         })
     }
 
     async fn make_stream(&self, address: ServerAddress) -> Result<AsyncStream> {
-        AsyncStream::connect(address, self.tls_config.as_ref()).await
+        runtime::timeout(
+            self.connect_timeout,
+            AsyncStream::connect(address, self.tls_config.as_ref()),
+        )
+        .await?
     }
 
     /// Establishes a connection.

src/runtime/mod.rs

@@ -109,7 +109,13 @@ pub(crate) async fn delay_for(delay: Duration) {
 
     #[cfg(feature = "async-std-runtime")]
     {
-        async_std::task::sleep(delay).await
+        // This avoids a panic in async-std when the provided duration is too large.
+        // See: https://github.com/async-rs/async-std/issues/1037.
+        if delay == Duration::MAX {
+            std::future::pending().await
+        } else {
+            async_std::task::sleep(delay).await
+        }
     }
 }
 
@@ -124,9 +130,15 @@ pub(crate) async fn timeout<F: Future>(timeout: Duration, future: F) -> Result<F
 
     #[cfg(feature = "async-std-runtime")]
     {
-        async_std::future::timeout(timeout, future)
-            .await
-            .map_err(|_| std::io::ErrorKind::TimedOut.into())
+        // This avoids a panic on async-std when the provided duration is too large.
+        // See: https://github.com/async-rs/async-std/issues/1037.
+        if timeout == Duration::MAX {
+            Ok(future.await)
+        } else {
+            async_std::future::timeout(timeout, future)
+                .await
+                .map_err(|_| std::io::ErrorKind::TimedOut.into())
+        }
     }
 }
 

src/sdam/monitor.rs

@@ -1,5 +1,4 @@
 use std::{
-    future,
     sync::Arc,
     time::{Duration, Instant},
 };
@@ -185,7 +184,7 @@ impl Monitor {
         });
 
         let heartbeat_frequency = self.heartbeat_frequency();
-        let timeout = if self.connect_timeout().as_millis() == 0 {
+        let timeout = if self.connect_timeout().is_zero() {
             // If connectTimeoutMS = 0, then the socket timeout for monitoring is unlimited.
             Duration::MAX
         } else if self.topology_version.is_some() {
@@ -198,16 +197,6 @@ impl Monitor {
             // Otherwise, just use connectTimeoutMS.
             self.connect_timeout()
         };
-        let timeout_future = async {
-            // If timeout is infinite, don't bother creating a delay future.
-            // This also avoids a panic on async-std when the provided duration is too large.
-            // See: https://github.com/async-rs/async-std/issues/1037.
-            if timeout == Duration::MAX {
-                future::pending().await
-            } else {
-                runtime::delay_for(timeout).await
-            }
-        };
 
         let execute_hello = async {
             match self.connection {
@@ -268,7 +257,7 @@ impl Monitor {
                 };
                 HelloResult::Cancelled { reason: reason_error }
             }
-            _ = timeout_future => {
+            _ = runtime::delay_for(timeout) => {
                 HelloResult::Err(Error::network_timeout())
             }
         };