[make-initrd] [PATCH v1 02/11] feature/procacct: Use epoll

[Alexey Gladkov <gladkov.alexey@gmail.com>]

This is preparation for adding more event sources.

Signed-off-by: Alexey Gladkov <gladkov.alexey@gmail.com>
---
 features/debug-procacct/src/procacct.c | 288 +++++++++++++++++--------
 1 file changed, 200 insertions(+), 88 deletions(-)

diff --git a/features/debug-procacct/src/procacct.c b/features/debug-procacct/src/procacct.c
index 9302edb0..cfb4de1f 100644
--- a/features/debug-procacct/src/procacct.c
+++ b/features/debug-procacct/src/procacct.c
@@ -31,6 +31,7 @@
 #include <sys/stat.h>
 #include <sys/socket.h>
 #include <sys/wait.h>
+#include <sys/epoll.h>
 
 #include <stdio.h>
 #include <stdlib.h>
@@ -43,6 +44,7 @@
 #include <search.h>
 
 #include "rd/logging.h"
+#include "rd/memory.h"
 
 /* Maximum size of response requested or message sent */
 #define MAX_MSG_SIZE 1024
@@ -63,6 +65,7 @@
 static pid_t current_pid;
 static int rcvbufsz;
 static char name[100];
+static int fd_out = 1;
 
 struct msgtemplate {
 	struct nlmsghdr n;
@@ -77,14 +80,44 @@ struct proc_cmdline {
 	char *cmdline;
 };
 
+struct ctx_netlink {
+	uint16_t cpuid;
+	uint16_t cpumask_len;
+	char cpumask[100 + 6 * MAX_CPUS];
+};
+
+enum {
+	FD_NETLINK,
+	FD_MAX,
+};
+
+#define EV_MAX (FD_MAX * 32)
+
+struct fd_handler;
+typedef int (*fdhandler_t)(struct fd_handler *);
+
+struct fd_handler {
+	int fd;
+	fdhandler_t fd_prepare;
+	fdhandler_t fd_handler;
+	fdhandler_t fd_finish;
+	void *data;
+};
+
+static struct fd_handler fd_handler_list[FD_MAX];
+
 static void usage(void)                                                       __attribute__((noreturn));
 static int proc_compare(const void *a, const void *b)                         __attribute__((nonnull(1, 2)));
-static int create_nl_socket(int protocol);
+static int process_netlink_events(struct fd_handler *el);
+static int prepare_netlink(struct fd_handler *el);
+static int finish_netlink(struct fd_handler *el);
+static void setup_netlink_fd(struct fd_handler *el);
 static ssize_t send_cmd(int fd, uint16_t nlmsg_type, uint8_t genl_cmd,
                         uint16_t nla_type, void *nla_data, uint16_t nla_len);
 static uint16_t get_family_id(int fd);
 static void print_procacct(int fd, struct taskstats *t)                       __attribute__((nonnull(2)));
 static void handle_aggr(struct nlattr *na, int fd)                            __attribute__((nonnull(1)));
+static int setup_epoll_fd(struct fd_handler list[FD_MAX]);
 
 int proc_compare(const void *a, const void *b)
 {
@@ -100,44 +133,35 @@ int proc_compare(const void *a, const void *b)
 
 void usage(void)
 {
-	fprintf(stderr, "procacct [-o logfile] [-r bufsize] [-m cpumask]\n");
+	fprintf(stderr, "procacct [-o logfile] [-r bufsize]\n");
 	exit(1);
 }
 
 /*
  * Create a raw netlink socket and bind
  */
-int create_nl_socket(int protocol)
+void setup_netlink_fd(struct fd_handler *el)
 {
-	int fd;
 	struct sockaddr_nl local;
 
-	fd = socket(AF_NETLINK, SOCK_RAW, protocol);
-	if (fd < 0)
-		return -1;
+	el->fd = socket(AF_NETLINK, SOCK_RAW | SOCK_NONBLOCK | SOCK_CLOEXEC, NETLINK_GENERIC);
+	if (el->fd < 0)
+		rd_fatal("error creating Netlink socket: %m");
 
-	if (rcvbufsz) {
-		if (setsockopt(fd, SOL_SOCKET, SO_RCVBUF, &rcvbufsz, sizeof(rcvbufsz)) < 0) {
-			rd_err("unable to set socket rcv buf size to %d", rcvbufsz);
-			goto error;
-		}
-	}
+	if (rcvbufsz && setsockopt(el->fd, SOL_SOCKET, SO_RCVBUF, &rcvbufsz, sizeof(rcvbufsz)) < 0)
+		rd_fatal("unable to set socket rcv buf size to %d", rcvbufsz);
 
 	memset(&local, 0, sizeof(local));
 	local.nl_family = AF_NETLINK;
 
-	if (bind(fd, (struct sockaddr *) &local, sizeof(local)) < 0) {
-		rd_err("unable bind to socket");
-		goto error;
-	}
+	if (bind(el->fd, (struct sockaddr *) &local, sizeof(local)) < 0)
+		rd_fatal("unable bind to socket");
 
-	return fd;
-error:
-	close(fd);
-	return -1;
+	el->fd_handler = process_netlink_events;
+	el->fd_prepare = prepare_netlink;
+	el->fd_finish  = finish_netlink;
 }
 
-
 ssize_t send_cmd(int fd, uint16_t nlmsg_type, uint8_t genl_cmd, uint16_t nla_type, void *nla_data, uint16_t nla_len)
 {
 	struct nlattr *na;
@@ -268,89 +292,80 @@ void handle_aggr(struct nlattr *na, int fd)
 	}
 }
 
-int main(int argc, char *argv[])
+int prepare_netlink(struct fd_handler *el)
 {
+	struct ctx_netlink *ctx;
 	ssize_t ret;
-	uint16_t id;
 
-	char cpumask[100 + 6 * MAX_CPUS];
-	uint16_t cpumask_len;
+	ctx = rd_calloc_or_die(1, sizeof(*ctx));
 
-	int fd_nlink = -1;
-	int fd_out = 1;
-	int write_file = 0;
-	int maskset = 0;
-	char *logfile = NULL;
+	long np = sysconf(_SC_NPROCESSORS_ONLN);
+	if (np > 1)
+		snprintf(ctx->cpumask, sizeof(ctx->cpumask), "0-%ld", np - 1);
+	else
+		snprintf(ctx->cpumask, sizeof(ctx->cpumask), "1");
 
-	while (1) {
-		int c = getopt(argc, argv, "m:o:r:");
-		if (c < 0)
-			break;
-
-		switch (c) {
-			case 'o':
-				logfile = strdup(optarg);
-				write_file = 1;
-				break;
-			case 'r':
-				rcvbufsz = atoi(optarg);
-				if (rcvbufsz < 0)
-					rd_fatal("invalid rcv buf size");
-				break;
-			case 'm':
-				strlcpy(cpumask, optarg, sizeof(cpumask));
-				maskset = 1;
-				break;
-			default:
-				usage();
-		}
-	}
-	if (!maskset) {
-		long np = sysconf(_SC_NPROCESSORS_ONLN);
-		if (np > 1)
-			snprintf(cpumask, sizeof(cpumask), "0-%ld", np - 1);
-		else
-			snprintf(cpumask, sizeof(cpumask), "1");
+	if ((strlen(ctx->cpumask) + 1) > USHRT_MAX) {
+		rd_err("cpumask too long");
+		free(ctx);
+		return -1;
 	}
 
-	if ((strlen(cpumask) + 1) > USHRT_MAX)
-		rd_fatal("cpumask too long");
+	ctx->cpumask_len = (uint16_t) strlen(ctx->cpumask) + 1;
 
-	cpumask_len = (uint16_t) strlen(cpumask) + 1;
+	ctx->cpuid = get_family_id(el->fd);
+	if (!ctx->cpuid) {
+		rd_err("error getting family id, errno=%d", errno);
+		free(ctx);
+		return -1;
+	}
 
-	if (write_file) {
-		fd_out = open(logfile, O_WRONLY | O_CREAT | O_TRUNC | O_SYNC, 0644);
-		if (fd_out < 0)
-			rd_fatal("cannot open output file: %s: %m", logfile);
+	ret = send_cmd(el->fd, ctx->cpuid, TASKSTATS_CMD_GET,
+	               TASKSTATS_CMD_ATTR_REGISTER_CPUMASK,
+		       &ctx->cpumask, ctx->cpumask_len);
+	if (ret < 0) {
+		rd_err("error sending register cpumask");
+		free(ctx);
+		return -1;
 	}
 
-	fd_nlink = create_nl_socket(NETLINK_GENERIC);
-	if (fd_nlink < 0)
-		rd_fatal("error creating Netlink socket: %m");
+	el->data = ctx;
 
-	current_pid = getpid();
+	return 0;
+}
 
-	id = get_family_id(fd_nlink);
-	if (!id) {
-		rd_err("error getting family id, errno=%d", errno);
-		goto err;
-	}
+int finish_netlink(struct fd_handler *el)
+{
+	struct ctx_netlink *ctx = el->data;
+	ssize_t ret;
 
-	ret = send_cmd(fd_nlink, id, TASKSTATS_CMD_GET,
-	               TASKSTATS_CMD_ATTR_REGISTER_CPUMASK, &cpumask, cpumask_len);
+	ret = send_cmd(el->fd, ctx->cpuid, TASKSTATS_CMD_GET,
+	               TASKSTATS_CMD_ATTR_DEREGISTER_CPUMASK,
+		       &ctx->cpumask, ctx->cpumask_len);
 	if (ret < 0) {
-		rd_err("error sending register cpumask");
-		goto err;
+		rd_err("error sending deregister cpumask");
+		return -1;
 	}
+	free(ctx);
+	return 0;
+}
+
+int process_netlink_events(struct fd_handler *el)
+{
+	ssize_t ret;
 
 	while (1) {
 		struct msgtemplate msg;
 		struct nlattr *na;
 
-		ret = recv(fd_nlink, &msg, sizeof(msg), 0);
+		errno = 0;
+		ret = recv(el->fd, &msg, sizeof(msg), 0);
 		if (ret < 0) {
-			rd_err("nonfatal reply error: errno=%d", errno);
-			continue;
+			if (errno == EAGAIN || errno == EWOULDBLOCK)
+				return 0;
+			if (errno == EINTR)
+				continue;
+			rd_fatal("nonfatal reply error: %m");
 		}
 
 		if (msg.n.nlmsg_type == NLMSG_ERROR || !NLMSG_OK((&msg.n), ret)) {
@@ -381,12 +396,109 @@ int main(int argc, char *argv[])
 		}
 	}
 
-	ret = send_cmd(fd_nlink, id, TASKSTATS_CMD_GET,
-	               TASKSTATS_CMD_ATTR_DEREGISTER_CPUMASK, &cpumask, cpumask_len);
-	if (ret < 0)
-		rd_fatal("error sending deregister cpumask");
+	return 0;
+}
+
+int setup_epoll_fd(struct fd_handler list[FD_MAX])
+{
+	int epollfd;
+	struct epoll_event ev = { .events = EPOLLIN };
+
+	if ((epollfd = epoll_create1(EPOLL_CLOEXEC)) < 0)
+		rd_fatal("epoll_create1: %m");
+
+	for (int i = 0; i < FD_MAX; i++) {
+		ev.data.ptr = &list[i];
+
+		if (epoll_ctl(epollfd, EPOLL_CTL_ADD, list[i].fd, &ev) < 0)
+			rd_fatal("epoll_ctl: %m");
+	}
+
+	return epollfd;
+}
+
+int main(int argc, char *argv[])
+{
+	int fd_epoll = -1;
+	int write_file = 0;
+	char *logfile = NULL;
+
+	while (1) {
+		int c = getopt(argc, argv, "o:r:");
+		if (c < 0)
+			break;
+
+		switch (c) {
+			case 'o':
+				logfile = strdup(optarg);
+				write_file = 1;
+				break;
+			case 'r':
+				rcvbufsz = atoi(optarg);
+				if (rcvbufsz < 0)
+					rd_fatal("invalid rcv buf size");
+				break;
+			default:
+				usage();
+		}
+	}
+
+	if (write_file) {
+		fd_out = open(logfile, O_WRONLY | O_CREAT | O_TRUNC | O_SYNC, 0644);
+		if (fd_out < 0)
+			rd_fatal("cannot open output file: %s: %m", logfile);
+	}
+
+	current_pid = getpid();
+
+	setup_netlink_fd(&fd_handler_list[FD_NETLINK]);
+	fd_epoll = setup_epoll_fd(fd_handler_list);
+
+	for (int i = 0; i < FD_MAX; i++) {
+		if (fd_handler_list[i].fd < 0 || !fd_handler_list[i].fd_prepare)
+			continue;
+		if (fd_handler_list[i].fd_prepare(&fd_handler_list[i]) < 0)
+			rd_fatal("unable to prepare file descriptor");
+	}
+
+	while (1) {
+		struct epoll_event ev[EV_MAX];
+		int fdcount;
+
+		errno = 0;
+		fdcount = epoll_wait(fd_epoll, ev, EV_MAX, -1);
+
+		if (fdcount < 0) {
+			if (errno == EINTR)
+				continue;
+			rd_fatal("epoll_wait: %m");
+		}
+
+		for (int i = 0; i < fdcount; i++) {
+			if (!(ev[i].events & EPOLLIN))
+				continue;
+
+			struct fd_handler *fde = ev[i].data.ptr;
+
+			if (fde->fd_handler(fde) != 0)
+				goto err;
+		}
+	}
+
 err:
-	close(fd_nlink);
+	for (int i = 0; i < FD_MAX; i++) {
+		if (fd_handler_list[i].fd < 0)
+			continue;
+
+		if (epoll_ctl(fd_epoll, EPOLL_CTL_DEL, fd_handler_list[i].fd, NULL) < 0)
+			rd_err("epoll_ctl: %m");
+
+		if (fd_handler_list[i].fd_finish)
+			fd_handler_list[i].fd_finish(&fd_handler_list[i]);
+
+		close(fd_handler_list[i].fd);
+	}
+	close(fd_epoll);
 
 	if (fd_out)
 		close(fd_out);
-- 
2.33.8