initial commit

1 files changed, 624 insertions, 0 deletions

diff --git a/rvnit.c b/rvnit.c
new file mode 100644
index 0000000..a93d2d0
--- /dev/null
+++ b/rvnit.c
@@ -0,0 +1,624 @@
+/* accept4 */
+#define _GNU_SOURCE
+
+#include <sys/wait.h>
+#include <sys/socket.h>
+#include <sys/un.h>
+//#include <sys/stat.h>
+
+#include <ctype.h>
+#include <dirent.h>
+#include <errno.h>
+#include <fcntl.h>
+#include <poll.h>
+#include <pthread.h>
+#include <signal.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <time.h>
+#include <unistd.h>
+
+#define MAX_SV 512
+
+struct entry {
+	char level;
+	char name[62];
+	char logged;
+	pid_t pid;
+	pid_t display_pid;  // not reset to 0, needed for oneshots
+	int logfd[2];
+	time_t start;
+	int status;
+	enum {
+		UNDEF,
+		UP,
+		DOWN,
+	} state;
+};
+
+struct entry services[MAX_SV];
+int level;
+
+pthread_t main_thread;
+pthread_t socket_thread;
+pthread_t logger_thread;
+
+int selflogfd[2];
+
+sig_atomic_t got_intr;
+sig_atomic_t got_sigusr1;
+
+void
+on_sigint(int sig)
+{
+	(void)sig;
+	got_intr = 1;
+}
+
+void
+on_sigusr1(int sig)
+{
+	(void)sig;
+	got_sigusr1 = 1;
+}
+
+void
+on_sigusr2(int sig)
+{
+	(void)sig;
+}
+
+void
+restart(int i)
+{
+	if (services[i].pid > 0)
+		return;
+	if (services[i].state != UP)
+		return;
+
+	int delay = 0;
+	time_t now = time(0);
+	if (services[i].start == now)
+		delay = 1;
+
+	int loggerpipe[2] = { -1, -1 };
+
+	if (services[i].logfd[0] == -1) {
+		pipe(services[i].logfd);
+		fcntl(services[i].logfd[0], F_SETFL, O_NONBLOCK);
+		fcntl(services[i].logfd[1], F_SETFL, O_NONBLOCK);
+
+		pthread_kill(logger_thread, SIGUSR2);
+	}
+
+	if (services[i].name[2] == 'L') {
+		for (int j = 0; j < MAX_SV; j++) {
+			if (strcmp(services[i].name + 3,
+			    services[j].name + 3) == 0 &&
+			    services[j].name[2] == 'D') {
+				if (services[j].logfd[0] == -1) {
+					// hook up logger
+					pipe(loggerpipe);
+					services[j].logfd[0] = loggerpipe[0];
+					services[j].logfd[1] = loggerpipe[1];
+					services[j].logged = 1;
+				} else {
+					// recover loggerpipe
+					loggerpipe[0] = services[j].logfd[0];
+					loggerpipe[1] = services[j].logfd[1];
+				}
+				break;
+			}
+		}
+	}
+
+	pid_t child = fork();
+	if (child == 0) {
+		dup2(services[i].logfd[1], 1);
+		if (loggerpipe[0] != -1) { // loggers get read end of loggerpipe
+			dup2(loggerpipe[0], 0);
+			close(loggerpipe[0]);
+			close(loggerpipe[1]);
+		}
+
+		close(services[i].logfd[0]);
+		close(services[i].logfd[1]);
+		sleep(delay);
+		execl(services[i].name,
+		    services[i].name,
+		    (char *)0);
+		exit(-1);
+	} else if (child > 0) {
+		fcntl(services[i].logfd[0], F_SETFD, FD_CLOEXEC);
+		fcntl(services[i].logfd[1], F_SETFD, FD_CLOEXEC);
+		if (loggerpipe[0] != -1) {
+			fcntl(loggerpipe[0], F_SETFD, FD_CLOEXEC);
+			fcntl(loggerpipe[1], F_SETFD, FD_CLOEXEC);
+		}
+		services[i].pid = services[i].display_pid = child;
+		services[i].start = now;
+	} else {
+		abort();
+	}
+}
+
+void *
+socket_loop(void* ignored)
+{
+	(void)ignored;
+
+	const char *path = "/tmp/rvnit.sock";
+
+	struct sockaddr_un addr = { 0 };
+	addr.sun_family = AF_UNIX;
+	strncpy(addr.sun_path, path, sizeof addr.sun_path - 1);
+	int listenfd = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0);
+	if (listenfd < 0) {
+		perror("socket");
+		exit(111);
+	}
+	unlink(path);
+	int r = bind(listenfd, (struct sockaddr *)&addr, sizeof addr);
+	if (r < 0) {
+		perror("bind");
+		exit(111);
+	}
+
+	r = listen(listenfd, SOMAXCONN);
+        if (r < 0) {
+                perror("listen");
+                exit(111);
+        }
+
+	while (1) {
+		int connfd = accept4(listenfd, 0, 0, SOCK_CLOEXEC);
+		if (connfd < 0)
+			continue;
+
+		char cmd = 0;
+		read(connfd, &cmd, 1);
+		write(connfd, "ok\n", 3);
+
+		if (cmd == 'l') {
+			dprintf(connfd, "level=%d\n", level);
+		}
+
+		if (cmd == 's') {
+			for (int i = 0; i < MAX_SV; i++) {
+				if (services[i].name[2] == 'D' ||
+				    services[i].name[2] == 'L')
+					dprintf(connfd, "%s pid=%d status=%d\n",
+					    services[i].name,
+					    services[i].pid,
+					    services[i].status);
+			}
+		}
+
+		if (cmd == 'd' || cmd == 'u') {
+			char buf[63];
+			ssize_t rd = read(connfd, buf, sizeof buf);
+			if (rd > 0) {
+				buf[rd] = 0;
+				printf("got %c|%s|\n", cmd, buf);
+
+				for (int i = 0; i < MAX_SV; i++) {
+					if (strcmp(services[i].name, buf) == 0) {
+						if (cmd == 'd') {
+							printf("setting %s down\n",
+							    services[i].name);
+							services[i].state = DOWN;
+						} else if (cmd == 'u') {
+							printf("setting %s up\n",
+							    services[i].name);
+							services[i].state = UP;
+						}
+					}
+				}
+			}
+
+			pthread_kill(main_thread, SIGUSR1);
+		}
+
+		if (cmd == 'X') {
+			pthread_kill(main_thread, SIGPWR);
+		}
+
+		close(connfd);
+	}
+
+	return 0;
+}
+
+void *
+logger_loop(void* ignored)
+{
+	(void)ignored;
+
+//	sigaction(SIGUSR2, &(struct sigaction){.sa_handler=on_sigusr2}, 0);
+
+	while (1) {
+		struct pollfd fds[MAX_SV];
+		nfds_t nfds = 0;
+
+		fds[nfds].fd = selflogfd[0];
+		fds[nfds].events = POLLIN;
+		nfds++;
+
+		for (int i = 0; i < MAX_SV; i++) {
+			if (/*services[i].pid > 0 && */ !services[i].logged && services[i].logfd[0] > 0) {
+				fds[nfds].fd = services[i].logfd[0];
+				fds[nfds].events = POLLIN;
+				nfds++;
+			}
+		}
+
+//		printf("waiting for %d logs\n", (int)nfds);
+		int n = poll(fds, nfds, -1);
+		if (n < 0)
+			continue;
+
+		if (n == 1 && (fds[0].revents & POLLHUP) &&
+		    !(fds[0].revents & POLLIN)) {
+			/* selflog was closed, end thread */
+			break;
+		}
+
+		for (nfds_t j = 0; j < nfds; j++) {
+			// XXX HUP, NVAL
+			if (!(fds[j].revents & POLLIN))
+				continue;
+
+			char buf[4096];
+			ssize_t rd = read(fds[j].fd, buf, sizeof buf);
+			if (rd < 0) {
+				perror("read");
+				continue;
+			}
+			if (fds[j].fd == selflogfd[0]) {
+				printf("%s[%d]: %.*s",
+				    "rvnit",
+				    (int)getpid(),
+				    (int)rd,
+				    buf);
+				continue;
+			}
+			for (int i = 0; i < MAX_SV; i++) {
+				if (services[i].logfd[0] == fds[j].fd) {
+					printf("%s[%d]: %.*s",
+					    services[i].name,
+					    services[i].display_pid,
+					    (int)rd,
+					    buf
+						);
+					break;
+				}
+			}
+		}
+	}
+
+	return 0;
+}
+
+int
+reap(pid_t pid, int status)
+{
+	for (int i = 0; i < MAX_SV; i++) {
+		if (services[i].pid != pid)
+			continue;
+
+		services[i].pid = 0;
+		services[i].status = status;
+
+		dprintf(2, "reaped %s[%d] with status %d during level=%d\n",
+		    services[i].name, pid, status, level);
+
+		return i;
+	}
+
+	return -1;
+}
+
+int
+main()
+{
+	pipe(selflogfd);
+	fcntl(selflogfd[0], F_SETFL, O_NONBLOCK);
+	fcntl(selflogfd[1], F_SETFL, O_NONBLOCK);
+	fcntl(selflogfd[0], F_SETFD, FD_CLOEXEC);
+	fcntl(selflogfd[1], F_SETFD, FD_CLOEXEC);
+
+	sigaction(SIGUSR2, &(struct sigaction){.sa_handler=on_sigusr2}, 0);
+
+	main_thread = pthread_self();
+	pthread_create(&socket_thread, 0, socket_loop, 0);
+	pthread_create(&logger_thread, 0, logger_loop, 0);
+
+	if (chdir("sv") < 0) {
+		perror("chdir");
+		return 111;
+	}
+
+	DIR *dir = opendir(".");
+	if (!dir)
+		return 111;
+
+	int i = 0;
+
+	struct dirent *ent;
+	while ((ent = readdir(dir))) {
+		if (!isdigit(ent->d_name[0]) || !isdigit(ent->d_name[1]))
+			continue;
+
+		services[i].level = 10 * (ent->d_name[0] - '0') +
+		    (ent->d_name[1] - '0');
+		(void)snprintf(services[i].name, sizeof services[i].name,
+		    "%s", ent->d_name);
+		// follow symlinks
+		services[i].state = access(ent->d_name, X_OK) == 0 ? UP : DOWN;
+		services[i].logfd[0] = -1;
+		services[i].logfd[1] = -1;
+		services[i].logged = 0;
+		i++;
+	}
+
+	closedir(dir);
+
+	dprintf(selflogfd[1], "booting\n");
+
+	for (level = 0; level < 100; level++) {
+		/* spawn all of level */
+		int oneshot = 0;
+
+		// spawn loggers first
+		for (i = 0; i < MAX_SV; i++) {
+			if (services[i].level != level)
+				continue;
+
+			if (services[i].state != UP)
+				continue;
+
+			if (services[i].name[2] == 'L')
+				restart(i);
+		}
+
+		// spawn oneshots and daemons
+		for (i = 0; i < MAX_SV; i++) {
+			if (services[i].level != level)
+				continue;
+
+			if (services[i].state != UP)
+				continue;
+
+			if (services[i].name[2] == 'S') {
+				restart(i);
+				oneshot++;
+			} else if (services[i].name[2] == 'D') {
+				restart(i);
+			}
+		}
+
+		while (oneshot) {
+			int status = 0;
+			int pid = wait(&status);
+
+			if (pid < 0 && errno == ECHILD)
+				break;
+
+			int i = reap(pid, status);
+			if (i < 0)
+				continue;
+
+			if (services[i].level != level)
+				continue;
+
+			if (services[i].name[2] == 'S') {
+				oneshot--;
+			} else if (services[i].name[2] == 'D' ||
+			    services[i].name[2] == 'L') {
+				restart(i);
+			}
+		}
+	}
+
+	sigaction(SIGINT, &(struct sigaction){.sa_handler=on_sigint}, 0);
+	sigaction(SIGPWR, &(struct sigaction){.sa_handler=on_sigint}, 0);
+	sigaction(SIGUSR1, &(struct sigaction){.sa_handler=on_sigusr1}, 0);
+
+	dprintf(selflogfd[1], "system up\n");
+
+	while (1) {
+		if (got_intr)
+			break;
+
+		if (got_sigusr1) {
+			printf("rescanning state\n");
+			for (i = 0; i < MAX_SV; i++) {
+				if (services[i].name[2] == 'D' ||
+				    services[i].name[2] == 'L') {
+					if (services[i].state == UP &&
+					    services[i].pid == 0)
+						restart(i);
+					if (services[i].state == DOWN &&
+					    services[i].pid > 0)
+						kill(services[i].pid, SIGTERM);
+				}
+			}
+
+			got_sigusr1 = 0;
+		}
+
+		int status = 0;
+		errno = 0;
+		int pid = wait(&status);
+
+		if (pid < 0) {
+			if (errno == ECHILD)
+				break;
+			if (errno == EINTR)
+				continue;
+		}
+
+		int i = reap(pid, status);
+		if (i < 0)
+			continue;
+
+		if (services[i].name[2] == 'D' ||
+		    services[i].name[2] == 'L') {
+			dprintf(selflogfd[1], "%s terminated with status %d\n", services[i].name, services[i].status);
+			if (services[i].state == UP) {
+				dprintf(selflogfd[1], "restarting %s\n", services[i].name);
+				restart(i);
+			}
+		}
+	}
+
+	dprintf(selflogfd[1], "shutting down\n");
+
+//	sigaction(SIGINT, &(struct sigaction){.sa_handler=SIG_DFL}, 0);
+
+
+	struct timespec timeout = {7, 0};
+	sigset_t childset;
+	sigemptyset(&childset);
+	sigaddset(&childset, SIGCHLD);
+
+
+	for (level = 99; level >= 0; level--) {
+		/* kill all of level */
+		int oneshot = 0;
+		int daemons = 0;
+
+		for (i = 0; i < MAX_SV; i++) {
+			if (services[i].level != level)
+				continue;
+
+			if (services[i].name[2] == 'K' &&
+			    services[i].state == UP) {
+				dprintf(2, "todo: shutdown oneshot %s\n", services[i].name);
+				oneshot++;
+			}
+			if (services[i].name[2] == 'D' &&
+			    services[i].pid > 0) {
+				dprintf(2, "todo: sending sigterm to %s\n", services[i].name);
+				daemons++;
+			}
+		}
+
+		dprintf(2, "level=%d oneshot=%d daemons=%d\n",
+		    level, oneshot, daemons);
+	}
+
+	for (level = 99; level >= 0; level--) {
+		/* kill all of level */
+		int oneshot = 0;
+		int daemons = 0;
+
+		for (i = 0; i < MAX_SV; i++) {
+			if (services[i].level != level)
+				continue;
+
+			if (services[i].name[2] == 'K' &&
+			    services[i].state == UP) {
+				dprintf(selflogfd[1], "starting shutdown oneshot %s\n", services[i].name);
+				restart(i);
+				oneshot++;
+			}
+			if (services[i].name[2] == 'D' &&
+			    services[i].pid > 0) {
+				dprintf(selflogfd[1], "sending sigterm to %s\n", services[i].name);
+				kill(services[i].pid, SIGTERM);
+				kill(services[i].pid, SIGCONT);
+				daemons++;
+			}
+		}
+
+		dprintf(2, "level=%d oneshot=%d daemons=%d\n",
+		    level, oneshot, daemons);
+		usleep(100000);
+
+		// XXX figure out logger shutdown
+
+		while (oneshot) {
+			int status = 0;
+			int pid = wait(&status);
+
+			if (pid < 0 && errno == ECHILD)
+				break;
+
+			int i = reap(pid, status);
+			if (i < 0)
+				continue;
+
+			if (services[i].level != level)
+				continue;
+
+			if (services[i].name[2] == 'K') {
+				dprintf(selflogfd[1], "oneshot %s exited with status %d\n", services[i].name, services[i].status);
+				oneshot--;
+			} else if (services[i].name[2] == 'D') {
+				dprintf(selflogfd[1], "daemon %s exited with status %d\n", services[i].name, services[i].status);
+//				if (services[i].logged)
+//					close(services[i].logfd[1]);
+				daemons--;
+			}
+		}
+
+		// only daemons are left, wait up to 7s before sending
+		// them SIGKILL.
+
+		if (!daemons)
+			continue;
+
+		sigprocmask(SIG_BLOCK, &childset, 0);
+
+		while (daemons) {
+			int status = 0;
+			int pid = waitpid(-1, &status, WNOHANG);
+			if (pid == 0) { // nothing to reap
+				if (sigtimedwait(&childset, 0, &timeout) == SIGCHLD)
+					continue;
+				if (errno == EAGAIN) { // hit timeout
+					printf("need kill\n");
+					for (i = 0; i < MAX_SV; i++) {
+						if (services[i].level != level)
+							continue;
+						if (services[i].name[2] == 'D' &&
+						    services[i].pid > 0)
+							kill(services[i].pid, SIGKILL);
+					}
+					// XXX ensure this is only run once,
+					// then force next level
+					continue;
+				}
+			}
+
+			if (pid < 0 && errno == ECHILD)
+				break;
+
+			int i = reap(pid, status);
+			if (i < 0)
+				continue;
+
+			if (services[i].level != level)
+				continue;
+
+			if (services[i].name[2] == 'K') {
+				// can't happen
+				dprintf(selflogfd[1], "oneshot %s exited with status %d\n", services[i].name, services[i].status);
+				oneshot--;
+			} else if (services[i].name[2] == 'D') {
+				dprintf(selflogfd[1], "daemon %s exited with status %d\n", services[i].name, services[i].status);
+//				if (services[i].logged)
+//					close(services[i].logfd[1]);
+				daemons--;
+			}
+		}
+
+		sigprocmask(SIG_UNBLOCK, &childset, 0);
+	}
+
+	dprintf(selflogfd[1], "shutdown\n");
+	close(selflogfd[1]);
+	pthread_join(logger_thread, 0);
+}