X-Git-Url: http://git.iain.cx/?a=blobdiff_plain;f=service.cpp;h=51c3937598824506c72af4735fa9ef0a6ac522c4;hb=f1d6394333ce0f90c423d74a9860620942f2bb6f;hp=4ec51333f4806ac54cc0ee91c3cab28d26815a2c;hpb=10eb89c84baf36bbb77f6aaf497dcd6f74644bc3;p=nssm.git

diff --git a/service.cpp b/service.cpp
index 4ec5133..51c3937 100644
--- a/service.cpp
+++ b/service.cpp
@@ -11,10 +11,160 @@ extern settings_t settings[];
 
 const TCHAR *exit_action_strings[] = { _T("Restart"), _T("Ignore"), _T("Exit"), _T("Suicide"), 0 };
 const TCHAR *startup_strings[] = { _T("SERVICE_AUTO_START"), _T("SERVICE_DELAYED_AUTO_START"), _T("SERVICE_DEMAND_START"), _T("SERVICE_DISABLED"), 0 };
+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 };
 
-static inline int throttle_milliseconds(unsigned long throttle) {
+typedef struct {
+  int first;
+  int last;
+} list_t;
+
+int affinity_mask_to_string(__int64 mask, TCHAR **string) {
+  if (! string) return 1;
+  if (! mask) {
+    *string = 0;
+    return 0;
+  }
+
+  __int64 i, n;
+
+  /* SetProcessAffinityMask() accepts a mask of up to 64 processors. */
+  list_t set[64];
+  for (n = 0; n < _countof(set); n++) set[n].first = set[n].last = -1;
+
+  for (i = 0, n = 0; i < _countof(set); i++) {
+    if (mask & (1LL << i)) {
+      if (set[n].first == -1) set[n].first = set[n].last = (int) i;
+      else if (set[n].last == (int) i - 1) set[n].last = (int) i;
+      else {
+        n++;
+        set[n].first = set[n].last = (int) i;
+      }
+    }
+  }
+
+  /* Worst case is 2x2 characters for first and last CPU plus - and/or , */
+  size_t len = (size_t) (n + 1) * 6;
+  *string = (TCHAR *) HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, len * sizeof(TCHAR));
+  if (! string) return 2;
+
+  size_t s = 0;
+  int ret;
+  for (i = 0; i <= n; i++) {
+    if (i) (*string)[s++] = _T(',');
+    ret = _sntprintf_s(*string + s, 3, _TRUNCATE, _T("%u"), set[i].first);
+    if (ret < 0) {
+      HeapFree(GetProcessHeap(), 0, *string);
+      *string = 0;
+      return 3;
+    }
+    else s += ret;
+    if (set[i].last != set[i].first) {
+      ret =_sntprintf_s(*string + s, 4, _TRUNCATE, _T("%c%u"), (set[i].last == set[i].first + 1) ? _T(',') : _T('-'), set[i].last);
+      if (ret < 0) {
+        HeapFree(GetProcessHeap(), 0, *string);
+        *string = 0;
+        return 4;
+      }
+      else s += ret;
+    }
+  }
+
+  return 0;
+}
+
+int affinity_string_to_mask(TCHAR *string, __int64 *mask) {
+  if (! mask) return 1;
+
+  *mask = 0LL;
+  if (! string) return 0;
+
+  list_t set[64];
+
+  TCHAR *s = string;
+  TCHAR *end;
+  int ret;
+  int i;
+  int n = 0;
+  unsigned long number;
+
+  for (n = 0; n < _countof(set); n++) set[n].first = set[n].last = -1;
+  n = 0;
+
+  while (*s) {
+    ret = str_number(s, &number, &end);
+    s = end;
+    if (ret == 0 || ret == 2) {
+      if (number >= _countof(set)) return 2;
+      set[n].first = set[n].last = (int) number;
+
+      switch (*s) {
+        case 0:
+          break;
+
+        case _T(','):
+          n++;
+          s++;
+          break;
+
+        case _T('-'):
+          if (! *(++s)) return 3;
+          ret = str_number(s, &number, &end);
+          if (ret == 0 || ret == 2) {
+            s = end;
+            if (! *s || *s == _T(',')) {
+              set[n].last = (int) number;
+              if (! *s) break;
+              n++;
+              s++;
+            }
+            else return 3;
+          }
+          else return 3;
+          break;
+
+        default:
+          return 3;
+      }
+    }
+    else return 4;
+  }
+
+  for (i = 0; i <= n; i++) {
+    for (int j = set[i].first; j <= set[i].last; j++) (__int64) *mask |= (1LL << (__int64) j);
+  }
+
+  return 0;
+}
+
+inline unsigned long priority_mask() {
+ return REALTIME_PRIORITY_CLASS | HIGH_PRIORITY_CLASS | ABOVE_NORMAL_PRIORITY_CLASS | NORMAL_PRIORITY_CLASS | BELOW_NORMAL_PRIORITY_CLASS | IDLE_PRIORITY_CLASS;
+}
+
+int priority_constant_to_index(unsigned long constant) {
+  switch (constant & priority_mask()) {
+    case REALTIME_PRIORITY_CLASS: return NSSM_REALTIME_PRIORITY;
+    case HIGH_PRIORITY_CLASS: return NSSM_HIGH_PRIORITY;
+    case ABOVE_NORMAL_PRIORITY_CLASS: return NSSM_ABOVE_NORMAL_PRIORITY;
+    case BELOW_NORMAL_PRIORITY_CLASS: return NSSM_BELOW_NORMAL_PRIORITY;
+    case IDLE_PRIORITY_CLASS: return NSSM_IDLE_PRIORITY;
+  }
+  return NSSM_NORMAL_PRIORITY;
+}
+
+unsigned long priority_index_to_constant(int index) {
+  switch (index) {
+    case NSSM_REALTIME_PRIORITY: return REALTIME_PRIORITY_CLASS;
+    case NSSM_HIGH_PRIORITY: return HIGH_PRIORITY_CLASS;
+    case NSSM_ABOVE_NORMAL_PRIORITY: return ABOVE_NORMAL_PRIORITY_CLASS;
+    case NSSM_BELOW_NORMAL_PRIORITY: return BELOW_NORMAL_PRIORITY_CLASS;
+    case NSSM_IDLE_PRIORITY: return IDLE_PRIORITY_CLASS;
+  }
+  return NORMAL_PRIORITY_CLASS;
+}
+
+static inline unsigned long throttle_milliseconds(unsigned long throttle) {
   /* pow() operates on doubles. */
-  int ret = 1; for (unsigned long i = 1; i < throttle; i++) ret *= 2;
+  unsigned long ret = 1; for (unsigned long i = 1; i < throttle; i++) ret *= 2;
   return ret * 1000;
 }
 
@@ -413,6 +563,7 @@ void set_nssm_service_defaults(nssm_service_t *service) {
   if (! service) return;
 
   service->type = SERVICE_WIN32_OWN_PROCESS;
+  service->priority = NORMAL_PRIORITY_CLASS;
   service->stdin_sharing = NSSM_STDIN_SHARING;
   service->stdin_disposition = NSSM_STDIN_DISPOSITION;
   service->stdin_flags = NSSM_STDIN_FLAGS;
@@ -550,14 +701,20 @@ int pre_edit_service(int argc, TCHAR **argv) {
       for (i = 0; settings[i].name; i++) _ftprintf(stderr, _T("%s\n"), settings[i].name);
       return 1;
     }
-    if (argc < mandatory) return usage(1);
 
     additional = 0;
     if (additional_mandatory) {
+      if (argc < mandatory) {
+        print_message(stderr, NSSM_MESSAGE_MISSING_SUBPARAMETER, parameter);
+        return 1;
+      }
       additional = argv[3];
       remainder = 4;
     }
-    else additional = argv[remainder];
+    else {
+      additional = argv[remainder];
+      if (argc < mandatory) return usage(1);
+    }
   }
 
   nssm_service_t *service = alloc_nssm_service();
@@ -902,7 +1059,7 @@ int control_service(unsigned long control, int argc, TCHAR **argv) {
   int ret;
   unsigned long error;
   SERVICE_STATUS service_status;
-  if (control == 0) {
+  if (control == NSSM_SERVICE_CONTROL_START) {
     ret = StartService(service_handle, (unsigned long) argc, (const TCHAR **) argv);
     error = GetLastError();
     CloseHandle(service_handle);
@@ -1124,12 +1281,13 @@ int monitor_service(nssm_service_t *service) {
 TCHAR *service_control_text(unsigned long control) {
   switch (control) {
     /* HACK: there is no SERVICE_CONTROL_START constant */
-    case 0: return _T("START");
+    case NSSM_SERVICE_CONTROL_START: return _T("START");
     case SERVICE_CONTROL_STOP: return _T("STOP");
     case SERVICE_CONTROL_SHUTDOWN: return _T("SHUTDOWN");
     case SERVICE_CONTROL_PAUSE: return _T("PAUSE");
     case SERVICE_CONTROL_CONTINUE: return _T("CONTINUE");
     case SERVICE_CONTROL_INTERROGATE: return _T("INTERROGATE");
+    case NSSM_SERVICE_CONTROL_ROTATE: return _T("ROTATE");
     default: return 0;
   }
 }
@@ -1217,6 +1375,12 @@ unsigned long WINAPI service_control_handler(unsigned long control, unsigned lon
       */
       log_service_control(service->name, control, false);
       return ERROR_CALL_NOT_IMPLEMENTED;
+
+    case NSSM_SERVICE_CONTROL_ROTATE:
+      log_service_control(service->name, control, true);
+      if (service->rotate_stdout_online) service->rotate_stdout_online = NSSM_ROTATE_ONLINE_ASAP;
+      if (service->rotate_stdout_online) service->rotate_stderr_online = NSSM_ROTATE_ONLINE_ASAP;
+      return NO_ERROR;
   }
 
   /* Unknown control */
@@ -1259,7 +1423,8 @@ int start_service(nssm_service_t *service) {
 
   bool inherit_handles = false;
   if (si.dwFlags & STARTF_USESTDHANDLES) inherit_handles = true;
-  unsigned long flags = 0;
+  unsigned long flags = service->priority & priority_mask();
+  if (service->affinity) flags |= CREATE_SUSPENDED;
 #ifdef UNICODE
   flags |= CREATE_UNICODE_ENVIRONMENT;
 #endif
@@ -1279,7 +1444,36 @@ int start_service(nssm_service_t *service) {
 
   if (get_process_creation_time(service->process_handle, &service->creation_time)) ZeroMemory(&service->creation_time, sizeof(service->creation_time));
 
-  close_output_handles(&si);
+  close_output_handles(&si, ! service->rotate_stdout_online, ! service->rotate_stderr_online);
+
+  if (service->affinity) {
+    /*
+      We are explicitly storing service->affinity as a 64-bit unsigned integer
+      so that we can parse it regardless of whether we're running in 32-bit
+      or 64-bit mode.  The arguments to SetProcessAffinityMask(), however, are
+      defined as type DWORD_PTR and hence limited to 32 bits on a 32-bit system
+      (or when running the 32-bit NSSM).
+
+      The result is a lot of seemingly-unnecessary casting throughout the code
+      and potentially confusion when we actually try to start the service.
+      Having said that, however, it's unlikely that we're actually going to
+      run in 32-bit mode on a system which has more than 32 CPUs so the
+      likelihood of seeing a confusing situation is somewhat diminished.
+    */
+    DWORD_PTR affinity, system_affinity;
+
+    if (GetProcessAffinityMask(service->process_handle, &affinity, &system_affinity)) affinity = service->affinity & system_affinity;
+    else {
+      affinity = (DWORD_PTR) service->affinity;
+      log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_GETPROCESSAFFINITYMASK_FAILED, service->name, error_string(GetLastError()), 0);
+    }
+
+    if (! SetProcessAffinityMask(service->process_handle, affinity)) {
+      log_event(EVENTLOG_WARNING_TYPE, NSSM_EVENT_SETPROCESSAFFINITYMASK_FAILED, service->name, error_string(GetLastError()), 0);
+    }
+
+    ResumeThread(pi.hThread);
+  }
 
   /*
     Wait for a clean startup before changing the service status to RUNNING
@@ -1309,6 +1503,9 @@ int start_service(nssm_service_t *service) {
     else service->throttle = 0;
   }
 
+  /* Ensure the restart delay is always applied. */
+  if (service->restart_delay && ! service->throttle) service->throttle++;
+
   return 0;
 }
 
@@ -1451,14 +1648,22 @@ void throttle_restart(nssm_service_t *service) {
   /* This can't be a restart if the service is already running. */
   if (! service->throttle++) return;
 
-  int ms = throttle_milliseconds(service->throttle);
+  unsigned long ms;
+  unsigned long throttle_ms = throttle_milliseconds(service->throttle);
+  TCHAR threshold[8], milliseconds[8];
+
+  if (service->restart_delay > throttle_ms) ms = service->restart_delay;
+  else ms = throttle_ms;
 
   if (service->throttle > 7) service->throttle = 8;
 
-  TCHAR threshold[8], milliseconds[8];
-  _sntprintf_s(threshold, _countof(threshold), _TRUNCATE, _T("%lu"), service->throttle_delay);
   _sntprintf_s(milliseconds, _countof(milliseconds), _TRUNCATE, _T("%lu"), ms);
-  log_event(EVENTLOG_WARNING_TYPE, NSSM_EVENT_THROTTLED, service->name, threshold, milliseconds, 0);
+
+  if (service->throttle == 1 && service->restart_delay > throttle_ms) log_event(EVENTLOG_INFORMATION_TYPE, NSSM_EVENT_RESTART_DELAY, service->name, milliseconds, 0);
+  else {
+    _sntprintf_s(threshold, _countof(threshold), _TRUNCATE, _T("%lu"), service->throttle_delay);
+    log_event(EVENTLOG_WARNING_TYPE, NSSM_EVENT_THROTTLED, service->name, threshold, milliseconds, 0);
+  }
 
   if (use_critical_section) EnterCriticalSection(&service->throttle_section);
   else if (service->throttle_timer) {