libbpf-tools/numamove: Add folio support (#4955)

In [0] kernel commit 73eab3ca481e ("mm: migrate: convert migrate_misplaced_page()
to migrate_misplaced_folio()") convert migrate_misplaced_page() to
migrate_misplaced_folio().

Kernel version:

    $ git describe 73eab3ca481e5be0f1fd8140365d604482f84ee1
    v6.6-rc4-109-g73eab3ca481e

[0] torvalds/linux@73eab3ca481e

Signed-off-by: Rong Tao <[email protected]>

libbpf-tools/numamove.bpf.c

@@ -14,7 +14,7 @@ struct {
 __u64 latency = 0;
 __u64 num = 0;
 
-static int __migrate_misplaced_page(void)
+static int __migrate_misplaced(void)
 {
 	u32 pid = bpf_get_current_pid_tgid();
 	u64 ts = bpf_ktime_get_ns();
@@ -26,16 +26,28 @@ static int __migrate_misplaced_page(void)
 SEC("fentry/migrate_misplaced_page")
 int BPF_PROG(fentry_migrate_misplaced_page)
 {
-	return __migrate_misplaced_page();
+	return __migrate_misplaced();
+}
+
+SEC("fentry/migrate_misplaced_folio")
+int BPF_PROG(fentry_migrate_misplaced_folio)
+{
+	return __migrate_misplaced();
 }
 
 SEC("kprobe/migrate_misplaced_page")
 int BPF_PROG(kprobe_migrate_misplaced_page)
 {
-	return __migrate_misplaced_page();
+	return __migrate_misplaced();
+}
+
+SEC("kprobe/migrate_misplaced_folio")
+int BPF_PROG(kprobe_migrate_misplaced_folio)
+{
+	return __migrate_misplaced();
 }
 
-static int __migrate_misplaced_page_exit(void)
+static int __migrate_misplaced_exit(void)
 {
 	u32 pid = bpf_get_current_pid_tgid();
 	u64 *tsp, ts = bpf_ktime_get_ns();
@@ -58,13 +70,25 @@ static int __migrate_misplaced_page_exit(void)
 SEC("fexit/migrate_misplaced_page")
 int BPF_PROG(fexit_migrate_misplaced_page_exit)
 {
-	return __migrate_misplaced_page_exit();
+	return __migrate_misplaced_exit();
+}
+
+SEC("fexit/migrate_misplaced_folio")
+int BPF_PROG(fexit_migrate_misplaced_folio_exit)
+{
+	return __migrate_misplaced_exit();
 }
 
 SEC("kretprobe/migrate_misplaced_page")
 int BPF_PROG(kretprobe_migrate_misplaced_page_exit)
 {
-	return __migrate_misplaced_page_exit();
+	return __migrate_misplaced_exit();
+}
+
+SEC("kretprobe/migrate_misplaced_folio")
+int BPF_PROG(kretprobe_migrate_misplaced_folio_exit)
+{
+	return __migrate_misplaced_exit();
 }
 
 char LICENSE[] SEC("license") = "GPL";

libbpf-tools/numamove.c

@@ -4,8 +4,10 @@
 // Based on numamove(8) from BPF-Perf-Tools-Book by Brendan Gregg.
 //  8-Jun-2020   Wenbo Zhang   Created this.
 // 30-Jan-2023   Rong Tao      Use fentry_can_attach() to decide use fentry/kprobe.
+// 06-Apr-2024   Rong Tao      Support migrate_misplaced_folio()
 #include <argp.h>
 #include <signal.h>
+#include <stdbool.h>
 #include <stdio.h>
 #include <unistd.h>
 #include <time.h>
@@ -76,6 +78,7 @@ int main(int argc, char **argv)
 	char ts[32];
 	time_t t;
 	int err;
+	bool use_folio, use_fentry;
 
 	err = argp_parse(&argp, argc, argv, 0, NULL, NULL);
 	if (err)
@@ -95,14 +98,33 @@ int main(int argc, char **argv)
 	}
 
 	/* It fallbacks to kprobes when kernel does not support fentry. */
-	if (fentry_can_attach("migrate_misplaced_page", NULL)) {
-		bpf_program__set_autoload(obj->progs.kprobe_migrate_misplaced_page, false);
-		bpf_program__set_autoload(obj->progs.kretprobe_migrate_misplaced_page_exit, false);
+	if (fentry_can_attach("migrate_misplaced_folio", NULL)) {
+		use_fentry = true;
+		use_folio = true;
+	} else if (kprobe_exists("migrate_misplaced_folio")) {
+		use_fentry = false;
+		use_folio = true;
+	} else if (fentry_can_attach("migrate_misplaced_page", NULL)) {
+		use_fentry = true;
+		use_folio = false;
+	} else if (kprobe_exists("migrate_misplaced_page")) {
+		use_fentry = false;
+		use_folio = false;
 	} else {
-		bpf_program__set_autoload(obj->progs.fentry_migrate_misplaced_page, false);
-		bpf_program__set_autoload(obj->progs.fexit_migrate_misplaced_page_exit, false);
+		fprintf(stderr, "can't found any fentry/kprobe of migrate misplaced folio/page\n");
+		return 1;
 	}
 
+	bpf_program__set_autoload(obj->progs.fentry_migrate_misplaced_folio, (use_fentry && use_folio));
+	bpf_program__set_autoload(obj->progs.fexit_migrate_misplaced_folio_exit, (use_fentry && use_folio));
+	bpf_program__set_autoload(obj->progs.kprobe_migrate_misplaced_folio, (!use_fentry && use_folio));
+	bpf_program__set_autoload(obj->progs.kretprobe_migrate_misplaced_folio_exit, (!use_fentry && use_folio));
+
+	bpf_program__set_autoload(obj->progs.fentry_migrate_misplaced_page, (use_fentry && !use_folio));
+	bpf_program__set_autoload(obj->progs.fexit_migrate_misplaced_page_exit, (use_fentry && !use_folio));
+	bpf_program__set_autoload(obj->progs.kprobe_migrate_misplaced_page, (!use_fentry && !use_folio));
+	bpf_program__set_autoload(obj->progs.kretprobe_migrate_misplaced_page_exit, (!use_fentry && !use_folio));
+
 	err = numamove_bpf__load(obj);
 	if (err) {
 		fprintf(stderr, "failed to load BPF skelect: %d\n", err);