[RFC PATCH v2 5/5] selftests/bpf: Add selftest for THP adjustment

[Yafang Shao <laoar.shao@gmail.com>]

This test case uses a BPF program to enforce the following THP allocation
policy:
- Only the current task is permitted to allocate THP.
- All other tasks are denied.

The expected behavior:
- Before the BPF prog is attached
  No tasks can allocate THP.
- After the BPF prog is attached
  Only the current task can allocate THP.
- Switch to "never" mode after the BPF prog is attached
  THP allocation is not allowed even for the current task.

The result is as follows,
  $ ./test_progs --name="thp_adjust"
  #437     thp_adjust:OK
  Summary: 1/0 PASSED, 0 SKIPPED, 0 FAILED

CONFIG_TRANSPARENT_HUGEPAGE=y is required for this test.

Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
---
 tools/testing/selftests/bpf/config            |   1 +
 .../selftests/bpf/prog_tests/thp_adjust.c     | 175 ++++++++++++++++++
 .../selftests/bpf/progs/test_thp_adjust.c     |  39 ++++
 3 files changed, 215 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/thp_adjust.c
 create mode 100644 tools/testing/selftests/bpf/progs/test_thp_adjust.c

diff --git a/tools/testing/selftests/bpf/config b/tools/testing/selftests/bpf/config
index c378d5d07e02..bb8a8a9d77a2 100644
--- a/tools/testing/selftests/bpf/config
+++ b/tools/testing/selftests/bpf/config
@@ -113,3 +113,4 @@ CONFIG_XDP_SOCKETS=y
 CONFIG_XFRM_INTERFACE=y
 CONFIG_TCP_CONG_DCTCP=y
 CONFIG_TCP_CONG_BBR=y
+CONFIG_TRANSPARENT_HUGEPAGE=y
diff --git a/tools/testing/selftests/bpf/prog_tests/thp_adjust.c b/tools/testing/selftests/bpf/prog_tests/thp_adjust.c
new file mode 100644
index 000000000000..6accd110d8ea
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/thp_adjust.c
@@ -0,0 +1,175 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include <sys/mman.h>
+#include <test_progs.h>
+#include "test_thp_adjust.skel.h"
+
+#define LEN (4 * 1024 * 1024) /* 4MB */
+#define THP_ENABLED_PATH "/sys/kernel/mm/transparent_hugepage/enabled"
+#define SMAPS_PATH "/proc/self/smaps"
+#define ANON_HUGE_PAGES "AnonHugePages:"
+
+static char *thp_addr;
+static char old_mode[32];
+
+int thp_mode_save(void)
+{
+	const char *start, *end;
+	char buf[128];
+	int fd, err;
+	size_t len;
+
+	fd = open(THP_ENABLED_PATH, O_RDONLY);
+	if (fd == -1)
+		return -1;
+
+	err = read(fd, buf, sizeof(buf) - 1);
+	if (err == -1)
+		goto close;
+
+	start = strchr(buf, '[');
+	end = start ? strchr(start, ']') : NULL;
+	if (!start || !end || end <= start) {
+		err = -1;
+		goto close;
+	}
+
+	len = end - start - 1;
+	if (len >= sizeof(old_mode))
+		len = sizeof(old_mode) - 1;
+	strncpy(old_mode, start + 1, len);
+	old_mode[len] = '\0';
+
+close:
+	close(fd);
+	return err;
+}
+
+int thp_set(const char *desired_mode)
+{
+	int fd, err;
+
+	fd = open(THP_ENABLED_PATH, O_RDWR);
+	if (fd == -1)
+		return -1;
+
+	err = write(fd, desired_mode, strlen(desired_mode));
+	close(fd);
+	return err;
+}
+
+int thp_reset(void)
+{
+	int fd, err;
+
+	fd = open(THP_ENABLED_PATH, O_WRONLY);
+	if (fd == -1)
+		return -1;
+
+	err = write(fd, old_mode, strlen(old_mode));
+	close(fd);
+	return err;
+}
+
+int thp_alloc(void)
+{
+	int err, i;
+
+	thp_addr = mmap(NULL, LEN, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANON, -1, 0);
+	if (thp_addr == MAP_FAILED)
+		return -1;
+
+	err = madvise(thp_addr, LEN, MADV_HUGEPAGE);
+	if (err == -1)
+		goto unmap;
+
+	for (i = 0; i < LEN; i += 4096)
+		thp_addr[i] = 1;
+	return 0;
+
+unmap:
+	munmap(thp_addr, LEN);
+	return -1;
+}
+
+void thp_free(void)
+{
+	if (!thp_addr)
+		return;
+	munmap(thp_addr, LEN);
+}
+
+void test_thp_adjust(void)
+{
+	struct bpf_link *fentry_link, *ops_link;
+	struct test_thp_adjust *skel;
+	int err, first_calls;
+
+	if (!ASSERT_NEQ(thp_mode_save(), -1, "THP mode save"))
+		return;
+	if (!ASSERT_GE(thp_set("bpf"), 0, "THP mode set"))
+		return;
+
+	skel = test_thp_adjust__open();
+	if (!ASSERT_OK_PTR(skel, "open"))
+		goto thp_reset;
+
+	skel->bss->target_pid = getpid();
+
+	err = test_thp_adjust__load(skel);
+	if (!ASSERT_OK(err, "load"))
+		goto destroy;
+
+	fentry_link = bpf_program__attach_trace(skel->progs.thp_run);
+	if (!ASSERT_OK_PTR(fentry_link, "attach fentry"))
+		goto destroy;
+
+	if (!ASSERT_NEQ(thp_alloc(), -1, "THP alloc"))
+		goto destroy;
+
+	/* Before attaching struct_ops, THP won't be allocated. */
+	if (!ASSERT_EQ(skel->bss->thp_calls, 0, "THP calls"))
+		goto thp_free;
+
+	if (!ASSERT_EQ(skel->bss->thp_wrong_calls, 0, "THP calls"))
+		goto thp_free;
+
+	thp_free();
+
+	ops_link = bpf_map__attach_struct_ops(skel->maps.thp);
+	if (!ASSERT_OK_PTR(ops_link, "attach struct_ops"))
+		goto destroy;
+
+	if (!ASSERT_NEQ(thp_alloc(), -1, "THP alloc"))
+		goto destroy;
+
+	/* After attaching struct_ops, THP will be allocated. */
+	if (!ASSERT_GT(skel->bss->thp_calls, 0, "THP calls"))
+		goto thp_free;
+
+	first_calls = skel->bss->thp_calls;
+
+	if (!ASSERT_EQ(skel->bss->thp_wrong_calls, 0, "THP calls"))
+		goto thp_free;
+
+	thp_free();
+
+	if (!ASSERT_GE(thp_set("never"), 0, "THP set"))
+		goto destroy;
+
+	if (!ASSERT_NEQ(thp_alloc(), -1, "THP alloc"))
+		goto destroy;
+
+	/* In "never" mode, THP won't be allocated even if the prog is attached. */
+	if (!ASSERT_EQ(skel->bss->thp_calls, first_calls, "THP calls"))
+		goto thp_free;
+
+	ASSERT_EQ(skel->bss->thp_wrong_calls, 0, "THP calls");
+
+thp_free:
+	thp_free();
+destroy:
+	test_thp_adjust__destroy(skel);
+thp_reset:
+	ASSERT_GE(thp_reset(), 0, "THP mode reset");
+}
diff --git a/tools/testing/selftests/bpf/progs/test_thp_adjust.c b/tools/testing/selftests/bpf/progs/test_thp_adjust.c
new file mode 100644
index 000000000000..69135380853c
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/test_thp_adjust.c
@@ -0,0 +1,39 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include "vmlinux.h"
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+
+char _license[] SEC("license") = "GPL";
+
+int target_pid;
+int thp_calls;
+int thp_wrong_calls;
+
+SEC("fentry/do_huge_pmd_anonymous_page")
+int BPF_PROG(thp_run)
+{
+	struct task_struct *current = bpf_get_current_task_btf();
+
+	if (current->pid == target_pid)
+		thp_calls++;
+	else
+		thp_wrong_calls++;
+	return 0;
+}
+
+SEC("struct_ops/thp_bpf_allowable")
+bool BPF_PROG(thp_bpf_allowable)
+{
+	struct task_struct *current = bpf_get_current_task_btf();
+
+	/* Permit the current task to allocate memory using THP. */
+	if (current->pid == target_pid)
+		return true;
+	return false;
+}
+
+SEC(".struct_ops.link")
+struct bpf_thp_ops thp = {
+	.thp_bpf_allowable = (void *)thp_bpf_allowable,
+};
-- 
2.43.5