/mordor/daemon.cpp
C++ | 453 lines | 396 code | 43 blank | 14 comment | 60 complexity | ec487a584bf5c7850db8622527ab8d9e MD5 | raw file
Possible License(s): BSD-3-Clause
- // Copyright (c) 2010 - Mozy, Inc.
- #include "daemon.h"
- #include "log.h"
- #include "main.h"
- namespace Mordor {
- namespace Daemon {
- static Logger::ptr g_log = Log::lookup("mordor:daemon");
- boost::signals2::signal<void ()> onTerminate;
- boost::signals2::signal<void ()> onInterrupt;
- boost::signals2::signal<void ()> onReload;
- boost::signals2::signal<void ()> onPause;
- boost::signals2::signal<void ()> onContinue;
- #define CTRL_CASE(ctrl) \
- case ctrl: \
- return os << #ctrl; \
- #ifdef WINDOWS
- static boost::function<int (int, char **)> g_daemonMain;
- namespace {
- struct ServiceStatus
- {
- SERVICE_STATUS_HANDLE serviceHandle;
- SERVICE_STATUS status;
- };
- }
- static DWORD acceptedControls()
- {
- DWORD accepted = 0;
- if (!onTerminate.empty())
- accepted |= SERVICE_ACCEPT_STOP;
- if (!onReload.empty())
- accepted |= SERVICE_ACCEPT_PARAMCHANGE;
- if (!onPause.empty())
- accepted |= SERVICE_ACCEPT_PAUSE_CONTINUE;
- return accepted;
- }
- namespace {
- enum ServiceCtrl {};
- std::ostream &operator <<(std::ostream &os, ServiceCtrl ctrl)
- {
- switch (ctrl) {
- CTRL_CASE(SERVICE_CONTROL_STOP);
- CTRL_CASE(SERVICE_CONTROL_PAUSE);
- CTRL_CASE(SERVICE_CONTROL_CONTINUE);
- CTRL_CASE(SERVICE_CONTROL_INTERROGATE);
- CTRL_CASE(SERVICE_CONTROL_SHUTDOWN);
- CTRL_CASE(SERVICE_CONTROL_PARAMCHANGE);
- CTRL_CASE(SERVICE_CONTROL_NETBINDADD);
- CTRL_CASE(SERVICE_CONTROL_NETBINDREMOVE);
- CTRL_CASE(SERVICE_CONTROL_NETBINDENABLE);
- CTRL_CASE(SERVICE_CONTROL_NETBINDDISABLE);
- CTRL_CASE(SERVICE_CONTROL_DEVICEEVENT);
- CTRL_CASE(SERVICE_CONTROL_HARDWAREPROFILECHANGE);
- CTRL_CASE(SERVICE_CONTROL_POWEREVENT);
- CTRL_CASE(SERVICE_CONTROL_SESSIONCHANGE);
- CTRL_CASE(SERVICE_CONTROL_PRESHUTDOWN);
- #ifdef SERVICE_CONTROL_TIMECHANGE
- CTRL_CASE(SERVICE_CONTROL_TIMECHANGE);
- #endif
- #ifdef SERVICE_CONTROL_TRIGGEREVENT
- CTRL_CASE(SERVICE_CONTROL_TRIGGEREVENT);
- #endif
- default:
- return os << (int)ctrl;
- }
- }
- }
- static DWORD WINAPI HandlerEx(DWORD dwControl, DWORD dwEventType,
- LPVOID lpEventData, LPVOID lpContext)
- {
- ServiceStatus *status = (ServiceStatus *)lpContext;
- MORDOR_LOG_DEBUG(g_log) << "Received service control "
- << (ServiceCtrl)dwControl;
- switch (dwControl) {
- case SERVICE_CONTROL_INTERROGATE:
- break;
- case SERVICE_CONTROL_PARAMCHANGE:
- onReload();
- break;
- case SERVICE_CONTROL_STOP:
- if (onTerminate.empty())
- ExitProcess(0);
- onTerminate();
- break;
- case SERVICE_CONTROL_PAUSE:
- if (!onPause.empty())
- status->status.dwCurrentState = SERVICE_PAUSED;
- onPause();
- break;
- case SERVICE_CONTROL_CONTINUE:
- onContinue();
- status->status.dwCurrentState = SERVICE_RUNNING;
- break;
- default:
- return ERROR_CALL_NOT_IMPLEMENTED;
- }
- MORDOR_LOG_TRACE(g_log) << "Service control " << (ServiceCtrl)dwControl
- << " handled";
- // If it's stop, onTerminate may cause ServiceMain to exit immediately,
- // before onTerminate returns, and our pointer to status on ServiceMain's
- // stack would then be invalid; we update the service status in ServiceMain
- // anyway, so it's okay to skip it here
- if (dwControl != SERVICE_CONTROL_STOP) {
- status->status.dwControlsAccepted = acceptedControls();
- SetServiceStatus(status->serviceHandle, &status->status);
- }
- return NO_ERROR;
- }
- static void WINAPI ServiceMain(DWORD argc, LPWSTR *argvW)
- {
- ServiceStatus status;
- status.status.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
- status.status.dwCurrentState = SERVICE_RUNNING;
- status.status.dwControlsAccepted = acceptedControls();
- status.status.dwWin32ExitCode = NO_ERROR;
- status.status.dwServiceSpecificExitCode = 0;
- status.status.dwCheckPoint = 0;
- status.status.dwWaitHint = 0;
- status.serviceHandle = RegisterServiceCtrlHandlerExW(NULL,
- &HandlerEx, &status);
- if (!status.serviceHandle)
- return;
- if (!SetServiceStatus(status.serviceHandle, &status.status))
- return;
- MORDOR_LOG_INFO(g_log) << "Starting service";
- char **argv = CommandLineToUtf8(argc, argvW);
- if (argv) {
- try {
- status.status.dwServiceSpecificExitCode =
- g_daemonMain(argc, argv);
- LocalFree(argv);
- } catch (...) {
- LocalFree(argv);
- throw;
- }
- } else {
- status.status.dwServiceSpecificExitCode = GetLastError();
- }
- status.status.dwCurrentState = SERVICE_STOPPED;
- if (status.status.dwServiceSpecificExitCode)
- status.status.dwWin32ExitCode = ERROR_SERVICE_SPECIFIC_ERROR;
- SetServiceStatus(status.serviceHandle, &status.status);
- MORDOR_LOG_INFO(g_log) << "Service stopped";
- }
- namespace {
- enum ConsoleCtrl {};
- std::ostream &operator <<(std::ostream &os, ConsoleCtrl ctrl)
- {
- switch (ctrl) {
- CTRL_CASE(CTRL_C_EVENT);
- CTRL_CASE(CTRL_BREAK_EVENT);
- CTRL_CASE(CTRL_CLOSE_EVENT);
- CTRL_CASE(CTRL_LOGOFF_EVENT);
- CTRL_CASE(CTRL_SHUTDOWN_EVENT);
- default:
- return os << (int)ctrl;
- }
- }
- }
- static BOOL WINAPI HandlerRoutine(DWORD dwCtrlType)
- {
- MORDOR_LOG_DEBUG(g_log) << "Received console control "
- << (ConsoleCtrl)dwCtrlType;
- switch (dwCtrlType) {
- case CTRL_C_EVENT:
- if (onInterrupt.empty())
- return FALSE;
- onInterrupt();
- break;
- case CTRL_CLOSE_EVENT:
- case CTRL_BREAK_EVENT:
- if (onTerminate.empty())
- return FALSE;
- onTerminate();
- break;
- default:
- return FALSE;
- }
- MORDOR_LOG_TRACE(g_log) << "Console control " << (ServiceCtrl)dwCtrlType
- << " handled";
- return TRUE;
- }
- int run(int argc, char **argv,
- boost::function<int (int, char **)> daemonMain, bool)
- {
- SERVICE_TABLE_ENTRYW ServiceStartTable[] = {
- { L"", &ServiceMain },
- { NULL, NULL }
- };
- g_daemonMain = daemonMain;
- if (!StartServiceCtrlDispatcherW(ServiceStartTable)) {
- if (GetLastError() == ERROR_FAILED_SERVICE_CONTROLLER_CONNECT) {
- MORDOR_LOG_INFO(g_log) << "Starting console process";
- SetConsoleCtrlHandler(&HandlerRoutine, TRUE);
- int result = daemonMain(argc, argv);
- SetConsoleCtrlHandler(&HandlerRoutine, FALSE);
- MORDOR_LOG_INFO(g_log) << "Console process stopped";
- return result;
- } else {
- return GetLastError();
- }
- } else {
- return 0;
- }
- }
- #else
- #include <errno.h>
- #include <signal.h>
- #include <sys/wait.h>
- boost::function<bool (pid_t, int)> onChildProcessExit;
- static sigset_t blockedSignals()
- {
- sigset_t result;
- sigemptyset(&result);
- sigaddset(&result, SIGTERM);
- sigaddset(&result, SIGINT);
- sigaddset(&result, SIGHUP);
- sigaddset(&result, SIGTSTP);
- sigaddset(&result, SIGCONT);
- return result;
- }
- static void unblockAndRaise(int signal)
- {
- sigset_t set;
- sigemptyset(&set);
- sigaddset(&set, signal);
- pthread_sigmask(SIG_UNBLOCK, &set, NULL);
- raise(signal);
- pthread_sigmask(SIG_BLOCK, &set, NULL);
- }
- namespace {
- enum Signal {};
- std::ostream &operator <<(std::ostream &os, Signal signal)
- {
- switch ((int)signal) {
- CTRL_CASE(SIGTERM);
- CTRL_CASE(SIGINT);
- CTRL_CASE(SIGTSTP);
- CTRL_CASE(SIGCONT);
- CTRL_CASE(SIGHUP);
- default:
- return os << (int)signal;
- }
- }
- }
- static void *signal_thread(void *arg)
- {
- sigset_t mask = blockedSignals();
- while (true) {
- int caught;
- int rc = sigwait(&mask, &caught);
- if (rc != 0)
- continue;
- MORDOR_LOG_DEBUG(g_log) << "Received signal " << (Signal)caught;
- switch (caught) {
- case SIGTERM:
- if (onTerminate.empty()) {
- unblockAndRaise(caught);
- continue;
- }
- onTerminate();
- break;
- case SIGINT:
- if (onInterrupt.empty()) {
- unblockAndRaise(caught);
- continue;
- }
- onInterrupt();
- break;
- case SIGTSTP:
- onPause();
- unblockAndRaise(caught);
- break;
- case SIGCONT:
- onContinue();
- unblockAndRaise(caught);
- break;
- case SIGHUP:
- if (onReload.empty()) {
- unblockAndRaise(caught);
- continue;
- }
- onReload();
- break;
- default:
- continue;
- }
- MORDOR_LOG_TRACE(g_log) << "Signal " << (Signal)caught
- << " handled";
- }
- return NULL;
- }
- #ifndef OSX
- static bool shouldDaemonize(char **enviro)
- {
- if (!enviro)
- return false;
- std::string parent;
- for (const char *env = *enviro; *env; env += strlen(env) + 1) {
- const char *equals = strchr(env, '=');
- if (equals != env + 1 || *env != '_')
- continue;
- parent = equals + 1;
- break;
- }
- if (parent.size() >= 12 &&
- strncmp(parent.c_str(), "/etc/init.d/", 12) == 0)
- return true;
- if (parent.size() >= 17 &&
- strcmp(parent.c_str() + parent.size() - 17, "start-stop-daemon") == 0)
- return true;
- return false;
- }
- static bool shouldDaemonizeDueToParent()
- {
- std::ostringstream os;
- os << "/proc/" << getppid() << "/environ";
- std::string parentEnviron;
- parentEnviron.resize(65536);
- int fd = open(os.str().c_str(), O_RDONLY);
- if (fd < 0)
- return false;
- int size = read(fd, &parentEnviron[0], 65536);
- close(fd);
- if (size < 0)
- return false;
- parentEnviron.resize(size + 1);
- const char *parentEnviro = parentEnviron.c_str();
- return shouldDaemonize((char **)&parentEnviro);
- }
- #endif
- // running user specified main
- static int startMain(int argc, char** argv,
- boost::function<int (int, char**)> userMain)
- {
- // Mask signals from other threads so we can handle them ourselves
- sigset_t mask = blockedSignals();
- int rc = pthread_sigmask(SIG_BLOCK, &mask, NULL);
- if (rc != 0)
- return errno;
- // Create the thread that will handle signals for us
- pthread_t signal_thread_id;
- rc = pthread_create(&signal_thread_id, NULL, &signal_thread, NULL);
- if (rc != 0)
- return errno;
- // Run the daemon's main
- MORDOR_LOG_INFO(g_log) << "Starting main in daemon";
- rc = userMain(argc, argv);
- MORDOR_LOG_INFO(g_log) << "Main stopped in daemon";
- return rc;
- }
- // start main in a forked process and monitor it, restart it in case
- // it dies abnormally.
- static int watchdog(int argc, char** argv,
- boost::function<int (int, char**)> userMain)
- {
- for (;;) {
- pid_t pid = fork();
- switch (pid) {
- case -1:
- MORDOR_LOG_ERROR(g_log) << "Unable to fork(2): " << errno;
- return errno;
- case 0:
- MORDOR_LOG_INFO(g_log) << "Child #" << getpid() << " started";
- return startMain(argc, argv, userMain);
- default:
- MORDOR_LOG_INFO(g_log) << "Watchdog starts monitoring child #" << pid;
- int status;
- if (waitpid(pid, &status, 0) == -1) {
- MORDOR_LOG_ERROR(g_log) << "Failed to waitpid(2): " << errno;
- return errno;
- }
- MORDOR_LOG_INFO(g_log) << "Child #" << pid << " dies with status:"
- << status;
- bool done = false;
- // check if program wish to quit
- if (!onChildProcessExit.empty()) {
- done = onChildProcessExit(pid, status);
- MORDOR_LOG_INFO(g_log) << "onChildProcessExit returns " << done;
- } else {
- MORDOR_LOG_INFO(g_log) << "onChildProcessExit is not set, "
- << "restart child by default";
- }
- if (done) {
- MORDOR_LOG_INFO(g_log) << "Watchdog stopped";
- return 0;
- }
- }
- }
- }
- int run(int argc, char **argv,
- boost::function<int (int, char **)> daemonMain,
- bool enableWatchdog)
- {
- #ifndef OSX
- // Check for being run from /etc/init.d or start-stop-daemon as a hint to
- // daemonize
- if (shouldDaemonize(environ) || shouldDaemonizeDueToParent()) {
- MORDOR_LOG_VERBOSE(g_log) << "Daemonizing";
- if (daemon(0, 0) == -1)
- return errno;
- }
- #endif
- if (enableWatchdog) {
- return watchdog(argc, argv, daemonMain);
- } else {
- return startMain(argc, argv, daemonMain);
- }
- }
- #endif
- }}