Browse Source

Allow setting processor affinity.

The REG_SZ AppPriority entry specifies a list of processors which will
be converted to a mask and passed to SetProcessAffinityMask().

Processors are numbered from 0, must not exceed 63, and are separated
with commas or dashes, to specify a range.

Thanks Robert Middleton.
Iain Patterson 8 years ago
parent
commit
53371f115d
13 changed files with 428 additions and 10 deletions
  1. 2 2
      ChangeLog.txt
  2. 31 3
      README.txt
  3. 76 0
      gui.cpp
  4. BIN
      messages.mc
  5. 7 0
      nssm.cpp
  6. 1 0
      nssm.h
  7. BIN
      nssm.rc
  8. 38 0
      registry.cpp
  9. 1 0
      registry.h
  10. 3 1
      resource.h
  11. 153 0
      service.cpp
  12. 3 0
      service.h
  13. 113 4
      settings.cpp

+ 2 - 2
ChangeLog.txt

@@ -3,8 +3,8 @@ Changes since 2.21
   * Existing services can now be managed using the GUI
     or on the command line.
 
-  * NSSM can now set the priority class of the managed
-    application.
+  * NSSM can now set the priority class and processor
+    affinity of the managed application.
 
   * NSSM can now optionally rotate existing files when
     redirecting I/O.

+ 31 - 3
README.txt

@@ -53,7 +53,8 @@ Since version 2.19, NSSM can add to the service's environment by setting
 AppEnvironmentExtra in place of or in addition to the srvany-compatible
 AppEnvironment.
 
-Since version 2.22, NSSM can set the managed application's process priority.
+Since version 2.22, NSSM can set the managed application's process priority
+and CPU affinity.
 
 Since version 2.22, NSSM can rotate existing output files when redirecting I/O.
 
@@ -171,6 +172,25 @@ SetPriorityClass().  If AppPriority() is missing or invalid the
 application will be launched with normal priority.
 
 
+Processor affinity
+------------------
+NSSM can set the CPU affinity of the managed application.  NSSM will look in
+the registry under HKLM\SYSTEM\CurrentControlSet\Services\<service>\Parameters
+for the REG_SZ entry AppAffinity.   It should specify a comma-separated listed
+of zero-indexed processor IDs.  A range of processors may optionally be
+specified with a dash.  No other characters are allowed in the string.
+
+For example, to specify the first; second; third and fifth CPUs, an appropriate
+AppAffinity would be 0-2,4.
+
+If AppAffinity is missing or invalid, NSSM will not attempt to restrict the
+application to specific CPUs.
+
+Note that the 64-bit version of NSSM can configure a maximum of 64 CPUs in this
+way and that the 32-bit version can configure a maxium of 32 CPUs even when
+running on 64-bit Windows.
+
+
 Stopping the service
 --------------------
 When stopping a service NSSM will attempt several different methods of killing
@@ -511,6 +531,10 @@ To configure the server to log to a file:
 
     nssm set UT2004 AppStdout c:\games\ut2004\service.log
 
+To restrict the server to a single CPU:
+
+    nssm set UT2004 AppAffinity 0
+
 To remove the server:
 
     nssm remove UT2004 confirm
@@ -542,7 +566,8 @@ Thanks to Joel Reingold for spotting a command line truncation bug.
 Thanks to Arve Knudsen for spotting that child processes of the monitored
 application could be left running on service shutdown, and that a missing
 registry value for AppDirectory confused NSSM.
-Thanks to Peter Wagemans and Laszlo Keresztfalvi for suggesting throttling restarts.
+Thanks to Peter Wagemans and Laszlo Keresztfalvi for suggesting throttling
+restarts.
 Thanks to Eugene Lifshitz for finding an edge case in CreateProcess() and for
 advising how to build messages.mc correctly in paths containing spaces.
 Thanks to Rob Sharp for pointing out that NSSM did not respect the
@@ -554,12 +579,15 @@ the default language when the user's display language was not translated.
 Thanks to Riccardo Gusmeroli for Italian translation.
 Thanks to Eric Cheldelin for the inspiration to generate a Control-C event
 on shutdown.
-Thanks to Brian Baxter for suggesting how to escape quotes from the command prompt.
+Thanks to Brian Baxter for suggesting how to escape quotes from the command
+prompt.
 Thanks to Russ Holmann for suggesting that the shutdown timeout be configurable.
 Thanks to Paul Spause for spotting a bug with default registry entries.
 Thanks to BUGHUNTER for spotting more GUI bugs.
 Thanks to Doug Watson for suggesting file rotation.
 Thanks to Арслан Сайдуганов for suggesting setting process priority.
+Thanks to Robert Middleton for suggestion and draft implementation of process
+affinity support.
 
 Licence
 -------

+ 76 - 0
gui.cpp

@@ -67,6 +67,7 @@ int nssm_gui(int resource, nssm_service_t *service) {
 
     /* Set existing details. */
     HWND combo;
+    HWND list;
 
     /* Application tab. */
     if (service->native) SetDlgItemText(tablist[NSSM_TAB_APPLICATION], IDC_PATH, service->image);
@@ -101,6 +102,21 @@ int nssm_gui(int resource, nssm_service_t *service) {
       SendMessage(combo, CB_SETCURSEL, priority, 0);
     }
 
+    if (service->affinity) {
+      list = GetDlgItem(tablist[NSSM_TAB_PROCESS], IDC_AFFINITY);
+      SendDlgItemMessage(tablist[NSSM_TAB_PROCESS], IDC_AFFINITY_ALL, BM_SETCHECK, BST_UNCHECKED, 0);
+      EnableWindow(GetDlgItem(tablist[NSSM_TAB_PROCESS], IDC_AFFINITY), 1);
+
+      DWORD_PTR affinity, system_affinity;
+      if (GetProcessAffinityMask(GetCurrentProcess(), &affinity, &system_affinity)) {
+        if ((service->affinity & (__int64) system_affinity) != service->affinity) popup_message(dlg, MB_OK | MB_ICONWARNING, NSSM_GUI_WARN_AFFINITY);
+      }
+
+      for (int i = 0; i < num_cpus(); i++) {
+        if (! (service->affinity & (1LL << (__int64) i))) SendMessage(list, LB_SETSEL, 0, i);
+      }
+    }
+
     /* Shutdown tab. */
     if (! (service->stop_method & NSSM_STOP_METHOD_CONSOLE)) {
       SendDlgItemMessage(tablist[NSSM_TAB_SHUTDOWN], IDC_METHOD_CONSOLE, BM_SETCHECK, BST_UNCHECKED, 0);
@@ -230,6 +246,10 @@ static inline void set_logon_enabled(unsigned char enabled) {
   EnableWindow(GetDlgItem(tablist[NSSM_TAB_LOGON], IDC_PASSWORD2), enabled);
 }
 
+static inline void set_affinity_enabled(unsigned char enabled) {
+  EnableWindow(GetDlgItem(tablist[NSSM_TAB_PROCESS], IDC_AFFINITY), enabled);
+}
+
 static inline void set_rotation_enabled(unsigned char enabled) {
   EnableWindow(GetDlgItem(tablist[NSSM_TAB_ROTATION], IDC_ROTATE_SECONDS), enabled);
   EnableWindow(GetDlgItem(tablist[NSSM_TAB_ROTATION], IDC_ROTATE_BYTES_LOW), enabled);
@@ -438,6 +458,22 @@ int configure(HWND window, nssm_service_t *service, nssm_service_t *orig_service
   combo = GetDlgItem(tablist[NSSM_TAB_PROCESS], IDC_PRIORITY);
   service->priority = priority_index_to_constant((unsigned long) SendMessage(combo, CB_GETCURSEL, 0, 0));
 
+  service->affinity = 0LL;
+  if (! (SendDlgItemMessage(tablist[NSSM_TAB_PROCESS], IDC_AFFINITY_ALL, BM_GETCHECK, 0, 0) & BST_CHECKED)) {
+    HWND list = GetDlgItem(tablist[NSSM_TAB_PROCESS], IDC_AFFINITY);
+    int selected = (int) SendMessage(list, LB_GETSELCOUNT, 0, 0);
+    int count = (int) SendMessage(list, LB_GETCOUNT, 0, 0);
+    if (! selected) {
+      popup_message(window, MB_OK | MB_ICONEXCLAMATION, NSSM_GUI_WARN_AFFINITY_NONE);
+      return 5;
+    }
+    else if (selected < count) {
+      for (int i = 0; i < count; i++) {
+        if (SendMessage(list, LB_GETSEL, i, 0)) service->affinity |= (1LL << (__int64) i);
+      }
+    }
+  }
+
   /* Get stop method stuff. */
   check_stop_method(service, NSSM_STOP_METHOD_CONSOLE, IDC_METHOD_CONSOLE);
   check_stop_method(service, NSSM_STOP_METHOD_WINDOW, IDC_METHOD_WINDOW);
@@ -766,6 +802,13 @@ INT_PTR CALLBACK tab_dlg(HWND tab, UINT message, WPARAM w, LPARAM l) {
           set_logon_enabled(1);
           break;
 
+        /* Affinity. */
+        case IDC_AFFINITY_ALL:
+          if (SendDlgItemMessage(tab, LOWORD(w), BM_GETCHECK, 0, 0) & BST_CHECKED) enabled = 0;
+          else enabled = 1;
+          set_affinity_enabled(enabled);
+          break;
+
         /* Shutdown methods. */
         case IDC_METHOD_CONSOLE:
           set_timeout_enabled(LOWORD(w), IDC_KILL_CONSOLE);
@@ -832,6 +875,8 @@ INT_PTR CALLBACK nssm_dlg(HWND window, UINT message, WPARAM w, LPARAM l) {
 
       HWND tabs;
       HWND combo;
+      HWND list;
+      int i, n;
       tabs = GetDlgItem(window, IDC_TAB1);
       if (! tabs) return 0;
 
@@ -901,6 +946,37 @@ INT_PTR CALLBACK nssm_dlg(HWND window, UINT message, WPARAM w, LPARAM l) {
       SendMessage(combo, CB_INSERTSTRING, NSSM_IDLE_PRIORITY, (LPARAM) message_string(NSSM_GUI_IDLE_PRIORITY_CLASS));
       SendMessage(combo, CB_SETCURSEL, NSSM_NORMAL_PRIORITY, 0);
 
+      list = GetDlgItem(tablist[NSSM_TAB_PROCESS], IDC_AFFINITY);
+      n = num_cpus();
+      SendMessage(list, LB_SETCOLUMNWIDTH, 16, 0);
+      for (i = 0; i < n; i++) {
+        TCHAR buffer[3];
+        _sntprintf_s(buffer, _countof(buffer), _TRUNCATE, _T("%d"), i);
+        SendMessage(list, LB_ADDSTRING, 0, (LPARAM) buffer);
+      }
+
+      /*
+        Size to fit.
+        The box is high enough for four rows.  It is wide enough for eight
+        columns without scrolling.  With scrollbars it shrinks to two rows.
+        Note that the above only holds if we set the column width BEFORE
+        adding the strings.
+      */
+      if (n < 32) {
+        int columns = (n - 1) / 4;
+        RECT rect;
+        GetWindowRect(list, &rect);
+        int width = rect.right - rect.left;
+        width -= (7 - columns) * 16;
+        int height = rect.bottom - rect.top;
+        if (n < 4) height -= SendMessage(list, LB_GETITEMHEIGHT, 0, 0) * (4 - n);
+        SetWindowPos(list, 0, 0, 0, width, height, SWP_NOMOVE | SWP_NOOWNERZORDER);
+      }
+      SendMessage(list, LB_SELITEMRANGE, 1, MAKELPARAM(0, n));
+
+      SendDlgItemMessage(tablist[NSSM_TAB_PROCESS], IDC_AFFINITY_ALL, BM_SETCHECK, BST_CHECKED, 0);
+      set_affinity_enabled(0);
+
       /* Shutdown tab. */
       tab.pszText = message_string(NSSM_GUI_TAB_SHUTDOWN);
       tab.cchTextMax = (int) _tcslen(tab.pszText);

BIN
messages.mc


+ 7 - 0
nssm.cpp

@@ -75,6 +75,13 @@ static void check_console() {
   FreeConsole();
 }
 
+int num_cpus() {
+  DWORD_PTR i, affinity, system_affinity;
+  if (! GetProcessAffinityMask(GetCurrentProcess(), &affinity, &system_affinity)) return 64;
+  for (i = 0; system_affinity & (1LL << i); i++);
+  return (int) i;
+}
+
 int _tmain(int argc, TCHAR **argv) {
   check_console();
 

+ 1 - 0
nssm.h

@@ -21,6 +21,7 @@ int str_equiv(const TCHAR *, const TCHAR *);
 void strip_basename(TCHAR *);
 int str_number(const TCHAR *, unsigned long *, TCHAR **);
 int str_number(const TCHAR *, unsigned long *);
+int num_cpus();
 int usage(int);
 
 #define NSSM _T("NSSM")

BIN
nssm.rc


+ 38 - 0
registry.cpp

@@ -53,6 +53,18 @@ int create_parameters(nssm_service_t *service, bool editing) {
   /* Other non-default parameters. May fail. */
   if (service->priority != NORMAL_PRIORITY_CLASS) set_number(key, NSSM_REG_PRIORITY, service->priority);
   else if (editing) RegDeleteValue(key, NSSM_REG_PRIORITY);
+  if (service->affinity) {
+    TCHAR *string;
+    if (! affinity_mask_to_string(service->affinity, &string)) {
+      if (RegSetValueEx(key, NSSM_REG_AFFINITY, 0, REG_SZ, (const unsigned char *) string, (unsigned long) (_tcslen(string) + 1) * sizeof(TCHAR)) != ERROR_SUCCESS) {
+        log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_SETVALUE_FAILED, NSSM_REG_AFFINITY, error_string(GetLastError()), 0);
+        HeapFree(GetProcessHeap(), 0, string);
+        return 5;
+      }
+    }
+    if (string) HeapFree(GetProcessHeap(), 0, string);
+  }
+  else if (editing) RegDeleteValue(key, NSSM_REG_AFFINITY);
   unsigned long stop_method_skip = ~service->stop_method;
   if (stop_method_skip) set_number(key, NSSM_REG_STOP_METHOD_SKIP, stop_method_skip);
   else if (editing) RegDeleteValue(key, NSSM_REG_STOP_METHOD_SKIP);
@@ -453,6 +465,32 @@ int get_parameters(nssm_service_t *service, STARTUPINFO *si) {
     log_event(EVENTLOG_WARNING_TYPE, NSSM_EVENT_NO_DIR, NSSM_REG_DIR, service->name, service->dir, 0);
   }
 
+  /* Try to get processor affinity - may fail. */
+  TCHAR buffer[512];
+  if (expand_parameter(key, NSSM_REG_AFFINITY, buffer, sizeof(buffer), false, false) || ! buffer[0]) service->affinity = 0LL;
+  else if (affinity_string_to_mask(buffer, &service->affinity)) {
+    log_event(EVENTLOG_WARNING_TYPE, NSSM_EVENT_BOGUS_AFFINITY_MASK, service->name, buffer);
+    service->affinity = 0LL;
+  }
+  else {
+    DWORD_PTR affinity, system_affinity;
+
+    if (GetProcessAffinityMask(GetCurrentProcess(), &affinity, &system_affinity)) {
+      _int64 effective_affinity = service->affinity & system_affinity;
+      if (effective_affinity != service->affinity) {
+        TCHAR *system = 0;
+        if (! affinity_mask_to_string(system_affinity, &system)) {
+          TCHAR *effective = 0;
+          if (! affinity_mask_to_string(effective_affinity, &effective)) {
+            log_event(EVENTLOG_WARNING_TYPE, NSSM_EVENT_EFFECTIVE_AFFINITY_MASK, service->name, buffer, system, effective, 0);
+          }
+          HeapFree(GetProcessHeap(), 0, effective);
+        }
+        HeapFree(GetProcessHeap(), 0, system);
+      }
+    }
+  }
+
   /* Try to get environment variables - may fail */
   set_environment(service->name, key, NSSM_REG_ENV, &service->env, &service->envlen);
   /* Environment variables to add to existing rather than replace - may fail. */

+ 1 - 0
registry.h

@@ -24,6 +24,7 @@
 #define NSSM_REG_ROTATE_BYTES_LOW _T("AppRotateBytes")
 #define NSSM_REG_ROTATE_BYTES_HIGH _T("AppRotateBytesHigh")
 #define NSSM_REG_PRIORITY _T("AppPriority")
+#define NSSM_REG_AFFINITY _T("AppAffinity")
 #define NSSM_STDIO_LENGTH 29
 
 HKEY open_registry(const TCHAR *, const TCHAR *, REGSAM sam);

+ 3 - 1
resource.h

@@ -57,6 +57,8 @@
 #define IDC_PASSWORD1                   1038
 #define IDC_PASSWORD2                   1039
 #define IDC_PRIORITY                    1040
+#define IDC_AFFINITY_ALL                1041
+#define IDC_AFFINITY                    1042
 
 // Next default values for new objects
 // 
@@ -64,7 +66,7 @@
 #ifndef APSTUDIO_READONLY_SYMBOLS
 #define _APS_NEXT_RESOURCE_VALUE        115
 #define _APS_NEXT_COMMAND_VALUE         40001
-#define _APS_NEXT_CONTROL_VALUE         1041
+#define _APS_NEXT_CONTROL_VALUE         1043
 #define _APS_NEXT_SYMED_VALUE           101
 #endif
 #endif

+ 153 - 0
service.cpp

@@ -13,6 +13,129 @@ const TCHAR *exit_action_strings[] = { _T("Restart"), _T("Ignore"), _T("Exit"),
 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 };
 
+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;
 }
@@ -1288,6 +1411,7 @@ int start_service(nssm_service_t *service) {
   bool inherit_handles = false;
   if (si.dwFlags & STARTF_USESTDHANDLES) inherit_handles = true;
   unsigned long flags = service->priority & priority_mask();
+  if (service->affinity) flags |= CREATE_SUSPENDED;
 #ifdef UNICODE
   flags |= CREATE_UNICODE_ENVIRONMENT;
 #endif
@@ -1309,6 +1433,35 @@ int start_service(nssm_service_t *service) {
 
   close_output_handles(&si);
 
+  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
     but be mindful of the fact that we are blocking the service control manager

+ 3 - 0
service.h

@@ -44,6 +44,7 @@ typedef struct {
   TCHAR flags[VALUE_LENGTH];
   TCHAR dir[MAX_PATH];
   TCHAR *env;
+  __int64 affinity;
   unsigned long envlen;
   TCHAR *env_extra;
   unsigned long env_extralen;
@@ -93,6 +94,8 @@ TCHAR *service_control_text(unsigned long);
 void log_service_control(TCHAR *, unsigned long, bool);
 unsigned long WINAPI service_control_handler(unsigned long, unsigned long, void *, void *);
 
+int affinity_mask_to_string(__int64, TCHAR **);
+int affinity_string_to_mask(TCHAR *, __int64 *);
 unsigned long priority_mask();
 int priority_constant_to_index(unsigned long);
 unsigned long priority_index_to_constant(int);

+ 113 - 4
settings.cpp

@@ -1,12 +1,15 @@
 #include "nssm.h"
 /* XXX: (value && value->string) is probably bogus because value is probably never null */
 
+/* Affinity. */
+#define NSSM_AFFINITY_ALL _T("All")
+
 extern const TCHAR *exit_action_strings[];
 extern const TCHAR *startup_strings[];
 extern const TCHAR *priority_strings[];
 
-/* Does the parameter refer to the default value of the AppExit setting? */
-static inline int is_default_exit_action(const TCHAR *value) {
+/* Does the parameter refer to the default value of the setting? */
+static inline int is_default(const TCHAR *value) {
   return (str_equiv(value, _T("default")) || str_equiv(value, _T("*")) || ! value[0]);
 }
 
@@ -110,7 +113,7 @@ static int setting_set_exit_action(const TCHAR *service_name, void *param, const
 
   if (additional) {
     /* Default action? */
-    if (is_default_exit_action(additional)) code = 0;
+    if (is_default(additional)) code = 0;
     else {
       if (str_number(additional, &exitcode)) return -1;
       code = (TCHAR *) additional;
@@ -167,7 +170,7 @@ static int setting_get_exit_action(const TCHAR *service_name, void *param, const
   unsigned long *code = 0;
 
   if (additional) {
-    if (! is_default_exit_action(additional)) {
+    if (! is_default(additional)) {
       if (str_number(additional, &exitcode)) return -1;
       code = &exitcode;
     }
@@ -183,6 +186,111 @@ static int setting_get_exit_action(const TCHAR *service_name, void *param, const
   return 1;
 }
 
+static int setting_set_affinity(const TCHAR *service_name, void *param, const TCHAR *name, void *default_value, value_t *value, const TCHAR *additional) {
+  HKEY key = (HKEY) param;
+  if (! key) return -1;
+
+  long error;
+  __int64 mask;
+  __int64 system_affinity = 0LL;
+
+  if (value && value->string) {
+    DWORD_PTR affinity;
+    if (! GetProcessAffinityMask(GetCurrentProcess(), &affinity, (DWORD_PTR *) &system_affinity)) system_affinity = ~0;
+
+    if (is_default(value->string) || str_equiv(value->string, NSSM_AFFINITY_ALL)) mask = 0LL;
+    else if (affinity_string_to_mask(value->string, &mask)) {
+      print_message(stderr, NSSM_MESSAGE_BOGUS_AFFINITY_MASK, value->string, num_cpus() - 1);
+      return -1;
+    }
+  }
+  else mask = 0LL;
+
+  if (! mask) {
+    error = RegDeleteValue(key, name);
+    if (error == ERROR_SUCCESS || error == ERROR_FILE_NOT_FOUND) return 0;
+    print_message(stderr, NSSM_MESSAGE_REGDELETEVALUE_FAILED, name, service_name, error_string(error));
+    return -1;
+  }
+
+  /* Canonicalise. */
+  TCHAR *canon = 0;
+  if (affinity_mask_to_string(mask, &canon)) canon = value->string;
+
+  __int64 effective_affinity = mask & system_affinity;
+  if (effective_affinity != mask) {
+    /* Requested CPUs did not intersect with available CPUs? */
+    if (! effective_affinity) mask = effective_affinity = system_affinity;
+
+    TCHAR *system = 0;
+    if (! affinity_mask_to_string(system_affinity, &system)) {
+      TCHAR *effective = 0;
+      if (! affinity_mask_to_string(effective_affinity, &effective)) {
+        print_message(stderr, NSSM_MESSAGE_EFFECTIVE_AFFINITY_MASK, value->string, system, effective);
+        HeapFree(GetProcessHeap(), 0, effective);
+      }
+      HeapFree(GetProcessHeap(), 0, system);
+    }
+  }
+
+  if (RegSetValueEx(key, name, 0, REG_SZ, (const unsigned char *) canon, (unsigned long) (_tcslen(canon) + 1) * sizeof(TCHAR)) != ERROR_SUCCESS) {
+    if (canon != value->string) HeapFree(GetProcessHeap(), 0, canon);
+    log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_SETVALUE_FAILED, name, error_string(GetLastError()), 0);
+    return -1;
+  }
+
+  if (canon != value->string) HeapFree(GetProcessHeap(), 0, canon);
+  return 1;
+}
+
+static int setting_get_affinity(const TCHAR *service_name, void *param, const TCHAR *name, void *default_value, value_t *value, const TCHAR *additional) {
+  HKEY key = (HKEY) param;
+  if (! key) return -1;
+
+  unsigned long type;
+  TCHAR *buffer = 0;
+  unsigned long buflen = 0;
+
+  int ret = RegQueryValueEx(key, name, 0, &type, 0, &buflen);
+  if (ret == ERROR_FILE_NOT_FOUND) {
+    if (value_from_string(name, value, NSSM_AFFINITY_ALL) == 1) return 0;
+    return -1;
+  }
+  if (ret != ERROR_SUCCESS) return -1;
+
+  if (type != REG_SZ) return -1;
+
+  buffer = (TCHAR *) HeapAlloc(GetProcessHeap(), 0, buflen);
+  if (! buffer) {
+    print_message(stderr, NSSM_MESSAGE_OUT_OF_MEMORY, _T("affinity"), _T("setting_get_affinity"));
+    return -1;
+  }
+
+  if (expand_parameter(key, (TCHAR *) name, buffer, buflen, false, true)) {
+    HeapFree(GetProcessHeap(), 0, buffer);
+    return -1;
+  }
+
+  __int64 affinity;
+  if (affinity_string_to_mask(buffer, &affinity)) {
+    print_message(stderr, NSSM_MESSAGE_BOGUS_AFFINITY_MASK, buffer, num_cpus() - 1);
+    HeapFree(GetProcessHeap(), 0, buffer);
+    return -1;
+  }
+
+  HeapFree(GetProcessHeap(), 0, buffer);
+
+  /* Canonicalise. */
+  if (affinity_mask_to_string(affinity, &buffer)) {
+    if (buffer) HeapFree(GetProcessHeap(), 0, buffer);
+    return -1;
+  }
+
+  ret = value_from_string(name, value, buffer);
+  HeapFree(GetProcessHeap(), 0, buffer);
+  return ret;
+}
+
 static int setting_set_environment(const TCHAR *service_name, void *param, const TCHAR *name, void *default_value, value_t *value, const TCHAR *additional) {
   HKEY key = (HKEY) param;
   if (! param) return -1;
@@ -699,6 +807,7 @@ settings_t settings[] = {
   { NSSM_REG_FLAGS, REG_EXPAND_SZ, (void *) _T(""), false, 0, setting_set_string, setting_get_string },
   { NSSM_REG_DIR, REG_EXPAND_SZ, (void *) _T(""), false, 0, setting_set_string, setting_get_string },
   { NSSM_REG_EXIT, REG_SZ, (void *) exit_action_strings[NSSM_EXIT_RESTART], false, ADDITIONAL_MANDATORY, setting_set_exit_action, setting_get_exit_action },
+  { NSSM_REG_AFFINITY, REG_SZ, 0, false, 0, setting_set_affinity, setting_get_affinity },
   { NSSM_REG_ENV, REG_MULTI_SZ, NULL, false, ADDITIONAL_CRLF, setting_set_environment, setting_get_environment },
   { NSSM_REG_ENV_EXTRA, REG_MULTI_SZ, NULL, false, ADDITIONAL_CRLF, setting_set_environment, setting_get_environment },
   { NSSM_REG_PRIORITY, REG_SZ, (void *) priority_strings[NSSM_NORMAL_PRIORITY], false, 0, setting_set_priority, setting_get_priority },