[PATCH v2] trace-cmd: Stop recording when processes traced with -P exit

[PATCH v2] trace-cmd: Stop recording when processes traced with -P exit

[Bijan Tabatabai]

From: Bijan Tabatabai <bijan311@yahoo.com>

When the -F flag is used in trace-cmd record, trace-cmd stops recording
when the executable it is tracing terminates.
This patch makes the -P flag act similarly.

Signed-off-by: Bijan Tabatabai <bijan311@yahoo.com>
---
 I changed this patch to use the pidfd interface to determine if a process had
 ended.
 I also refactored the patch a bit, so it looks much cleaner in my opinion.
---
 tracecmd/trace-record.c | 81 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 81 insertions(+)

diff --git a/tracecmd/trace-record.c b/tracecmd/trace-record.c
index d8c24eb..5b9c961 100644
--- a/tracecmd/trace-record.c
+++ b/tracecmd/trace-record.c
@@ -16,6 +16,7 @@
 #include <sys/time.h>
 #include <sys/wait.h>
 #include <sys/socket.h>
+#include <sys/syscall.h>
 #include <sys/utsname.h>
 #ifndef NO_PTRACE
 #include <sys/ptrace.h>
@@ -33,6 +34,7 @@
 #include <errno.h>
 #include <limits.h>
 #include <libgen.h>
+#include <poll.h>
 #include <pwd.h>
 #include <grp.h>
 #ifdef VSOCK
@@ -121,8 +123,10 @@ struct filter_pids {
 };
 
 static struct filter_pids *filter_pids;
+static struct filter_pids *process_pids;
 static int nr_filter_pids;
 static int len_filter_pids;
+static int nr_process_pids = 0;
 
 static int have_set_event_pid;
 static int have_event_fork;
@@ -1243,6 +1247,79 @@ static pid_t trace_waitpid(enum trace_type type, pid_t pid, int *status, int opt
 			trace_stream_read(pids, recorder_threads, &tv);
 	} while (1);
 }
+
+#ifndef __NR_pidfd_open
+#define __NR_pidfd_open 434
+#endif
+
+static int pidfd_open(pid_t pid, unsigned int flags) {
+	return syscall(__NR_pidfd_open, pid, flags);
+}
+
+static int trace_waitpidfd(id_t pidfd) {
+	struct pollfd pollfd;
+
+	pollfd.fd = pidfd;
+	pollfd.events = POLLIN;
+
+	while (!finished) {
+		int ret = poll(&pollfd, 1, -1);
+		/* If waitid was interrupted, keep waiting */
+		if (ret < 0 && errno == EINTR)
+			continue;
+		else if (ret < 0)
+			return 1;
+		else
+			break;
+	}
+
+	return 0;
+}
+
+static int trace_wait_for_processes() {
+	int ret = 0;
+	int nr_fds = 0;
+	int i;
+	int *pidfds;
+	struct filter_pids *pid;
+
+	pidfds = malloc(sizeof(int) * nr_process_pids);
+	if (!pidfds)
+		return 1;
+
+	for (pid = process_pids; pid && nr_process_pids; pid = pid->next) {
+		if (pid->exclude) {
+			nr_process_pids--;
+			continue;
+		}
+		pidfds[nr_fds] = pidfd_open(pid->pid, 0);
+
+		/* If the pid doesn't exist, the process has probably exited */
+		if (pidfds[nr_fds] < 0 && errno == ESRCH) {
+			nr_process_pids--;
+			continue;
+		} else if (pidfds[nr_fds] < 0) {
+			ret = 1;
+			goto out;
+		}
+
+		nr_fds++;
+		nr_process_pids--;
+	}
+
+	for (i = 0; i < nr_fds; i++) {
+		if (trace_waitpidfd(pidfds[i])) {
+			ret = 1;
+			goto out;
+		}
+	}
+
+out:
+	for (i = 0; i < nr_fds; i++)
+		close(pidfds[i]);
+	free(pidfds);
+	return ret;
+}
 #ifndef NO_PTRACE
 
 /**
@@ -5832,7 +5909,9 @@ static void parse_record_options(int argc,
 			while (pid) {
 				add_filter_pid(atoi(pid), 0);
 				pid = strtok_r(NULL, ",", &sav);
+				nr_process_pids++;
 			}
+			process_pids = filter_pids;
 			free(pids);
 			break;
 		case 'c':
@@ -6308,6 +6387,8 @@ static void record_trace(int argc, char **argv,
 		}
 		/* sleep till we are woken with Ctrl^C */
 		printf("Hit Ctrl^C to stop recording\n");
+		if (nr_process_pids && !trace_wait_for_processes())
+			finished = 1;
 		while (!finished)
 			trace_or_sleep(type);
 	}
-- 
2.25.1