From 957569f81f486d88c9c5d8e6436d76048b5753bd Mon Sep 17 00:00:00 2001
From: woodpenker <woodpenker5@gmail.com>
Date: Sun, 13 Nov 2022 17:05:32 +0800
Subject: [PATCH] libbpf-tools/cpudist: Allow cpudist to run on old kernels

---
 libbpf-tools/cpudist.bpf.c | 25 ++++++++++++++++++-------
 libbpf-tools/cpudist.c     |  5 +++++
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/libbpf-tools/cpudist.bpf.c b/libbpf-tools/cpudist.bpf.c
index a1d254c7852b..0901fcae5405 100644
--- a/libbpf-tools/cpudist.bpf.c
+++ b/libbpf-tools/cpudist.bpf.c
@@ -71,8 +71,7 @@ static __always_inline void update_hist(struct task_struct *task,
 		histp = bpf_map_lookup_elem(&hists, &id);
 		if (!histp)
 			return;
-		bpf_probe_read_kernel_str(&histp->comm, sizeof(task->comm),
-					  task->comm);
+		BPF_CORE_READ_STR_INTO(&histp->comm, task, comm);
 	}
 	delta = ts - *tsp;
 	if (targ_ms)
@@ -85,12 +84,10 @@ static __always_inline void update_hist(struct task_struct *task,
 	__sync_fetch_and_add(&histp->slots[slot], 1);
 }
 
-SEC("tp_btf/sched_switch")
-int BPF_PROG(sched_switch, bool preempt, struct task_struct *prev,
-	     struct task_struct *next)
+static int handle_switch(struct task_struct *prev, struct task_struct *next)
 {
-	u32 prev_tgid = prev->tgid, prev_pid = prev->pid;
-	u32 tgid = next->tgid, pid = next->pid;
+	u32 prev_tgid = BPF_CORE_READ(prev, tgid), prev_pid = BPF_CORE_READ(prev, pid);
+	u32 tgid = BPF_CORE_READ(next, tgid), pid = BPF_CORE_READ(next, pid);
 	u64 ts = bpf_ktime_get_ns();
 
 	if (filter_cg && !bpf_current_task_under_cgroup(&cgroup_map, 0))
@@ -107,4 +104,18 @@ int BPF_PROG(sched_switch, bool preempt, struct task_struct *prev,
 	return 0;
 }
 
+SEC("tp_btf/sched_switch")
+int BPF_PROG(sched_switch_btf, bool preempt, struct task_struct *prev,
+	     struct task_struct *next)
+{
+	return handle_switch(prev, next);
+}
+
+SEC("raw_tp/sched_switch")
+int BPF_PROG(sched_switch_tp, bool preempt, struct task_struct *prev,
+	     struct task_struct *next)
+{
+	return handle_switch(prev, next);
+}
+
 char LICENSE[] SEC("license") = "GPL";
diff --git a/libbpf-tools/cpudist.c b/libbpf-tools/cpudist.c
index f24d845f4053..72682c5c294e 100644
--- a/libbpf-tools/cpudist.c
+++ b/libbpf-tools/cpudist.c
@@ -217,6 +217,11 @@ int main(int argc, char **argv)
 		return 1;
 	}
 
+	if (probe_tp_btf("sched_switch"))
+		bpf_program__set_autoload(obj->progs.sched_switch_tp, false);
+	else
+		bpf_program__set_autoload(obj->progs.sched_switch_btf, false);
+
 	/* initialize global data (filtering options) */
 	obj->rodata->filter_cg = env.cg;
 	obj->rodata->targ_per_process = env.per_process;