Add a little more

1 files changed, 109 insertions, 12 deletions

diff --git a/svcsupervise.c b/svcsupervise.c
index 1fce356..0770f87 100644
--- a/svcsupervise.c
+++ b/svcsupervise.c
@@ -11,6 +11,7 @@
 #include <sys/stat.h>
 #include <sys/file.h>
 #include <sys/wait.h>
+#include <sys/procctl.h>
 
 #include "status.h"
 
@@ -94,6 +95,9 @@ handle_ctrl_command(int fd)
 		return 1;
 	} else {
 		switch (cmd) {
+		case '\n':
+			/* Ignore newline */
+			break;
 		case 'd':
 			/* Down */
 			break;
@@ -101,23 +105,31 @@ handle_ctrl_command(int fd)
 			/* Up */
 			break;
 		case 'o':
+			/* Once, up do not restart */
 			break;
 		case 'a':
+			/* SIGALRM */
 			break;
 		case 'h':
+			/* SIGHUP */
 			break;
 		case 'k':
+			/* SIGKILL */
 			break;
 		case 't':
+			/* SIGTERM */
 			break;
 		case 'i':
+			/* SIGINT */
 			break;
 		case 'p':
+			/* SIGSTOP */
 			break;
 		case 'c':
+			/* SIGCONT */
 			break;
 		case 'x':
-			return 0;
+			return 0; /* Exit */
 			break;
 		default:
 			fprintf(stderr, "Unknown command \"%c\"", cmd);
@@ -158,11 +170,61 @@ reset_signals()
 		err(1, "sigaction()");
 }
 
+void
+reap_all()
+{
+	int r, m = 0;
+
+	for (;;) {
+		r = waitpid(-1, NULL, WNOHANG);
+
+		if (r == 0) {
+			if (m >= 3)
+				break;
+			sleep(1);
+			m++;
+		}
+
+		if (r == -1) {
+			if (errno == EINTR)
+				continue;
+			if (errno != ECHILD)
+				perror("waitpid()");
+			break;
+		}
+	}
+}
+
+int
+try_wait()
+{
+	int p, s, r = 0;
+	for (;;) {
+		p = waitpid(-1, &s, WNOHANG);
+		if (p == 0)
+			break;
+		if (p == -1) {
+			if (errno == EINTR) /* Impossible? */
+				continue;
+			if (errno != ECHILD)
+				perror("waitpid()");
+			break;
+		}
+
+		r = 1;
+	}
+
+	return r;
+}
+
 int
 main(int argc, char **argv)
 {
-	int ctrl_fd, ok_fd, lock_fd, kq, e, i;
-	struct kevent evt[6], revt[6];
+	int ctrl_fd, ok_fd, lock_fd, kq, e, i, wantup;
+	struct stat sb;
+	struct kevent evt[7], revt[7];
+	struct procctl_reaper_kill rk;
+	struct procctl_reaper_status rs;
 
 	if (argc != 2)
 		errx(1, "Usage: %s <dir>\n", argv[0]);
@@ -181,6 +243,9 @@ main(int argc, char **argv)
 	if (mkfifo("supervise/ok", 0600) == -1 && errno != EEXIST)
 		err(1, "mkfifo()");
 
+	if (procctl(P_PID, getpid(), PROC_REAP_ACQUIRE, NULL) == -1)
+		err(1, "procctl()");
+
 	setup_signals();
 
 	kq = kqueue();
@@ -193,17 +258,28 @@ main(int argc, char **argv)
 	if ((ok_fd = open("supervise/ok", O_RDONLY | O_NONBLOCK | O_CLOEXEC)) == -1)
 		err(1, "open()");
 
+	if (stat("./down", &sb) == -1) {
+		if (errno == ENOENT) {
+			wantup = 1;
+		} else
+			err(1, "lstat()");
+	} else {
+		wantup = 0;
+	}
+
 	EV_SET(&evt[0], SIGHUP,  EVFILT_SIGNAL, EV_ADD | EV_ENABLE, 0, 0, 0);
 	EV_SET(&evt[1], SIGINT,  EVFILT_SIGNAL, EV_ADD | EV_ENABLE, 0, 0, 0);
 	EV_SET(&evt[2], SIGTERM, EVFILT_SIGNAL, EV_ADD | EV_ENABLE, 0, 0, 0);
 	EV_SET(&evt[3], SIGCHLD, EVFILT_SIGNAL, EV_ADD | EV_ENABLE, 0, 0, 0);
 	EV_SET(&evt[4], ctrl_fd, EVFILT_READ,   EV_ADD | EV_ENABLE, 0, 0, 0);
-	EV_SET(&evt[5], 1,       EVFILT_TIMER,  EV_ADD | EV_ENABLE | EV_ONESHOT, NOTE_SECONDS, 5, 0);
+	EV_SET(&evt[5], ok_fd,   EVFILT_READ,   EV_ADD | EV_ENABLE, 0, 0, 0);
+	EV_SET(&evt[6], 1,       EVFILT_TIMER,  EV_ADD | EV_ENABLE | EV_ONESHOT, NOTE_SECONDS, 5, 0);
 
-	if (kevent(kq, evt, 5, NULL, 0, NULL) == -1)
+	if (kevent(kq, evt, 6, NULL, 0, NULL) == -1)
 		err(1, "kevent()");
 
-	start_proc();
+	if (wantup)
+		start_proc();
 
 	for (;;) {
 		e = kevent(kq, NULL, 0, revt, 6, NULL);
@@ -212,14 +288,15 @@ main(int argc, char **argv)
 				perror("kevent()");
 			continue;
 		} else if (e > 0) {
-			/* XXX: Do the things */
 			for (i = 0; i < e; i++) {
 				if (revt[i].filter == EVFILT_SIGNAL) {
 					if (revt[i].ident == SIGCHLD) {
-						wait(NULL); /* XXX: We have to this better */
-						if (kevent(kq, &evt[5], 1, NULL, 0, NULL) == -1)
-							perror("kevent()");
-					} else if (revt[i].ident == SIGHUP ||
+						if (try_wait()) {
+							if (kevent(kq, &evt[6], 1, NULL, 0, NULL) == -1)
+								perror("kevent()");
+						}
+					}
+					if (revt[i].ident == SIGHUP ||
 					    revt[i].ident == SIGINT ||
 					    revt[i].ident == SIGTERM) {
 						goto end;
@@ -231,6 +308,11 @@ main(int argc, char **argv)
 						if (handle_ctrl_command(ctrl_fd) == 0)
 							goto end;
 					}
+					if (revt[i].ident == ok_fd) {
+						char buf[1024];
+						read(ok_fd, buf, sizeof(buf));
+						/* Make sure the ok pipe doesn't fill up */
+					}
 				} else {
 					fprintf(stderr, "Unknown event\n");
 				}
@@ -239,7 +321,20 @@ main(int argc, char **argv)
 	}
 
 end:
-	fprintf(stderr, "Sup bye\n");
+	if (procctl(P_PID, getpid(), PROC_REAP_STATUS, &rs) != -1) {
+		if (rs.rs_children > 0) {
+			rk.rk_sig = SIGTERM;
+			rk.rk_flags = REAPER_KILL_CHILDREN;
+			if (procctl(P_PID, getpid(), PROC_REAP_KILL, &rk))
+				perror("procctl()");
+		}
+	} else {
+		perror("procctl()");
+	}
+
+	if (close(kq) == -1)
+		perror("close()");
+
 	if (close(ctrl_fd) == -1)
 		perror("close()");
 
@@ -248,5 +343,7 @@ end:
 
 	reset_signals();
 
+	reap_all();
+
 	return 0;
 }