From f3d91adc48a7618fd9c94cbc45143f89b47b59e4 Mon Sep 17 00:00:00 2001
From: Iain Patterson <me@iain.cx>
Date: Fri, 15 Nov 2013 11:24:45 +0000
Subject: [PATCH] Added await_shutdown() function.

The system expects service STOP requests to be honoured promptly.  If
service shutdown will take longer than 30 seconds we must update the
service status checkpoint variable before then.

We also need to ensure that the service status wait hint time is
strictly increasing every time we update the checkpoint, otherwise the
service will be considered to be hung after 60 seconds.
---
 messages.mc | 16 ++++++++++
 nssm.h      |  3 ++
 service.cpp | 86 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 service.h   |  1 +
 4 files changed, 106 insertions(+)

diff --git a/messages.mc b/messages.mc
index fdf295b..a2da26f 100644
--- a/messages.mc
+++ b/messages.mc
@@ -1344,3 +1344,19 @@ The registry value %2, used to specify the maximum number of milliseconds to wai
 Language = Italian
 The registry value %2, used to specify the maximum number of milliseconds to wait for service %1 to stop after posting a WM_QUIT message to the message queues of threads managed by the application, was not of type REG_DWORD.  The default time of %3 milliseconds will be used.
 .
+
+MessageId = +1
+SymbolicName = NSSM_EVENT_AWAITING_SHUTDOWN
+Severity = Informational
+Language = English
+%1 has waited %3 of %5 milliseconds for the %2 service to exit.
+Next update in %4 milliseconds.
+.
+Language = French
+%1 has waited %3 of %5 milliseconds for the %2 service to exit.
+Next update in %4 milliseconds.
+.
+Language = Italian
+%1 has waited %3 of %5 milliseconds for the %2 service to exit.
+Next update in %4 milliseconds.
+.
diff --git a/nssm.h b/nssm.h
index 01228bd..bfde164 100644
--- a/nssm.h
+++ b/nssm.h
@@ -66,4 +66,7 @@ int str_equiv(const char *, const char *);
 #define NSSM_STOP_METHOD_THREADS (1 << 2)
 #define NSSM_STOP_METHOD_TERMINATE (1 << 3)
 
+/* How many milliseconds to wait before updating service status. */
+#define NSSM_SHUTDOWN_CHECKPOINT 20000
+
 #endif
diff --git a/service.cpp b/service.cpp
index f14390f..028b035 100644
--- a/service.cpp
+++ b/service.cpp
@@ -599,3 +599,89 @@ void throttle_restart() {
     else Sleep(ms);
   }
 }
+
+/*
+  When responding to a stop (or any other) request we need to set dwWaitHint to
+  the number of milliseconds we expect the operation to take, and optionally
+  increase dwCheckPoint.  If dwWaitHint milliseconds elapses without the
+  operation completing or dwCheckPoint increasing, the system will consider the
+  service to be hung.
+
+  However the system will consider the service to be hung after 30000
+  milliseconds regardless of the value of dwWaitHint if dwCheckPoint has not
+  changed.  Therefore if we want to wait longer than that we must periodically
+  increase dwCheckPoint.
+
+  Furthermore, it will consider the service to be hung after 60000 milliseconds
+  regardless of the value of dwCheckPoint unless dwWaitHint is increased every
+  time dwCheckPoint is also increased.
+
+  Our strategy then is to retrieve the initial dwWaitHint and wait for
+  NSSM_SHUTDOWN_CHECKPOINT milliseconds.  If the process is still running and
+  we haven't finished waiting we increment dwCheckPoint and add whichever is
+  smaller of NSSM_SHUTDOWN_CHECKPOINT or the remaining timeout to dwWaitHint.
+
+  Only doing both these things will prevent the system from killing the service.
+
+  Returns: 1 if the wait timed out.
+           0 if the wait completed.
+          -1 on error.
+*/
+int await_shutdown(char *function_name, char *service_name, SERVICE_STATUS_HANDLE service_handle, SERVICE_STATUS *service_status, HANDLE process_handle, unsigned long timeout) {
+  unsigned long interval;
+  unsigned long waithint;
+  unsigned long ret;
+  unsigned long waited;
+  char interval_milliseconds[16];
+  char timeout_milliseconds[16];
+  char waited_milliseconds[16];
+  char *function = function_name;
+
+  /* Add brackets to function name. */
+  size_t funclen = strlen(function_name) + 3;
+  char *func = (char *) HeapAlloc(GetProcessHeap(), 0, funclen);
+  if (func) {
+    if (_snprintf_s(func, funclen, _TRUNCATE, "%s()", function_name) > -1) function = func;
+  }
+
+  _snprintf_s(timeout_milliseconds, sizeof(timeout_milliseconds), _TRUNCATE, "%lu", timeout);
+
+  waithint = service_status->dwWaitHint;
+  waited = 0;
+  while (waited < timeout) {
+    interval = timeout - waited;
+    if (interval > NSSM_SHUTDOWN_CHECKPOINT) interval = NSSM_SHUTDOWN_CHECKPOINT;
+
+    service_status->dwCurrentState = SERVICE_STOP_PENDING;
+    service_status->dwWaitHint += interval;
+    service_status->dwCheckPoint++;
+    SetServiceStatus(service_handle, service_status);
+
+    if (waited) {
+      _snprintf_s(waited_milliseconds, sizeof(waited_milliseconds), _TRUNCATE, "%lu", waited);
+      _snprintf_s(interval_milliseconds, sizeof(interval_milliseconds), _TRUNCATE, "%lu", interval);
+      log_event(EVENTLOG_INFORMATION_TYPE, NSSM_EVENT_AWAITING_SHUTDOWN, function, service_name, waited_milliseconds, interval_milliseconds, timeout_milliseconds, 0);
+    }
+
+    switch (WaitForSingleObject(process_handle, interval)) {
+      case WAIT_OBJECT_0:
+        ret = 0;
+        goto awaited;
+
+      case WAIT_TIMEOUT:
+        ret = 1;
+      break;
+
+      default:
+        ret = -1;
+        goto awaited;
+    }
+
+    waited += interval;
+  }
+
+awaited:
+  if (func) HeapFree(GetProcessHeap(), 0, func);
+
+  return ret;
+}
diff --git a/service.h b/service.h
index 2db92fd..49760ba 100644
--- a/service.h
+++ b/service.h
@@ -19,5 +19,6 @@ int start_service();
 int stop_service(unsigned long, bool, bool);
 void CALLBACK end_service(void *, unsigned char);
 void throttle_restart();
+int await_shutdown(char *, char *, SERVICE_STATUS_HANDLE, SERVICE_STATUS *, HANDLE, unsigned long);
 
 #endif
-- 
2.20.1