libbpf-tools/offcputime: Add multi process/thread support

ekyooo · ekyooo · commit 68d5ac5b0d8b · 2024-09-20T18:42:20.000+09:00
This is a test example.

  # ./offcputime -p 16,48
    Tracing off-CPU time (us) of PID [16, 48]... Hit Ctrl-C to end.
        bpf_prog_a42aae11c0bc18f2_sched_switch
        bpf_prog_a42aae11c0bc18f2_sched_switch
        bpf_trace_run4
        __traceiter_sched_switch
        __schedule
        schedule
        worker_thread
        kthread
        ret_from_fork
        ret_from_fork_asm
        -                kworker/2:1 (48)
            3353019

        bpf_prog_a42aae11c0bc18f2_sched_switch
        bpf_prog_a42aae11c0bc18f2_sched_switch
        bpf_trace_run4
        __traceiter_sched_switch
        __schedule
        schedule
        rcu_gp_kthread
        kthread
        ret_from_fork
        ret_from_fork_asm
        -                rcu_preempt (16)
            1720974
diff --git a/libbpf-tools/offcputime.bpf.c b/libbpf-tools/offcputime.bpf.c
@@ -14,8 +14,8 @@ const volatile bool kernel_threads_only = false;
 const volatile bool user_threads_only = false;
 const volatile __u64 max_block_ns = -1;
 const volatile __u64 min_block_ns = 1;
-const volatile pid_t targ_tgid = -1;
-const volatile pid_t targ_pid = -1;
+const volatile bool filter_by_tgid = false;
+const volatile bool filter_by_pid = false;
 const volatile long state = -1;
 
 struct internal_key {
@@ -42,11 +42,28 @@ struct {
 	__uint(max_entries, MAX_ENTRIES);
 } info SEC(".maps");
 
+struct {
+	__uint(type, BPF_MAP_TYPE_HASH);
+	__type(key, u32);
+	__type(value, u8);
+	__uint(max_entries, MAX_PID_NR);
+} tgids SEC(".maps");
+
+struct {
+	__uint(type, BPF_MAP_TYPE_HASH);
+	__type(key, u32);
+	__type(value, u8);
+	__uint(max_entries, MAX_TID_NR);
+} pids SEC(".maps");
+
 static bool allow_record(struct task_struct *t)
 {
-	if (targ_tgid != -1 && targ_tgid != t->tgid)
+	u32 tgid = t->tgid;
+	u32 pid = t->pid;
+
+	if (filter_by_tgid && !bpf_map_lookup_elem(&tgids, &tgid))
 		return false;
-	if (targ_pid != -1 && targ_pid != t->pid)
+	if (filter_by_pid && !bpf_map_lookup_elem(&pids, &pid))
 		return false;
 	if (user_threads_only && t->flags & PF_KTHREAD)
 		return false;
diff --git a/libbpf-tools/offcputime.c b/libbpf-tools/offcputime.c
@@ -16,8 +16,8 @@
 #include "trace_helpers.h"
 
 static struct env {
-	pid_t pid;
-	pid_t tid;
+	pid_t pids[MAX_PID_NR];
+	pid_t tids[MAX_TID_NR];
 	bool user_threads_only;
 	bool kernel_threads_only;
 	int stack_storage_size;
@@ -28,8 +28,6 @@ static struct env {
 	int duration;
 	bool verbose;
 } env = {
-	.pid = -1,
-	.tid = -1,
 	.stack_storage_size = 1024,
 	.perf_max_stack_depth = 127,
 	.min_block_time = 1,
@@ -82,9 +80,31 @@ static const struct argp_option opts[] = {
 	{},
 };
 
+static int split_pidstr(char *s, char* sep, int max_split, pid_t *pids)
+{
+	char *pid;
+	int nr = 0;
+
+	errno = 0;
+	pid = strtok(s, sep);
+	while (pid) {
+		if (nr >= max_split)
+			return -ENOBUFS;
+
+		pids[nr++] = strtol(pid, NULL, 10);
+		if (errno)
+			return -errno;
+
+		pid = strtok(NULL, ",");
+	}
+
+	return 0;
+}
+
 static error_t parse_arg(int key, char *arg, struct argp_state *state)
 {
 	static int pos_args;
+	int ret;
 
 	switch (key) {
 	case 'h':
@@ -94,18 +114,26 @@ static error_t parse_arg(int key, char *arg, struct argp_state *state)
 		env.verbose = true;
 		break;
 	case 'p':
-		errno = 0;
-		env.pid = strtol(arg, NULL, 10);
-		if (errno) {
-			fprintf(stderr, "invalid PID: %s\n", arg);
+		ret = split_pidstr(strdup(arg), ",", MAX_PID_NR, env.pids);
+		if (ret) {
+			if (ret == -ENOBUFS)
+				fprintf(stderr, "the number of pid is too big, please "
+					"increase MAX_PID_NR's value and recompile\n");
+			else
+				fprintf(stderr, "invalid PID: %s\n", arg);
+
 			argp_usage(state);
 		}
 		break;
 	case 't':
-		errno = 0;
-		env.tid = strtol(arg, NULL, 10);
-		if (errno || env.tid <= 0) {
-			fprintf(stderr, "Invalid TID: %s\n", arg);
+		ret = split_pidstr(strdup(arg), ",", MAX_TID_NR, env.tids);
+		if (ret) {
+			if (ret == -ENOBUFS)
+				fprintf(stderr, "the number of tid is too big, please "
+					"increase MAX_TID_NR's value and recompile\n");
+			else
+				fprintf(stderr, "invalid TID: %s\n", arg);
+
 			argp_usage(state);
 		}
 		break;
@@ -281,6 +309,30 @@ static void print_map(struct ksyms *ksyms, struct syms_cache *syms_cache,
 	free(ip);
 }
 
+static void print_headers()
+{
+	int i;
+
+	printf("Tracing off-CPU time (us) of");
+
+	if (env.pids[0]) {
+		printf(" PID [");
+		for (i = 0; i < MAX_PID_NR && env.pids[i]; i++)
+			printf("%d%s", env.pids[i], (i < MAX_PID_NR - 1 && env.pids[i + 1]) ? ", " : "]");
+	} else if (env.tids[0]) {
+		printf(" TID [");
+		for (i = 0; i < MAX_TID_NR && env.tids[i]; i++)
+			printf("%d%s", env.tids[i], (i < MAX_TID_NR - 1 && env.tids[i + 1]) ? ", " : "]");
+	} else {
+		printf(" all threads");
+	}
+
+	if (env.duration < 99999999)
+		printf(" for %d secs.\n", env.duration);
+	else
+		printf("... Hit Ctrl-C to end.\n");
+}
+
 int main(int argc, char **argv)
 {
 	static const struct argp argp = {
@@ -291,7 +343,9 @@ int main(int argc, char **argv)
 	struct syms_cache *syms_cache = NULL;
 	struct ksyms *ksyms = NULL;
 	struct offcputime_bpf *obj;
-	int err;
+	int pids_fd, tids_fd;
+	int err, i;
+	__u8 val = 0;
 
 	err = argp_parse(&argp, argc, argv, 0, NULL, NULL);
 	if (err)
@@ -314,14 +368,18 @@ int main(int argc, char **argv)
 	}
 
 	/* initialize global data (filtering options) */
-	obj->rodata->targ_tgid = env.pid;
-	obj->rodata->targ_pid = env.tid;
 	obj->rodata->user_threads_only = env.user_threads_only;
 	obj->rodata->kernel_threads_only = env.kernel_threads_only;
 	obj->rodata->state = env.state;
 	obj->rodata->min_block_ns = env.min_block_time;
 	obj->rodata->max_block_ns = env.max_block_time;
 
+	/* User space PID and TID correspond to TGID and PID in the kernel, respectively */
+	if (env.pids[0])
+		obj->rodata->filter_by_tgid = true;
+	else if (env.tids[0])
+		obj->rodata->filter_by_pid = true;
+
 	bpf_map__set_value_size(obj->maps.stackmap,
 				env.perf_max_stack_depth * sizeof(unsigned long));
 	bpf_map__set_max_entries(obj->maps.stackmap, env.stack_storage_size);
@@ -331,6 +389,28 @@ int main(int argc, char **argv)
 		fprintf(stderr, "failed to load BPF programs\n");
 		goto cleanup;
 	}
+
+	if (env.pids[0]) {
+		/* User pids_fd points to the tgids map in the BPF program */
+		pids_fd = bpf_map__fd(obj->maps.tgids);
+		for (i = 0; i < MAX_PID_NR && env.pids[i]; i++) {
+			if (bpf_map_update_elem(pids_fd, &(env.pids[i]), &val, BPF_ANY) != 0) {
+				fprintf(stderr, "failed to init pids map: %s\n", strerror(errno));
+				goto cleanup;
+			}
+		}
+	}
+	else if (env.tids[0]) {
+		/* User tids_fd points to the pids map in the BPF program */
+		tids_fd = bpf_map__fd(obj->maps.pids);
+		for (i = 0; i < MAX_TID_NR && env.tids[i]; i++) {
+			if (bpf_map_update_elem(tids_fd, &(env.tids[i]), &val, BPF_ANY) != 0) {
+				fprintf(stderr, "failed to init tids map: %s\n", strerror(errno));
+				goto cleanup;
+			}
+		}
+	}
+
 	ksyms = ksyms__load();
 	if (!ksyms) {
 		fprintf(stderr, "failed to load kallsyms\n");
@@ -349,11 +429,8 @@ int main(int argc, char **argv)
 
 	signal(SIGINT, sig_handler);
 
-	printf("Tracing off-CPU time (us)");
-	if (env.duration < 99999999)
-		printf(" for %d secs.\n", env.duration);
-	else
-		printf("... Hit Ctrl-C to end.\n");
+	print_headers();
+
 	/*
 	 * We'll get sleep interrupted when someone presses Ctrl-C (which will
 	 * be "handled" with noop by sig_handler).
diff --git a/libbpf-tools/offcputime.h b/libbpf-tools/offcputime.h
@@ -3,6 +3,8 @@
 #define __OFFCPUTIME_H
 
 #define TASK_COMM_LEN		16
+#define MAX_PID_NR		30
+#define MAX_TID_NR		30
 
 struct key_t {
 	__u32 pid;
