PageRenderTime 39ms CodeModel.GetById 1ms RepoModel.GetById 2ms app.codeStats 0ms

/mordor/daemon.cpp

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