SERVER-19856: Register for PRESHUTDOWN notifications on Windows Vista+

1 files changed, 34 insertions, 9 deletions

diff --git a/src/mongo/util/ntservice.cpp b/src/mongo/util/ntservice.cpp
index 18a44334b40..02f03b7e989 100644
--- a/src/mongo/util/ntservice.cpp
+++ b/src/mongo/util/ntservice.cpp
@@ -78,7 +78,8 @@ bool shouldStartService() {
     return _startService;
 }
 
-static void WINAPI serviceCtrl(DWORD ctrlCode);
+static DWORD WINAPI
+serviceCtrl(DWORD dwControl, DWORD dwEventType, LPVOID lpEventData, LPVOID lpContext);
 
 void configureService(ServiceCallback serviceCallback,
                       const moe::Environment& params,
@@ -412,6 +413,21 @@ void installServiceOrDie(const wstring& serviceName,
         log() << "Could not set service description. Check the Windows Event Log for more details.";
     }
 
+    // Set the pre-shutdown notification with a timeout of 10 minutes.
+    // Windows will either wait for us to finish with SERVICE_STOPPED or it will timeout, whichever
+    // is first.
+    SERVICE_PRESHUTDOWN_INFO servicePreshutdownInfo;
+    servicePreshutdownInfo.dwPreshutdownTimeout = 10 * 60 * 1000;  // 10 minutes
+
+    BOOL ret = ::ChangeServiceConfig2(
+        schService, SERVICE_CONFIG_PRESHUTDOWN_INFO, &servicePreshutdownInfo);
+    if (!ret) {
+        DWORD gle = ::GetLastError();
+        error() << "Failed to set timeout for pre-shutdown notification with error: "
+                << errnoWithDescription(gle);
+        serviceInstalled = false;
+    }
+
     ::CloseServiceHandle(schService);
     ::CloseServiceHandle(schSCManager);
 
@@ -483,7 +499,7 @@ bool reportStatus(DWORD reportState, DWORD waitHint, DWORD exitCode) {
             dwControlsAccepted = 0;
             break;
         default:
-            dwControlsAccepted = SERVICE_ACCEPT_STOP | SERVICE_ACCEPT_SHUTDOWN;
+            dwControlsAccepted = SERVICE_ACCEPT_STOP | SERVICE_ACCEPT_PRESHUTDOWN;
             break;
     }
 
@@ -538,7 +554,7 @@ static void serviceStop() {
 }
 
 static void WINAPI initService(DWORD argc, LPTSTR* argv) {
-    _statusHandle = RegisterServiceCtrlHandler(_serviceName.c_str(), serviceCtrl);
+    _statusHandle = RegisterServiceCtrlHandlerEx(_serviceName.c_str(), serviceCtrl, NULL);
     if (!_statusHandle)
         return;
 
@@ -570,15 +586,24 @@ static void serviceShutdown(const char* controlCodeName) {
     // Note: we will report exit status in initService
 }
 
-static void WINAPI serviceCtrl(DWORD ctrlCode) {
-    switch (ctrlCode) {
+static DWORD WINAPI
+serviceCtrl(DWORD dwControl, DWORD dwEventType, LPVOID lpEventData, LPVOID lpContext) {
+    switch (dwControl) {
+        case SERVICE_CONTROL_INTERROGATE:
+            // Return NO_ERROR per MSDN even though we do nothing for this control code.
+            return NO_ERROR;
         case SERVICE_CONTROL_STOP:
             serviceShutdown("SERVICE_CONTROL_STOP");
-            break;
-        case SERVICE_CONTROL_SHUTDOWN:
-            serviceShutdown("SERVICE_CONTROL_SHUTDOWN");
-            break;
+            // Return NO_ERROR since we handle the STOP
+            return NO_ERROR;
+        case SERVICE_CONTROL_PRESHUTDOWN:
+            serviceShutdown("SERVICE_CONTROL_PRESHUTDOWN");
+            // Return NO_ERROR since we handle the PRESHUTDOWN
+            return NO_ERROR;
     }
+
+    // Return ERROR_CALL_NOT_IMPLEMENTED as the default
+    return ERROR_CALL_NOT_IMPLEMENTED;
 }
 
 void startService() {