From 92bf318448f2a5fa1118244bae8d977636f5e52c Mon Sep 17 00:00:00 2001
From: zhanghan2021 <zhanghan@kylinos.cn>
Date: Wed, 21 Jun 2023 15:19:24 +0800
Subject: [PATCH] Add the main function function function of numamove

---
 observation/src/numamove/numamove.bpf.c |  35 ++++++++
 observation/src/numamove/numamove.c     | 107 ++++++++++++++++++++++++
 2 files changed, 142 insertions(+)

diff --git a/observation/src/numamove/numamove.bpf.c b/observation/src/numamove/numamove.bpf.c
index de2b1633..18a5c2ef 100644
--- a/observation/src/numamove/numamove.bpf.c
+++ b/observation/src/numamove/numamove.bpf.c
@@ -54,6 +54,41 @@ int BPF_KPROBE(kprobe_migrate_misplaced_page)
 
 static u64 zero;
 
+static int __migrate_misplaced_page_exit(void)
+{
+    pid_t pid = (pid_t)bpf_get_current_pid_tgid();
+    s64 delta;
+    u64 *tsp, *value;
+    u32 key = 0;
+
+    tsp = bpf_map_lookup_elem(&start, &pid);
+    if (!tsp)
+        return 0;
+
+    delta = (s64)(bpf_ktime_get_ns() - *tsp);
+    if (delta < 0)
+        goto cleanup;
+
+#if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 7, 0)
+    __sync_fetch_and_add(&latency, delta / 1000000UL);
+    __sync_fetch_and_add(&num, 1);
+#else
+    value = bpf_map_lookup_or_try_init(&latency_map, &key, &zero);
+    if (!value)
+        goto cleanup;
+    __sync_fetch_and_add(value, delta / 1000000UL);
+
+    value = bpf_map_lookup_or_try_init(&num_map, &key, &zero);
+    if (!value)
+        goto cleanup;
+    __sync_fetch_and_add(value, 1);
+#endif
+
+cleanup:
+    bpf_map_delete_elem(&start, &pid);
+    return 0;
+}
+
 SEC("fexit/migrate_misplaced_page")
 int BPF_PROG(fexit_migrate_misplaced_page)
 {
diff --git a/observation/src/numamove/numamove.c b/observation/src/numamove/numamove.c
index 6e0740a7..0a6707af 100644
--- a/observation/src/numamove/numamove.c
+++ b/observation/src/numamove/numamove.c
@@ -57,4 +57,111 @@ static int libbpf_print_fn(enum libbpf_print_level level, const char *format,
 static void sig_handler(int sig)
 {
 	exiting = 1;
+}
+
+int main(int argc, char *argv[])
+{
+	static const struct argp argp = {
+		.options = opts,
+		.parser = parse_arg,
+		.doc = argp_program_doc,
+	};
+
+	struct numamove_bpf *bpf_obj;
+	int err;
+	bool has_bss;
+
+	err = argp_parse(&argp, argc, argv, 0, NULL, NULL);
+	if (err)
+		return err;
+
+	if (numa_max_node() == 0)
+	{
+		printf("We only have one NUMA Node, so we don't need to track NUMA moving.\n");
+		return 0;
+	}
+
+	if (!bpf_is_root())
+		return 1;
+
+	libbpf_set_print(libbpf_print_fn);
+
+	bpf_obj = numamove_bpf__open();
+	if (!bpf_obj)
+	{
+		warning("Failed to open BPF object\n");
+		return 1;
+	}
+
+	/* It's fallback to kprobe when kernel does not support fentry. */
+	if (fentry_can_attach("migrate_misplaced_page", NULL))
+	{
+		bpf_program__set_autoload(bpf_obj->progs.kprobe_migrate_misplaced_page, false);
+		bpf_program__set_autoload(bpf_obj->progs.kretprobe_migrate_misplaced_page, false);
+	}
+	else
+	{
+		bpf_program__set_autoload(bpf_obj->progs.fentry_migrate_misplaced_page, false);
+		bpf_program__set_autoload(bpf_obj->progs.fexit_migrate_misplaced_page, false);
+	}
+
+	err = numamove_bpf__load(bpf_obj);
+	if (err)
+	{
+		warning("Failed to load BPF object: %d\n", err);
+		goto cleanup;
+	}
+
+	has_bss = bpf_obj->bss;
+	if (!has_bss)
+		warning("Memory-mapping BPF maps is supported starting from Linux 5.7, please upgrade.\n\n");
+
+	err = numamove_bpf__attach(bpf_obj);
+	if (err)
+	{
+		warning("Failed to attach BPF programs: %d\n", err);
+		goto cleanup;
+	}
+
+	if (signal(SIGINT, sig_handler) == SIG_ERR)
+	{
+		warning("Failed to set signal hander\n");
+		err = 1;
+		goto cleanup;
+	}
+
+	printf("%-10s %19s %18s\n", "TIME", "NUMA_migrations_num", "NUMA_migrateions_ms");
+	while (!exiting)
+	{
+		char ts[32];
+
+		sleep(1);
+
+		strftime_now(ts, sizeof(ts), "%H:%M:%S");
+		if (has_bss)
+		{
+			printf("%-10s %18lld %18lld\n", ts,
+				   __atomic_exchange_n(&bpf_obj->bss->num, 0, __ATOMIC_RELAXED),
+				   __atomic_exchange_n(&bpf_obj->bss->latency, 0, __ATOMIC_RELAXED));
+		}
+		else
+		{
+			int num_fd = bpf_map__fd(bpf_obj->maps.num_map);
+			int latency_fd = bpf_map__fd(bpf_obj->maps.latency_map);
+			int key = 0, value = 0;
+			__u64 num = 0, latency = 0;
+
+			bpf_map_lookup_elem(num_fd, &key, &num);
+			bpf_map_lookup_elem(latency_fd, &key, &latency);
+
+			printf("%-10s %19lld %18lld\n", ts, num, latency);
+			bpf_map_update_elem(num_fd, &key, &value, BPF_ANY);
+			bpf_map_update_elem(latency_fd, &key, &value, BPF_ANY);
+		}
+	}
+
+cleanup:
+	numamove_bpf__destroy(bpf_obj);
+
+	return err != 0;
 }
\ No newline at end of file
-- 
Gitee