SERVER-29341 Set KeepAlive params on windows and mac

(cherry picked from commit bd153cba11ab82ffbf8d2bc0b801eb808369eb9e)

7 files changed, 205 insertions, 43 deletions

diff --git a/src/mongo/util/SConscript b/src/mongo/util/SConscript
index fe0380c1345..7918ca945ac 100644
--- a/src/mongo/util/SConscript
+++ b/src/mongo/util/SConscript
@@ -300,6 +300,12 @@ if env['MONGO_ALLOCATOR'] == 'tcmalloc':
         ],
     )
 
+env.Library(
+    target='winutil',
+    source=[
+        'winutil.cpp',
+    ],
+)
 
 env.Library(
     target='ntservice',
diff --git a/src/mongo/util/net/SConscript b/src/mongo/util/net/SConscript
index b53be63d404..59b546fad65 100644
--- a/src/mongo/util/net/SConscript
+++ b/src/mongo/util/net/SConscript
@@ -59,6 +59,7 @@ networkEnv.Library(
         '$BUILD_DIR/mongo/util/options_parser/options_parser',
         '$BUILD_DIR/third_party/shim_asio',
         'hostandport',
+        '$BUILD_DIR/mongo/util/winutil',
     ],
 )
 
diff --git a/src/mongo/util/net/sock.cpp b/src/mongo/util/net/sock.cpp
index dcc9b371892..457cd9d2bc8 100644
--- a/src/mongo/util/net/sock.cpp
+++ b/src/mongo/util/net/sock.cpp
@@ -49,6 +49,10 @@
 #if defined(__OpenBSD__)
 #include <sys/uio.h>
 #endif
+#else
+#include <mstcpip.h>
+#include <winsock2.h>
+#include <ws2tcpip.h>
 #endif
 
 #include "mongo/config.h"
@@ -65,6 +69,7 @@
 #include "mongo/util/net/socket_poll.h"
 #include "mongo/util/net/ssl_manager.h"
 #include "mongo/util/quick_exit.h"
+#include "mongo/util/winutil.h"
 
 namespace mongo {
 
@@ -105,6 +110,7 @@ void networkWarnWithDescription(const Socket& socket, StringData call, int error
 
 const double kMaxConnectTimeoutMS = 5000;
 
+
 }  // namespace
 
 static bool ipv6 = false;
@@ -140,23 +146,104 @@ void setSockTimeouts(int sock, double secs) {
 #endif
 }
 
-#if defined(_WIN32)
-void disableNagle(int sock) {
-    int x = 1;
-    if (setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, (char*)&x, sizeof(x)))
-        error() << "disableNagle failed: " << errnoWithDescription();
-    if (setsockopt(sock, SOL_SOCKET, SO_KEEPALIVE, (char*)&x, sizeof(x)))
-        error() << "SO_KEEPALIVE failed: " << errnoWithDescription();
-}
+#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 disableNagle(int sock) {
-    int x = 1;
+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
-    int level = 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
+}
+
+void disableNagle(int sock) {
+    int x = 1;
+#ifdef _WIN32
+    const int level = IPPROTO_TCP;
+#elif defined(SOL_TCP)
+    const int level = SOL_TCP;
 #else
-    int level = SOL_SOCKET;
+    const int level = SOL_SOCKET;
 #endif
 
     if (setsockopt(sock, level, TCP_NODELAY, (char*)&x, sizeof(x)))
@@ -165,34 +252,10 @@ void disableNagle(int sock) {
 #ifdef SO_KEEPALIVE
     if (setsockopt(sock, SOL_SOCKET, SO_KEEPALIVE, (char*)&x, sizeof(x)))
         error() << "SO_KEEPALIVE failed: " << errnoWithDescription();
-
-#ifdef __linux__
-    socklen_t len = sizeof(x);
-    if (getsockopt(sock, level, TCP_KEEPIDLE, (char*)&x, &len))
-        error() << "can't get TCP_KEEPIDLE: " << errnoWithDescription();
-
-    if (x > 300) {
-        x = 300;
-        if (setsockopt(sock, level, TCP_KEEPIDLE, (char*)&x, sizeof(x))) {
-            error() << "can't set TCP_KEEPIDLE: " << errnoWithDescription();
-        }
-    }
-
-    len = sizeof(x);  // just in case it changed
-    if (getsockopt(sock, level, TCP_KEEPINTVL, (char*)&x, &len))
-        error() << "can't get TCP_KEEPINTVL: " << errnoWithDescription();
-
-    if (x > 300) {
-        x = 300;
-        if (setsockopt(sock, level, TCP_KEEPINTVL, (char*)&x, sizeof(x))) {
-            error() << "can't set TCP_KEEPINTVL: " << errnoWithDescription();
-        }
-    }
-#endif
 #endif
-}
 
-#endif
+    setSocketKeepAliveParams(sock);
+}
 
 string getAddrInfoStrError(int code) {
 #if !defined(_WIN32)
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
diff --git a/src/mongo/util/ntservice.cpp b/src/mongo/util/ntservice.cpp
index 083b2fcc210..03cf4727cb2 100644
--- a/src/mongo/util/ntservice.cpp
+++ b/src/mongo/util/ntservice.cpp
@@ -287,7 +287,7 @@ void installServiceOrDie(const wstring& serviceName,
     SC_HANDLE schSCManager = ::OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
     if (schSCManager == NULL) {
         DWORD err = ::GetLastError();
-        log() << "Error connecting to the Service Control Manager: " << GetWinErrMsg(err);
+        log() << "Error connecting to the Service Control Manager: " << windows::GetErrMsg(err);
         quickExit(EXIT_NTSERVICE_ERROR);
     }
 
@@ -336,7 +336,7 @@ void installServiceOrDie(const wstring& serviceName,
                                   NULL);                      // user account password
     if (schService == NULL) {
         DWORD err = ::GetLastError();
-        log() << "Error creating service: " << GetWinErrMsg(err);
+        log() << "Error creating service: " << windows::GetErrMsg(err);
         ::CloseServiceHandle(schSCManager);
         quickExit(EXIT_NTSERVICE_ERROR);
     }
@@ -441,7 +441,7 @@ void removeServiceOrDie(const wstring& serviceName) {
     SC_HANDLE schSCManager = ::OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
     if (schSCManager == NULL) {
         DWORD err = ::GetLastError();
-        log() << "Error connecting to the Service Control Manager: " << GetWinErrMsg(err);
+        log() << "Error connecting to the Service Control Manager: " << windows::GetErrMsg(err);
         quickExit(EXIT_NTSERVICE_ERROR);
     }
 
diff --git a/src/mongo/util/winutil.cpp b/src/mongo/util/winutil.cpp
new file mode 100644
index 00000000000..9d82a029b1d
--- /dev/null
+++ b/src/mongo/util/winutil.cpp
@@ -0,0 +1,64 @@
+// winutil.cpp
+
+/*    Copyright 2017 10gen Inc.
+ *
+ *    This program is free software: you can redistribute it and/or  modify
+ *    it under the terms of the GNU Affero General Public License, version 3,
+ *    as published by the Free Software Foundation.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    GNU Affero General Public License for more details.
+ *
+ *    You should have received a copy of the GNU Affero General Public License
+ *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the GNU Affero General Public License in all respects
+ *    for all of the code used other than as permitted herein. If you modify
+ *    file(s) with this exception, you may extend this exception to your
+ *    version of the file(s), but you are not obligated to do so. If you do not
+ *    wish to do so, delete this exception statement from your version. If you
+ *    delete this exception statement from all source files in the program,
+ *    then also delete it in the license file.
+ */
+
+#if defined(_WIN32)
+
+#include "mongo/platform/basic.h"
+
+#include "mongo/util/winutil.h"
+
+#include "mongo/base/error_codes.h"
+#include "mongo/base/status.h"
+
+namespace mongo {
+
+StatusWith<boost::optional<DWORD>> windows::getDWORDRegistryKey(const CString& group,
+                                                                const CString& key) {
+    CRegKey regkey;
+    if (ERROR_SUCCESS != regkey.Open(HKEY_LOCAL_MACHINE, group, KEY_READ)) {
+        return Status(ErrorCodes::InternalError, "Unable to access windows registry");
+    }
+
+    DWORD val;
+    const auto res = regkey.QueryDWORDValue(key, val);
+    if (ERROR_INVALID_DATA == res) {
+        return Status(ErrorCodes::TypeMismatch,
+                      "Invalid data type in windows registry, expected DWORD");
+    }
+
+    if (ERROR_SUCCESS != res) {
+        return boost::none;
+    }
+
+    return val;
+}
+
+}  // namespace mongo
+
+#endif
diff --git a/src/mongo/util/winutil.h b/src/mongo/util/winutil.h
index a49fcc2f406..a9138e77113 100644
--- a/src/mongo/util/winutil.h
+++ b/src/mongo/util/winutil.h
@@ -33,11 +33,19 @@
 
 #if defined(_WIN32)
 #include "text.h"
+#include <atlbase.h>
+#include <atlstr.h>
+#include <boost/optional.hpp>
+#include <sstream>
+#include <string>
 #include <windows.h>
 
+#include <mongo/base/status_with.h>
+
 namespace mongo {
+namespace windows {
 
-inline std::string GetWinErrMsg(DWORD err) {
+inline std::string GetErrMsg(DWORD err) {
     LPTSTR errMsg;
     ::FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM,
                     NULL,
@@ -56,6 +64,23 @@ inline std::string GetWinErrMsg(DWORD err) {
 
     return output.str();
 }
-}
+
+/**
+ * Retrieve a DWORD value from the Local Machine Windows Registry for element:
+ * group\key
+ * e.g. HKEY_LOCAL_MACHINE\SYSTEM\CurrentControlSet\Services\Tcpip\Parameters\KeepAliveTime
+ *
+ * On success, returns:
+ *   boost::none if the key does not exist.
+ *   The value read from the registry.
+ *
+ * On failure, returns:
+ *   ErrorCodes::InternalError - Unable to access the registry group.
+ *   ErrorCodes::TypeMismatch - Key exists, but is of the wrong type.
+ */
+StatusWith<boost::optional<DWORD>> getDWORDRegistryKey(const CString& group, const CString& key);
+
+}  // namespace windows
+}  // namespace mongo
 
 #endif