/mordor/daemon.cpp
C++ | 453 lines | 396 code | 43 blank | 14 comment | 60 complexity | ec487a584bf5c7850db8622527ab8d9e MD5 | raw file
1// Copyright (c) 2010 - Mozy, Inc. 2 3#include "daemon.h" 4 5#include "log.h" 6#include "main.h" 7 8namespace Mordor { 9namespace Daemon { 10 11static Logger::ptr g_log = Log::lookup("mordor:daemon"); 12 13boost::signals2::signal<void ()> onTerminate; 14boost::signals2::signal<void ()> onInterrupt; 15boost::signals2::signal<void ()> onReload; 16boost::signals2::signal<void ()> onPause; 17boost::signals2::signal<void ()> onContinue; 18 19#define CTRL_CASE(ctrl) \ 20 case ctrl: \ 21 return os << #ctrl; \ 22 23#ifdef WINDOWS 24static boost::function<int (int, char **)> g_daemonMain; 25 26namespace { 27struct ServiceStatus 28{ 29 SERVICE_STATUS_HANDLE serviceHandle; 30 SERVICE_STATUS status; 31}; 32} 33 34static DWORD acceptedControls() 35{ 36 DWORD accepted = 0; 37 if (!onTerminate.empty()) 38 accepted |= SERVICE_ACCEPT_STOP; 39 if (!onReload.empty()) 40 accepted |= SERVICE_ACCEPT_PARAMCHANGE; 41 if (!onPause.empty()) 42 accepted |= SERVICE_ACCEPT_PAUSE_CONTINUE; 43 return accepted; 44} 45 46namespace { 47 48enum ServiceCtrl {}; 49 50std::ostream &operator <<(std::ostream &os, ServiceCtrl ctrl) 51{ 52 switch (ctrl) { 53 CTRL_CASE(SERVICE_CONTROL_STOP); 54 CTRL_CASE(SERVICE_CONTROL_PAUSE); 55 CTRL_CASE(SERVICE_CONTROL_CONTINUE); 56 CTRL_CASE(SERVICE_CONTROL_INTERROGATE); 57 CTRL_CASE(SERVICE_CONTROL_SHUTDOWN); 58 CTRL_CASE(SERVICE_CONTROL_PARAMCHANGE); 59 CTRL_CASE(SERVICE_CONTROL_NETBINDADD); 60 CTRL_CASE(SERVICE_CONTROL_NETBINDREMOVE); 61 CTRL_CASE(SERVICE_CONTROL_NETBINDENABLE); 62 CTRL_CASE(SERVICE_CONTROL_NETBINDDISABLE); 63 CTRL_CASE(SERVICE_CONTROL_DEVICEEVENT); 64 CTRL_CASE(SERVICE_CONTROL_HARDWAREPROFILECHANGE); 65 CTRL_CASE(SERVICE_CONTROL_POWEREVENT); 66 CTRL_CASE(SERVICE_CONTROL_SESSIONCHANGE); 67 CTRL_CASE(SERVICE_CONTROL_PRESHUTDOWN); 68#ifdef SERVICE_CONTROL_TIMECHANGE 69 CTRL_CASE(SERVICE_CONTROL_TIMECHANGE); 70#endif 71#ifdef SERVICE_CONTROL_TRIGGEREVENT 72 CTRL_CASE(SERVICE_CONTROL_TRIGGEREVENT); 73#endif 74 default: 75 return os << (int)ctrl; 76 } 77} 78 79} 80 81static DWORD WINAPI HandlerEx(DWORD dwControl, DWORD dwEventType, 82 LPVOID lpEventData, LPVOID lpContext) 83{ 84 ServiceStatus *status = (ServiceStatus *)lpContext; 85 MORDOR_LOG_DEBUG(g_log) << "Received service control " 86 << (ServiceCtrl)dwControl; 87 88 switch (dwControl) { 89 case SERVICE_CONTROL_INTERROGATE: 90 break; 91 case SERVICE_CONTROL_PARAMCHANGE: 92 onReload(); 93 break; 94 case SERVICE_CONTROL_STOP: 95 if (onTerminate.empty()) 96 ExitProcess(0); 97 onTerminate(); 98 break; 99 case SERVICE_CONTROL_PAUSE: 100 if (!onPause.empty()) 101 status->status.dwCurrentState = SERVICE_PAUSED; 102 onPause(); 103 break; 104 case SERVICE_CONTROL_CONTINUE: 105 onContinue(); 106 status->status.dwCurrentState = SERVICE_RUNNING; 107 break; 108 default: 109 return ERROR_CALL_NOT_IMPLEMENTED; 110 } 111 MORDOR_LOG_TRACE(g_log) << "Service control " << (ServiceCtrl)dwControl 112 << " handled"; 113 // If it's stop, onTerminate may cause ServiceMain to exit immediately, 114 // before onTerminate returns, and our pointer to status on ServiceMain's 115 // stack would then be invalid; we update the service status in ServiceMain 116 // anyway, so it's okay to skip it here 117 if (dwControl != SERVICE_CONTROL_STOP) { 118 status->status.dwControlsAccepted = acceptedControls(); 119 SetServiceStatus(status->serviceHandle, &status->status); 120 } 121 return NO_ERROR; 122} 123 124static void WINAPI ServiceMain(DWORD argc, LPWSTR *argvW) 125{ 126 ServiceStatus status; 127 status.status.dwServiceType = SERVICE_WIN32_OWN_PROCESS; 128 status.status.dwCurrentState = SERVICE_RUNNING; 129 status.status.dwControlsAccepted = acceptedControls(); 130 status.status.dwWin32ExitCode = NO_ERROR; 131 status.status.dwServiceSpecificExitCode = 0; 132 status.status.dwCheckPoint = 0; 133 status.status.dwWaitHint = 0; 134 135 status.serviceHandle = RegisterServiceCtrlHandlerExW(NULL, 136 &HandlerEx, &status); 137 if (!status.serviceHandle) 138 return; 139 140 if (!SetServiceStatus(status.serviceHandle, &status.status)) 141 return; 142 143 MORDOR_LOG_INFO(g_log) << "Starting service"; 144 char **argv = CommandLineToUtf8(argc, argvW); 145 if (argv) { 146 try { 147 status.status.dwServiceSpecificExitCode = 148 g_daemonMain(argc, argv); 149 LocalFree(argv); 150 } catch (...) { 151 LocalFree(argv); 152 throw; 153 } 154 } else { 155 status.status.dwServiceSpecificExitCode = GetLastError(); 156 } 157 status.status.dwCurrentState = SERVICE_STOPPED; 158 if (status.status.dwServiceSpecificExitCode) 159 status.status.dwWin32ExitCode = ERROR_SERVICE_SPECIFIC_ERROR; 160 SetServiceStatus(status.serviceHandle, &status.status); 161 MORDOR_LOG_INFO(g_log) << "Service stopped"; 162} 163 164namespace { 165 166enum ConsoleCtrl {}; 167 168std::ostream &operator <<(std::ostream &os, ConsoleCtrl ctrl) 169{ 170 switch (ctrl) { 171 CTRL_CASE(CTRL_C_EVENT); 172 CTRL_CASE(CTRL_BREAK_EVENT); 173 CTRL_CASE(CTRL_CLOSE_EVENT); 174 CTRL_CASE(CTRL_LOGOFF_EVENT); 175 CTRL_CASE(CTRL_SHUTDOWN_EVENT); 176 default: 177 return os << (int)ctrl; 178 } 179} 180 181} 182 183static BOOL WINAPI HandlerRoutine(DWORD dwCtrlType) 184{ 185 MORDOR_LOG_DEBUG(g_log) << "Received console control " 186 << (ConsoleCtrl)dwCtrlType; 187 switch (dwCtrlType) { 188 case CTRL_C_EVENT: 189 if (onInterrupt.empty()) 190 return FALSE; 191 onInterrupt(); 192 break; 193 case CTRL_CLOSE_EVENT: 194 case CTRL_BREAK_EVENT: 195 if (onTerminate.empty()) 196 return FALSE; 197 onTerminate(); 198 break; 199 default: 200 return FALSE; 201 } 202 MORDOR_LOG_TRACE(g_log) << "Console control " << (ServiceCtrl)dwCtrlType 203 << " handled"; 204 return TRUE; 205} 206 207int run(int argc, char **argv, 208 boost::function<int (int, char **)> daemonMain, bool) 209{ 210 SERVICE_TABLE_ENTRYW ServiceStartTable[] = { 211 { L"", &ServiceMain }, 212 { NULL, NULL } 213 }; 214 g_daemonMain = daemonMain; 215 if (!StartServiceCtrlDispatcherW(ServiceStartTable)) { 216 if (GetLastError() == ERROR_FAILED_SERVICE_CONTROLLER_CONNECT) { 217 MORDOR_LOG_INFO(g_log) << "Starting console process"; 218 SetConsoleCtrlHandler(&HandlerRoutine, TRUE); 219 int result = daemonMain(argc, argv); 220 SetConsoleCtrlHandler(&HandlerRoutine, FALSE); 221 MORDOR_LOG_INFO(g_log) << "Console process stopped"; 222 return result; 223 } else { 224 return GetLastError(); 225 } 226 } else { 227 return 0; 228 } 229} 230#else 231 232#include <errno.h> 233#include <signal.h> 234#include <sys/wait.h> 235boost::function<bool (pid_t, int)> onChildProcessExit; 236 237static sigset_t blockedSignals() 238{ 239 sigset_t result; 240 sigemptyset(&result); 241 sigaddset(&result, SIGTERM); 242 sigaddset(&result, SIGINT); 243 sigaddset(&result, SIGHUP); 244 sigaddset(&result, SIGTSTP); 245 sigaddset(&result, SIGCONT); 246 return result; 247} 248 249static void unblockAndRaise(int signal) 250{ 251 sigset_t set; 252 sigemptyset(&set); 253 sigaddset(&set, signal); 254 pthread_sigmask(SIG_UNBLOCK, &set, NULL); 255 raise(signal); 256 pthread_sigmask(SIG_BLOCK, &set, NULL); 257} 258 259namespace { 260 261enum Signal {}; 262 263std::ostream &operator <<(std::ostream &os, Signal signal) 264{ 265 switch ((int)signal) { 266 CTRL_CASE(SIGTERM); 267 CTRL_CASE(SIGINT); 268 CTRL_CASE(SIGTSTP); 269 CTRL_CASE(SIGCONT); 270 CTRL_CASE(SIGHUP); 271 default: 272 return os << (int)signal; 273 } 274} 275 276} 277 278static void *signal_thread(void *arg) 279{ 280 sigset_t mask = blockedSignals(); 281 while (true) { 282 int caught; 283 int rc = sigwait(&mask, &caught); 284 if (rc != 0) 285 continue; 286 MORDOR_LOG_DEBUG(g_log) << "Received signal " << (Signal)caught; 287 switch (caught) { 288 case SIGTERM: 289 if (onTerminate.empty()) { 290 unblockAndRaise(caught); 291 continue; 292 } 293 onTerminate(); 294 break; 295 case SIGINT: 296 if (onInterrupt.empty()) { 297 unblockAndRaise(caught); 298 continue; 299 } 300 onInterrupt(); 301 break; 302 case SIGTSTP: 303 onPause(); 304 unblockAndRaise(caught); 305 break; 306 case SIGCONT: 307 onContinue(); 308 unblockAndRaise(caught); 309 break; 310 case SIGHUP: 311 if (onReload.empty()) { 312 unblockAndRaise(caught); 313 continue; 314 } 315 onReload(); 316 break; 317 default: 318 continue; 319 } 320 MORDOR_LOG_TRACE(g_log) << "Signal " << (Signal)caught 321 << " handled"; 322 } 323 324 return NULL; 325} 326 327#ifndef OSX 328static bool shouldDaemonize(char **enviro) 329{ 330 if (!enviro) 331 return false; 332 std::string parent; 333 for (const char *env = *enviro; *env; env += strlen(env) + 1) { 334 const char *equals = strchr(env, '='); 335 if (equals != env + 1 || *env != '_') 336 continue; 337 parent = equals + 1; 338 break; 339 } 340 if (parent.size() >= 12 && 341 strncmp(parent.c_str(), "/etc/init.d/", 12) == 0) 342 return true; 343 if (parent.size() >= 17 && 344 strcmp(parent.c_str() + parent.size() - 17, "start-stop-daemon") == 0) 345 return true; 346 return false; 347} 348 349static bool shouldDaemonizeDueToParent() 350{ 351 std::ostringstream os; 352 os << "/proc/" << getppid() << "/environ"; 353 std::string parentEnviron; 354 parentEnviron.resize(65536); 355 int fd = open(os.str().c_str(), O_RDONLY); 356 if (fd < 0) 357 return false; 358 int size = read(fd, &parentEnviron[0], 65536); 359 close(fd); 360 if (size < 0) 361 return false; 362 parentEnviron.resize(size + 1); 363 const char *parentEnviro = parentEnviron.c_str(); 364 return shouldDaemonize((char **)&parentEnviro); 365} 366#endif 367 368// running user specified main 369static int startMain(int argc, char** argv, 370 boost::function<int (int, char**)> userMain) 371{ 372 // Mask signals from other threads so we can handle them ourselves 373 sigset_t mask = blockedSignals(); 374 int rc = pthread_sigmask(SIG_BLOCK, &mask, NULL); 375 if (rc != 0) 376 return errno; 377 378 // Create the thread that will handle signals for us 379 pthread_t signal_thread_id; 380 rc = pthread_create(&signal_thread_id, NULL, &signal_thread, NULL); 381 if (rc != 0) 382 return errno; 383 384 // Run the daemon's main 385 MORDOR_LOG_INFO(g_log) << "Starting main in daemon"; 386 rc = userMain(argc, argv); 387 MORDOR_LOG_INFO(g_log) << "Main stopped in daemon"; 388 return rc; 389} 390 391// start main in a forked process and monitor it, restart it in case 392// it dies abnormally. 393static int watchdog(int argc, char** argv, 394 boost::function<int (int, char**)> userMain) 395{ 396 for (;;) { 397 pid_t pid = fork(); 398 switch (pid) { 399 case -1: 400 MORDOR_LOG_ERROR(g_log) << "Unable to fork(2): " << errno; 401 return errno; 402 case 0: 403 MORDOR_LOG_INFO(g_log) << "Child #" << getpid() << " started"; 404 return startMain(argc, argv, userMain); 405 default: 406 MORDOR_LOG_INFO(g_log) << "Watchdog starts monitoring child #" << pid; 407 int status; 408 if (waitpid(pid, &status, 0) == -1) { 409 MORDOR_LOG_ERROR(g_log) << "Failed to waitpid(2): " << errno; 410 return errno; 411 } 412 MORDOR_LOG_INFO(g_log) << "Child #" << pid << " dies with status:" 413 << status; 414 bool done = false; 415 // check if program wish to quit 416 if (!onChildProcessExit.empty()) { 417 done = onChildProcessExit(pid, status); 418 MORDOR_LOG_INFO(g_log) << "onChildProcessExit returns " << done; 419 } else { 420 MORDOR_LOG_INFO(g_log) << "onChildProcessExit is not set, " 421 << "restart child by default"; 422 } 423 if (done) { 424 MORDOR_LOG_INFO(g_log) << "Watchdog stopped"; 425 return 0; 426 } 427 } 428 } 429} 430 431int run(int argc, char **argv, 432 boost::function<int (int, char **)> daemonMain, 433 bool enableWatchdog) 434{ 435#ifndef OSX 436 // Check for being run from /etc/init.d or start-stop-daemon as a hint to 437 // daemonize 438 if (shouldDaemonize(environ) || shouldDaemonizeDueToParent()) { 439 MORDOR_LOG_VERBOSE(g_log) << "Daemonizing"; 440 if (daemon(0, 0) == -1) 441 return errno; 442 } 443#endif 444 445 if (enableWatchdog) { 446 return watchdog(argc, argv, daemonMain); 447 } else { 448 return startMain(argc, argv, daemonMain); 449 } 450} 451#endif 452 453}}