SERVER-41397 Allow specifying a connection timeout

2 files changed, 14 insertions, 5 deletions

diff --git a/src/mongo/util/net/sock.cpp b/src/mongo/util/net/sock.cpp
index 39210d57a40..7345ed6194e 100644
--- a/src/mongo/util/net/sock.cpp
+++ b/src/mongo/util/net/sock.cpp
@@ -270,6 +270,12 @@ std::string Socket::getSNIServerName() const {
 #endif
 
 bool Socket::connect(SockAddr& remote) {
+    const Milliseconds connectTimeoutMillis(static_cast<int64_t>(
+        _timeout > 0 ? std::min(kMaxConnectTimeoutMS, (_timeout * 1000)) : kMaxConnectTimeoutMS));
+    return connect(remote, connectTimeoutMillis);
+}
+
+bool Socket::connect(SockAddr& remote, Milliseconds connectTimeoutMillis) {
     _remote = remote;
 
     _fd = ::socket(remote.getType(), SOCK_STREAM, 0);
@@ -283,10 +289,7 @@ bool Socket::connect(SockAddr& remote) {
         return false;
     }
 
-    const Milliseconds connectTimeoutMillis(static_cast<int64_t>(
-        _timeout > 0 ? std::min(kMaxConnectTimeoutMS, (_timeout * 1000)) : kMaxConnectTimeoutMS));
     const Date_t expiration = Date_t::now() + connectTimeoutMillis;
-
     bool connectSucceeded = ::connect(_fd, _remote.raw(), _remote.addressSize) == 0;
 
     if (!connectSucceeded) {
diff --git a/src/mongo/util/net/sock.h b/src/mongo/util/net/sock.h
index a9edfc71de6..6493806816a 100644
--- a/src/mongo/util/net/sock.h
+++ b/src/mongo/util/net/sock.h
@@ -54,6 +54,7 @@
 #include "mongo/logger/log_severity.h"
 #include "mongo/platform/compiler.h"
 #include "mongo/util/assert_util.h"
+#include "mongo/util/duration.h"
 #include "mongo/util/net/sockaddr.h"
 
 namespace mongo {
@@ -88,7 +89,7 @@ class Socket {
 public:
     static const int errorPollIntervalSecs;
 
-    Socket(int sock, const SockAddr& farEnd);
+    Socket(int sock, const SockAddr& remote);
 
     /** In some cases the timeout will actually be 2x this value - eg we do a partial send,
         then the timeout fires, then we try to send again, then the timeout fires again with
@@ -108,7 +109,12 @@ public:
      *  an error, or due to a timeout on connection, or due to the system socket deciding the
      *  socket is invalid.
      */
-    bool connect(SockAddr& farEnd);
+    bool connect(SockAddr& remote, Milliseconds connectTimeoutMillis);
+
+    /**
+     * Connect using a default connect timeout of min(_timeout * 1000, kMaxConnectTimeoutMS)
+     */
+    bool connect(SockAddr& remote);
 
     void close();
     void send(const char* data, int len, const char* context);