Allow setting processor affinity.
[nssm.git] / service.cpp
index 621277d..93b7d4a 100644 (file)
@@ -11,6 +11,156 @@ extern settings_t settings[];
 \r
 const TCHAR *exit_action_strings[] = { _T("Restart"), _T("Ignore"), _T("Exit"), _T("Suicide"), 0 };\r
 const TCHAR *startup_strings[] = { _T("SERVICE_AUTO_START"), _T("SERVICE_DELAYED_AUTO_START"), _T("SERVICE_DEMAND_START"), _T("SERVICE_DISABLED"), 0 };\r
+const TCHAR *priority_strings[] = { _T("REALTIME_PRIORITY_CLASS"), _T("HIGH_PRIORITY_CLASS"), _T("ABOVE_NORMAL_PRIORITY_CLASS"), _T("NORMAL_PRIORITY_CLASS"), _T("BELOW_NORMAL_PRIORITY_CLASS"), _T("IDLE_PRIORITY_CLASS"), 0 };\r
+\r
+typedef struct {\r
+  int first;\r
+  int last;\r
+} list_t;\r
+\r
+int affinity_mask_to_string(__int64 mask, TCHAR **string) {\r
+  if (! string) return 1;\r
+  if (! mask) {\r
+    *string = 0;\r
+    return 0;\r
+  }\r
+\r
+  __int64 i, n;\r
+\r
+  /* SetProcessAffinityMask() accepts a mask of up to 64 processors. */\r
+  list_t set[64];\r
+  for (n = 0; n < _countof(set); n++) set[n].first = set[n].last = -1;\r
+\r
+  for (i = 0, n = 0; i < _countof(set); i++) {\r
+    if (mask & (1LL << i)) {\r
+      if (set[n].first == -1) set[n].first = set[n].last = (int) i;\r
+      else if (set[n].last == (int) i - 1) set[n].last = (int) i;\r
+      else {\r
+        n++;\r
+        set[n].first = set[n].last = (int) i;\r
+      }\r
+    }\r
+  }\r
+\r
+  /* Worst case is 2x2 characters for first and last CPU plus - and/or , */\r
+  size_t len = (size_t) (n + 1) * 6;\r
+  *string = (TCHAR *) HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, len * sizeof(TCHAR));\r
+  if (! string) return 2;\r
+\r
+  size_t s = 0;\r
+  int ret;\r
+  for (i = 0; i <= n; i++) {\r
+    if (i) (*string)[s++] = _T(',');\r
+    ret = _sntprintf_s(*string + s, 3, _TRUNCATE, _T("%u"), set[i].first);\r
+    if (ret < 0) {\r
+      HeapFree(GetProcessHeap(), 0, *string);\r
+      *string = 0;\r
+      return 3;\r
+    }\r
+    else s += ret;\r
+    if (set[i].last != set[i].first) {\r
+      ret =_sntprintf_s(*string + s, 4, _TRUNCATE, _T("%c%u"), (set[i].last == set[i].first + 1) ? _T(',') : _T('-'), set[i].last);\r
+      if (ret < 0) {\r
+        HeapFree(GetProcessHeap(), 0, *string);\r
+        *string = 0;\r
+        return 4;\r
+      }\r
+      else s += ret;\r
+    }\r
+  }\r
+\r
+  return 0;\r
+}\r
+\r
+int affinity_string_to_mask(TCHAR *string, __int64 *mask) {\r
+  if (! mask) return 1;\r
+\r
+  *mask = 0LL;\r
+  if (! string) return 0;\r
+\r
+  list_t set[64];\r
+\r
+  TCHAR *s = string;\r
+  TCHAR *end;\r
+  int ret;\r
+  int i;\r
+  int n = 0;\r
+  unsigned long number;\r
+\r
+  for (n = 0; n < _countof(set); n++) set[n].first = set[n].last = -1;\r
+  n = 0;\r
+\r
+  while (*s) {\r
+    ret = str_number(s, &number, &end);\r
+    s = end;\r
+    if (ret == 0 || ret == 2) {\r
+      if (number >= _countof(set)) return 2;\r
+      set[n].first = set[n].last = (int) number;\r
+\r
+      switch (*s) {\r
+        case 0:\r
+          break;\r
+\r
+        case _T(','):\r
+          n++;\r
+          s++;\r
+          break;\r
+\r
+        case _T('-'):\r
+          if (! *(++s)) return 3;\r
+          ret = str_number(s, &number, &end);\r
+          if (ret == 0 || ret == 2) {\r
+            s = end;\r
+            if (! *s || *s == _T(',')) {\r
+              set[n].last = (int) number;\r
+              if (! *s) break;\r
+              n++;\r
+              s++;\r
+            }\r
+            else return 3;\r
+          }\r
+          else return 3;\r
+          break;\r
+\r
+        default:\r
+          return 3;\r
+      }\r
+    }\r
+    else return 4;\r
+  }\r
+\r
+  for (i = 0; i <= n; i++) {\r
+    for (int j = set[i].first; j <= set[i].last; j++) (__int64) *mask |= (1LL << (__int64) j);\r
+  }\r
+\r
+  return 0;\r
+}\r
+\r
+inline unsigned long priority_mask() {\r
+ return REALTIME_PRIORITY_CLASS | HIGH_PRIORITY_CLASS | ABOVE_NORMAL_PRIORITY_CLASS | NORMAL_PRIORITY_CLASS | BELOW_NORMAL_PRIORITY_CLASS | IDLE_PRIORITY_CLASS;\r
+}\r
+\r
+int priority_constant_to_index(unsigned long constant) {\r
+  switch (constant & priority_mask()) {\r
+    case REALTIME_PRIORITY_CLASS: return NSSM_REALTIME_PRIORITY;\r
+    case HIGH_PRIORITY_CLASS: return NSSM_HIGH_PRIORITY;\r
+    case ABOVE_NORMAL_PRIORITY_CLASS: return NSSM_ABOVE_NORMAL_PRIORITY;\r
+    case BELOW_NORMAL_PRIORITY_CLASS: return NSSM_BELOW_NORMAL_PRIORITY;\r
+    case IDLE_PRIORITY_CLASS: return NSSM_IDLE_PRIORITY;\r
+  }\r
+  return NSSM_NORMAL_PRIORITY;\r
+}\r
+\r
+unsigned long priority_index_to_constant(int index) {\r
+  switch (index) {\r
+    case NSSM_REALTIME_PRIORITY: return REALTIME_PRIORITY_CLASS;\r
+    case NSSM_HIGH_PRIORITY: return HIGH_PRIORITY_CLASS;\r
+    case NSSM_ABOVE_NORMAL_PRIORITY: return ABOVE_NORMAL_PRIORITY_CLASS;\r
+    case NSSM_BELOW_NORMAL_PRIORITY: return BELOW_NORMAL_PRIORITY_CLASS;\r
+    case NSSM_IDLE_PRIORITY: return IDLE_PRIORITY_CLASS;\r
+  }\r
+  return NORMAL_PRIORITY_CLASS;\r
+}\r
 \r
 static inline int throttle_milliseconds(unsigned long throttle) {\r
   /* pow() operates on doubles. */\r
@@ -262,6 +412,7 @@ int get_service_username(const TCHAR *service_name, const QUERY_SERVICE_CONFIG *
 }\r
 \r
 int grant_logon_as_service(const TCHAR *username) {\r
+  if (! username) return 0;\r
   if (str_equiv(username, NSSM_LOCALSYSTEM_ACCOUNT)) return 0;\r
 \r
   /* Open Policy object. */\r
@@ -285,7 +436,7 @@ int grant_logon_as_service(const TCHAR *username) {
 #else\r
   size_t buflen;\r
   mbstowcs_s(&buflen, NULL, 0, username, _TRUNCATE);\r
-  lsa_username.MaximumLength = buflen * sizeof(wchar_t);\r
+  lsa_username.MaximumLength = (unsigned short) buflen * sizeof(wchar_t);\r
   lsa_username.Length = lsa_username.MaximumLength - sizeof(wchar_t);\r
   lsa_username.Buffer = (wchar_t *) HeapAlloc(GetProcessHeap(), 0, lsa_username.MaximumLength);\r
   if (lsa_username.Buffer) mbstowcs_s(&buflen, lsa_username.Buffer, lsa_username.MaximumLength, username, _TRUNCATE);\r
@@ -412,6 +563,7 @@ void set_nssm_service_defaults(nssm_service_t *service) {
   if (! service) return;\r
 \r
   service->type = SERVICE_WIN32_OWN_PROCESS;\r
+  service->priority = NORMAL_PRIORITY_CLASS;\r
   service->stdin_sharing = NSSM_STDIN_SHARING;\r
   service->stdin_disposition = NSSM_STDIN_DISPOSITION;\r
   service->stdin_flags = NSSM_STDIN_FLAGS;\r
@@ -781,9 +933,9 @@ int install_service(nssm_service_t *service) {
   GetModuleFileName(0, service->image, _countof(service->image));\r
 \r
   /* Create the service - settings will be changed in edit_service() */\r
-  service->handle = CreateService(services, service->name, service->name, SC_MANAGER_ALL_ACCESS, SERVICE_WIN32_OWN_PROCESS, SERVICE_AUTO_START, SERVICE_ERROR_NORMAL, service->image, 0, 0, 0, 0, 0);\r
+  service->handle = CreateService(services, service->name, service->name, SERVICE_ALL_ACCESS, SERVICE_WIN32_OWN_PROCESS, SERVICE_AUTO_START, SERVICE_ERROR_NORMAL, service->image, 0, 0, 0, 0, 0);\r
   if (! service->handle) {\r
-    print_message(stderr, NSSM_MESSAGE_CREATESERVICE_FAILED);\r
+    print_message(stderr, NSSM_MESSAGE_CREATESERVICE_FAILED, error_string(GetLastError()));\r
     CloseServiceHandle(services);\r
     return 5;\r
   }\r
@@ -907,6 +1059,21 @@ int control_service(unsigned long control, int argc, TCHAR **argv) {
     CloseHandle(service_handle);\r
     CloseServiceHandle(services);\r
 \r
+    if (error == ERROR_IO_PENDING) {\r
+      /*\r
+        Older versions of Windows return immediately with ERROR_IO_PENDING\r
+        indicate that the operation is still in progress.  Newer versions\r
+        will return it if there really is a delay.  As far as we're\r
+        concerned the operation is a success.  We don't claim to offer a\r
+        fully-feature service control method; it's just a quick 'n' dirty\r
+        interface.\r
+\r
+        In the future we may identify and handle this situation properly.\r
+      */\r
+      ret = 1;\r
+      error = ERROR_SUCCESS;\r
+    }\r
+\r
     if (ret) {\r
       _tprintf(_T("%s: %s"), canonical_name, error_string(error));\r
       return 0;\r
@@ -949,6 +1116,11 @@ int control_service(unsigned long control, int argc, TCHAR **argv) {
     CloseHandle(service_handle);\r
     CloseServiceHandle(services);\r
 \r
+    if (error == ERROR_IO_PENDING) {\r
+      ret = 1;\r
+      error = ERROR_SUCCESS;\r
+    }\r
+\r
     if (ret) {\r
       _tprintf(_T("%s: %s"), canonical_name, error_string(error));\r
       return 0;\r
@@ -1238,7 +1410,8 @@ int start_service(nssm_service_t *service) {
 \r
   bool inherit_handles = false;\r
   if (si.dwFlags & STARTF_USESTDHANDLES) inherit_handles = true;\r
-  unsigned long flags = 0;\r
+  unsigned long flags = service->priority & priority_mask();\r
+  if (service->affinity) flags |= CREATE_SUSPENDED;\r
 #ifdef UNICODE\r
   flags |= CREATE_UNICODE_ENVIRONMENT;\r
 #endif\r
@@ -1260,6 +1433,35 @@ int start_service(nssm_service_t *service) {
 \r
   close_output_handles(&si);\r
 \r
+  if (service->affinity) {\r
+    /*\r
+      We are explicitly storing service->affinity as a 64-bit unsigned integer\r
+      so that we can parse it regardless of whether we're running in 32-bit\r
+      or 64-bit mode.  The arguments to SetProcessAffinityMask(), however, are\r
+      defined as type DWORD_PTR and hence limited to 32 bits on a 32-bit system\r
+      (or when running the 32-bit NSSM).\r
+\r
+      The result is a lot of seemingly-unnecessary casting throughout the code\r
+      and potentially confusion when we actually try to start the service.\r
+      Having said that, however, it's unlikely that we're actually going to\r
+      run in 32-bit mode on a system which has more than 32 CPUs so the\r
+      likelihood of seeing a confusing situation is somewhat diminished.\r
+    */\r
+    DWORD_PTR affinity, system_affinity;\r
+\r
+    if (GetProcessAffinityMask(service->process_handle, &affinity, &system_affinity)) affinity = service->affinity & system_affinity;\r
+    else {\r
+      affinity = (DWORD_PTR) service->affinity;\r
+      log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_GETPROCESSAFFINITYMASK_FAILED, service->name, error_string(GetLastError()), 0);\r
+    }\r
+\r
+    if (! SetProcessAffinityMask(service->process_handle, affinity)) {\r
+      log_event(EVENTLOG_WARNING_TYPE, NSSM_EVENT_SETPROCESSAFFINITYMASK_FAILED, service->name, error_string(GetLastError()), 0);\r
+    }\r
+\r
+    ResumeThread(pi.hThread);\r
+  }\r
+\r
   /*\r
     Wait for a clean startup before changing the service status to RUNNING\r
     but be mindful of the fact that we are blocking the service control manager\r