Skip to content

Commit

Permalink
Try to throttle using a critical section.
Browse files Browse the repository at this point in the history
The first implementation of service restart throttling used a condition
variable in a critical section to sleep for the required amount of time.
The implementation was changed to use a waitable timer because Windows
2000 does not support SleepConditionVariableCS() or
WakeConditionVariable().

Since we are now using LoadLibrary() and GetProcAddress() to use newer
functions dynamically without having to build OS-specific binaries, we
can now use a critical section where it's supported and fall back to a
waitable timer when running on Windows 2000.
  • Loading branch information
Iain Patterson committed Nov 12, 2013
1 parent e42e690 commit 4550eb2
Show file tree
Hide file tree
Showing 3 changed files with 45 additions and 9 deletions.
10 changes: 10 additions & 0 deletions imports.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,16 @@ int get_imports() {
if (! imports.AttachConsole) {
if (error != ERROR_PROC_NOT_FOUND) return 2;
}

imports.SleepConditionVariableCS = (SleepConditionVariableCS_ptr) get_import(imports.kernel32, "SleepConditionVariableCS", &error);
if (! imports.SleepConditionVariableCS) {
if (error != ERROR_PROC_NOT_FOUND) return 3;
}

imports.WakeConditionVariable = (WakeConditionVariable_ptr) get_import(imports.kernel32, "WakeConditionVariable", &error);
if (! imports.WakeConditionVariable) {
if (error != ERROR_PROC_NOT_FOUND) return 4;
}
}
else if (error != ERROR_MOD_NOT_FOUND) return 1;

Expand Down
4 changes: 4 additions & 0 deletions imports.h
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,14 @@
#define IMPORTS_H

typedef BOOL (WINAPI *AttachConsole_ptr)(DWORD);
typedef BOOL (WINAPI *SleepConditionVariableCS_ptr)(PCONDITION_VARIABLE, PCRITICAL_SECTION, DWORD);
typedef void (WINAPI *WakeConditionVariable_ptr)(PCONDITION_VARIABLE);

typedef struct {
HMODULE kernel32;
AttachConsole_ptr AttachConsole;
SleepConditionVariableCS_ptr SleepConditionVariableCS;
WakeConditionVariable_ptr WakeConditionVariable;
} imports_t;

HMODULE get_dll(const char *, unsigned long *);
Expand Down
40 changes: 31 additions & 9 deletions service.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,15 @@ bool stopping;
bool allow_restart;
unsigned long throttle_delay;
unsigned long stop_method;
CRITICAL_SECTION throttle_section;
CONDITION_VARIABLE throttle_condition;
HANDLE throttle_timer;
LARGE_INTEGER throttle_duetime;
bool use_critical_section;
FILETIME creation_time;

extern imports_t imports;

static enum { NSSM_EXIT_RESTART, NSSM_EXIT_IGNORE, NSSM_EXIT_REALLY, NSSM_EXIT_UNCLEAN } exit_actions;
static const char *exit_action_strings[] = { "Restart", "Ignore", "Exit", "Suicide", 0 };

Expand Down Expand Up @@ -184,6 +189,10 @@ void WINAPI service_main(unsigned long argc, char **argv) {
return;
}

/* We can use a condition variable in a critical section on Vista or later. */
if (imports.SleepConditionVariableCS && imports.WakeConditionVariable) use_critical_section = true;
else use_critical_section = false;

/* Initialise status */
ZeroMemory(&service_status, sizeof(service_status));
service_status.dwServiceType = SERVICE_WIN32_OWN_PROCESS | SERVICE_INTERACTIVE_PROCESS;
Expand Down Expand Up @@ -218,9 +227,12 @@ void WINAPI service_main(unsigned long argc, char **argv) {
}

/* Used for signalling a resume if the service pauses when throttled. */
throttle_timer = CreateWaitableTimer(0, 1, 0);
if (! throttle_timer) {
log_event(EVENTLOG_WARNING_TYPE, NSSM_EVENT_CREATEWAITABLETIMER_FAILED, service_name, error_string(GetLastError()), 0);
if (use_critical_section) InitializeCriticalSection(&throttle_section);
else {
throttle_timer = CreateWaitableTimer(0, 1, 0);
if (! throttle_timer) {
log_event(EVENTLOG_WARNING_TYPE, NSSM_EVENT_CREATEWAITABLETIMER_FAILED, service_name, error_string(GetLastError()), 0);
}
}

monitor_service();
Expand Down Expand Up @@ -333,10 +345,13 @@ unsigned long WINAPI service_control_handler(unsigned long control, unsigned lon

case SERVICE_CONTROL_CONTINUE:
log_service_control(service_name, control, true);
if (! throttle_timer) return ERROR_CALL_NOT_IMPLEMENTED;
throttle = 0;
ZeroMemory(&throttle_duetime, sizeof(throttle_duetime));
SetWaitableTimer(throttle_timer, &throttle_duetime, 0, 0, 0, 0);
if (use_critical_section) imports.WakeConditionVariable(&throttle_condition);
else {
if (! throttle_timer) return ERROR_CALL_NOT_IMPLEMENTED;
ZeroMemory(&throttle_duetime, sizeof(throttle_duetime));
SetWaitableTimer(throttle_timer, &throttle_duetime, 0, 0, 0, 0);
}
service_status.dwCurrentState = SERVICE_CONTINUE_PENDING;
service_status.dwWaitHint = throttle_milliseconds() + NSSM_WAITHINT_MARGIN;
log_event(EVENTLOG_INFORMATION_TYPE, NSSM_EVENT_RESET_THROTTLE, service_name, 0);
Expand Down Expand Up @@ -558,7 +573,8 @@ void throttle_restart() {
_snprintf_s(milliseconds, sizeof(milliseconds), _TRUNCATE, "%d", ms);
log_event(EVENTLOG_WARNING_TYPE, NSSM_EVENT_THROTTLED, service_name, threshold, milliseconds, 0);

if (throttle_timer) {
if (use_critical_section) EnterCriticalSection(&throttle_section);
else if (throttle_timer) {
ZeroMemory(&throttle_duetime, sizeof(throttle_duetime));
throttle_duetime.QuadPart = 0 - (ms * 10000LL);
SetWaitableTimer(throttle_timer, &throttle_duetime, 0, 0, 0, 0);
Expand All @@ -567,6 +583,12 @@ void throttle_restart() {
service_status.dwCurrentState = SERVICE_PAUSED;
SetServiceStatus(service_handle, &service_status);

if (throttle_timer) WaitForSingleObject(throttle_timer, INFINITE);
else Sleep(ms);
if (use_critical_section) {
imports.SleepConditionVariableCS(&throttle_condition, &throttle_section, ms);
LeaveCriticalSection(&throttle_section);
}
else {
if (throttle_timer) WaitForSingleObject(throttle_timer, INFINITE);
else Sleep(ms);
}
}

0 comments on commit 4550eb2

Please sign in to comment.