SERVER-29359 SO_KEEPALIVE for nia sockets

Previously had not set keepalives for egress sockets on executors, this fixes that.

7 files changed, 122 insertions, 0 deletions

diff --git a/src/mongo/executor/SConscript b/src/mongo/executor/SConscript
index ed040af2910..0359c40a020 100644
--- a/src/mongo/executor/SConscript
+++ b/src/mongo/executor/SConscript
@@ -168,6 +168,7 @@ env.Library(
     LIBDEPS=[
         '$BUILD_DIR/mongo/base/system_error',
         '$BUILD_DIR/mongo/client/authentication',
+        '$BUILD_DIR/mongo/util/net/network',
         '$BUILD_DIR/third_party/shim_asio',
         'task_executor_interface',
     ]
diff --git a/src/mongo/executor/async_secure_stream.cpp b/src/mongo/executor/async_secure_stream.cpp
index 0c7969024f7..a7c051bc34b 100644
--- a/src/mongo/executor/async_secure_stream.cpp
+++ b/src/mongo/executor/async_secure_stream.cpp
@@ -74,6 +74,11 @@ void AsyncSecureStream::connect(asio::ip::tcp::resolver::iterator endpoints,
                 return _userHandler(ec);
             }
 
+            ec = setStreamKeepAlive(&_stream.next_layer());
+            if (ec) {
+                return _userHandler(ec);
+            }
+
             _connected = true;
             return _handleConnect(std::move(iter));
         }));
diff --git a/src/mongo/executor/async_stream.cpp b/src/mongo/executor/async_stream.cpp
index 075a9094dd4..aeafd8a0b8b 100644
--- a/src/mongo/executor/async_stream.cpp
+++ b/src/mongo/executor/async_stream.cpp
@@ -70,6 +70,11 @@ void AsyncStream::connect(tcp::resolver::iterator iter, ConnectHandler&& connect
                 return connectHandler(ec);
             }
 
+            ec = setStreamKeepAlive(&_stream);
+            if (ec) {
+                return connectHandler(ec);
+            }
+
             _connected = true;
             return connectHandler(ec);
         }));
diff --git a/src/mongo/executor/async_stream_common.cpp b/src/mongo/executor/async_stream_common.cpp
index 92708862c94..97cfb3115b0 100644
--- a/src/mongo/executor/async_stream_common.cpp
+++ b/src/mongo/executor/async_stream_common.cpp
@@ -57,6 +57,11 @@ void logFailureInSetStreamNoDelay(std::error_code ec) {
     severe() << "Failed to set no-delay mode on stream: " << ec.message();
 }
 
+void logFailureInSetStreamKeepAlive(std::error_code ec) {
+    invariant(ec);
+    severe() << "Failed to set keep-alive mode on stream: " << ec.message();
+}
+
 void logUnexpectedErrorInCheckOpen(std::error_code ec) {
     invariant(ec);
     log() << "unexpected error when checking if a stream was open: " << ec.message()
diff --git a/src/mongo/executor/async_stream_common.h b/src/mongo/executor/async_stream_common.h
index 67c12bbff3b..5ca7dcad5fe 100644
--- a/src/mongo/executor/async_stream_common.h
+++ b/src/mongo/executor/async_stream_common.h
@@ -31,6 +31,7 @@
 #include <utility>
 
 #include "mongo/util/assert_util.h"
+#include "mongo/util/net/sock.h"
 
 namespace mongo {
 namespace executor {
@@ -93,6 +94,7 @@ void cancelStream(ASIOStream* stream) {
 
 void logFailureInSetStreamNonBlocking(std::error_code ec);
 void logFailureInSetStreamNoDelay(std::error_code ec);
+void logFailureInSetStreamKeepAlive(std::error_code ec);
 
 template <typename ASIOStream>
 std::error_code setStreamNonBlocking(ASIOStream* stream) {
@@ -114,6 +116,17 @@ std::error_code setStreamNoDelay(ASIOStream* stream) {
     return ec;
 }
 
+template <typename ASIOStream>
+std::error_code setStreamKeepAlive(ASIOStream* stream) {
+    std::error_code ec;
+    stream->set_option(asio::socket_base::keep_alive(true), ec);
+    if (ec) {
+        logFailureInSetStreamKeepAlive(ec);
+    }
+    setSocketKeepAliveParams(stream->native_handle());
+    return ec;
+}
+
 void logUnexpectedErrorInCheckOpen(std::error_code ec);
 
 template <typename ASIOStream>
diff --git a/src/mongo/util/net/sock.cpp b/src/mongo/util/net/sock.cpp
index dcc9b371892..c44fe3c21f6 100644
--- a/src/mongo/util/net/sock.cpp
+++ b/src/mongo/util/net/sock.cpp
@@ -140,6 +140,96 @@ void setSockTimeouts(int sock, double secs) {
 #endif
 }
 
+#ifdef _WIN32
+#ifdef _UNICODE
+#define X_STR_CONST(str) (L##str)
+#else
+#define X_STR_CONST(str) (str)
+#endif
+const CString kKeepAliveGroup(
+    X_STR_CONST("SYSTEM\\CurrentControlSet\\Services\\Tcpip\\Parameters"));
+const CString kKeepAliveTime(X_STR_CONST("KeepAliveTime"));
+const CString kKeepAliveInterval(X_STR_CONST("KeepAliveInterval"));
+#undef X_STR_CONST
+#endif
+
+void setSocketKeepAliveParams(int sock,
+                              unsigned int maxKeepIdleSecs,
+                              unsigned int maxKeepIntvlSecs) {
+#ifdef _WIN32
+    // Defaults per MSDN when registry key does not exist.
+    // Expressed in seconds here to be consistent with posix,
+    // though Windows uses milliseconds.
+    const DWORD kWindowsKeepAliveTimeSecsDefault = 2 * 60 * 60;
+    const DWORD kWindowsKeepAliveIntervalSecsDefault = 1;
+
+    const auto getKey = [](const CString& key, DWORD default_value) {
+        auto withval = windows::getDWORDRegistryKey(kKeepAliveGroup, key);
+        if (withval.isOK()) {
+            auto val = withval.getValue();
+            // Return seconds
+            return val ? (val.get() / 1000) : default_value;
+        }
+        error() << "can't get KeepAlive parameter: " << withval.getStatus();
+        return default_value;
+    };
+
+    const auto keepIdleSecs = getKey(kKeepAliveTime, kWindowsKeepAliveTimeSecsDefault);
+    const auto keepIntvlSecs = getKey(kKeepAliveInterval, kWindowsKeepAliveIntervalSecsDefault);
+
+    if ((keepIdleSecs > maxKeepIdleSecs) || (keepIntvlSecs > maxKeepIntvlSecs)) {
+        DWORD sent = 0;
+        struct tcp_keepalive keepalive;
+        keepalive.onoff = TRUE;
+        keepalive.keepalivetime = std::min<DWORD>(keepIdleSecs, maxKeepIdleSecs) * 1000;
+        keepalive.keepaliveinterval = std::min<DWORD>(keepIntvlSecs, maxKeepIntvlSecs) * 1000;
+        if (WSAIoctl(sock,
+                     SIO_KEEPALIVE_VALS,
+                     &keepalive,
+                     sizeof(keepalive),
+                     nullptr,
+                     0,
+                     &sent,
+                     nullptr,
+                     nullptr)) {
+            error() << "failed setting keepalive values: " << WSAGetLastError();
+        }
+    }
+#elif defined(__APPLE__) || defined(__linux__)
+    const auto updateSockOpt =
+        [sock](int level, int optnum, unsigned int maxval, StringData optname) {
+            unsigned int optval = 1;
+            socklen_t len = sizeof(optval);
+
+            if (getsockopt(sock, level, optnum, (char*)&optval, &len)) {
+                error() << "can't get " << optname << ": " << errnoWithDescription();
+            }
+
+            if (optval > maxval) {
+                optval = maxval;
+                if (setsockopt(sock, level, optnum, (char*)&optval, sizeof(optval))) {
+                    error() << "can't set " << optname << ": " << errnoWithDescription();
+                }
+            }
+        };
+
+#ifdef __APPLE__
+    updateSockOpt(IPPROTO_TCP, TCP_KEEPALIVE, maxKeepIdleSecs, "TCP_KEEPALIVE");
+#endif
+
+#ifdef __linux__
+#ifdef SOL_TCP
+    const int level = SOL_TCP;
+#else
+    const int level = SOL_SOCKET;
+#endif
+    updateSockOpt(level, TCP_KEEPIDLE, maxKeepIdleSecs, "TCP_KEEPIDLE");
+    updateSockOpt(level, TCP_KEEPINTVL, maxKeepIntvlSecs, "TCP_KEEPINTVL");
+#endif
+
+#endif
+}
+
 #if defined(_WIN32)
 void disableNagle(int sock) {
     int x = 1;
diff --git a/src/mongo/util/net/sock.h b/src/mongo/util/net/sock.h
index 44a83b6ef5b..f6fea719279 100644
--- a/src/mongo/util/net/sock.h
+++ b/src/mongo/util/net/sock.h
@@ -70,6 +70,9 @@ extern const int portRecvFlags;
 
 const int SOCK_FAMILY_UNKNOWN_ERROR = 13078;
 
+void setSocketKeepAliveParams(int sock,
+                              unsigned int maxKeepIdleSecs = 300,
+                              unsigned int maxKeepIntvlSecs = 300);
 void disableNagle(int sock);
 
 // Generate a string representation for getaddrinfo return codes