From 53371f115d94fbbc7e5cb60853b9e4a5d356d4b0 Mon Sep 17 00:00:00 2001 From: Iain Patterson Date: Wed, 8 Jan 2014 14:55:59 +0000 Subject: [PATCH] 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. --- ChangeLog.txt | 4 +- README.txt | 34 ++++++++++- gui.cpp | 76 +++++++++++++++++++++++++ messages.mc | Bin 126048 -> 134690 bytes nssm.cpp | 7 +++ nssm.h | 1 + nssm.rc | Bin 40722 -> 42814 bytes registry.cpp | 38 +++++++++++++ registry.h | 1 + resource.h | 4 +- service.cpp | 153 ++++++++++++++++++++++++++++++++++++++++++++++++++ service.h | 3 + settings.cpp | 117 ++++++++++++++++++++++++++++++++++++-- 13 files changed, 428 insertions(+), 10 deletions(-) diff --git a/ChangeLog.txt b/ChangeLog.txt index 141f31b..671b7d0 100644 --- a/ChangeLog.txt +++ b/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. diff --git a/README.txt b/README.txt index 7ed473e..685807a 100644 --- a/README.txt +++ b/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\\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 ------- diff --git a/gui.cpp b/gui.cpp index 7697e26..935198f 100644 --- a/gui.cpp +++ b/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); diff --git a/messages.mc b/messages.mc index 952012596261e00fd6e6b28bd9c8b397340c05e3..2ee1f8047585c0896bed186bc8ab1f2566b93e5f 100644 GIT binary patch delta 4766 zcmeHLUrdu%6h9|oA^_ppT#T$k)kxCE%als#YFH$ z5?yAx`*344@y{N7a2XhlF}@fQpACI6BufmKi7!6Lvp(#1ZYk1Q)-WALX1?_MzMgw~ z&OP^>^ZT9Kx!L?H7Ye5GCRB~f;oVDL;Mz;g)J`qbiMMWQGO`v5K+;Z!Ao&KeHrl5y zSAD0ZUgYUJ-E8MRH63wQMkz*a>Y<~!LKMMgn0|sJ23drTQ6JttWT_sXb8L_XA&;7S z2;;*T-9dI3={NJjko^b^@*K#iHAGfYH||l`azkG~el1c9Yp{&|@JkCHVIT%WL-Ygg z$>AVm8H?9^i=EC29gniB zLN=z4dNC_0@cqV`kiE%`t;2!z)$hqiKibDx!2P+XlS7z3NT^{P%uhrQ#6DD0|EkWIs%y+BPk4#Z!7Uj^n0Ks zE6x-cC4CfD@j351<hTB;;V;ld@yn1Tlgi^w@h4)bLF&?sg=>pn>_2;X$Dn$uMC@$);v%MNNoII7*mtDTDD8PL%n{99_<_ zQlngcvVzRi`DY)pT_>(_xte`e_?c)4;u5Ax)<%*!U)q!0mdri|MPf(RG@=v7*{^06 z_WUnA-LA$T1-M5~7qRnQi0bG6GOCDfP$9k5_(QuZBfEjd0U&cgEp~R53$?DyNy)td ze56Sc9YYP?oO2G;`;_C)Jl86E)a0qI8{-e>N|fW=&Mb>D399@PU7JuCs_C|GQwKGx z>AQ92`}GH&-wrQ6*NvC>aM`P#EVZM$;wJB3?nf)Ja(I=|%$clT*@@PwvoRO%(8c*& zrX5e&;g^aTz>Y}gCK`s#?!>?jQl-hxQZz_cn_M=4-gIi~XHs7;VSim(|FT}*>iH4w zTH84xJuqwPXKVk%gF@YMI9ER?v>^#%x-6fIB{}R3sWcch5Zez8De(TM15cEW>cu>+ zTYEphq5Su}UfaEK_M+Uzz54!XE`0?FQ&1KU@z6f$Qw4!zU}@oO=^Rq&mhJ4E4ejLY zY|-~Rp}10ib%q_~8GSYm^JWHpOrKCAx%JQ~_Nhy|OW%_D1NCS{zHWM3=I86VTSv~a z&C^N;!H~`P3#1KKChMdEb<)mJ^~a>mE6r4ne-aSBR5{sw)9=fVI0 delta 28 mcmV+%0OS9noCx692e1T(lTf4xvmTOuSF@Jz))J?t1OfZO}`6heK(wO|fQYq4rAqNN*7z!AQ z81fmC8B!UFfiNG)D`rq)&;iPLGPp1}GsH7EGPnVuCxaglhcHBfMSK{1z-kN`%oq%z zEDJDgHaRdueeyq9C3a_q0ESS8;>m&rsz96BCR>?tO*XL9VD)4OW(b-5QA2$4234+T zbkjjLBTO}hnr^`WG2H}8`!F~GO^Ii4WrzT}$rb1}7a-{Z6mtc-Oo5>Wh#bkX;&mkL_YH4wsd*oXfH*1x%sbC4w z2k0RhUkVK85@2X6fH7wB%%muJfOV diff --git a/registry.cpp b/registry.cpp index 6576753..c6189fe 100644 --- a/registry.cpp +++ b/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. */ diff --git a/registry.h b/registry.h index 6afa536..e0055b3 100644 --- a/registry.h +++ b/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); diff --git a/resource.h b/resource.h index baf15d8..0589119 100644 --- a/resource.h +++ b/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 diff --git a/service.cpp b/service.cpp index 72e0192..93b7d4a 100644 --- a/service.cpp +++ b/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 diff --git a/service.h b/service.h index 170ab42..0e5da20 100644 --- a/service.h +++ b/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); diff --git a/settings.cpp b/settings.cpp index 17a1cb3..25c7921 100644 --- a/settings.cpp +++ b/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 }, -- 2.20.1