[PATCH RFC v3 2/2] selftests/bpf: Test basic workflow of task local data

[Amery Hung <ameryhung@gmail.com>]

Test the workflow of task local data. A user space program first declares
task local data using two different APIs. As the test starts, it calls
bpf_tld_thread_init() for every new thread that would access the
storage. Then, values can be accessed directly. The user space triggers
two bpf programs: prog_init and prog_main. prog_init simulates a
sched_ext_ops::init_task, which runs only once for every new task. It
caches the offsets of values of the task. prog_main represents bpf
programs for normal operation. It reads the task local data and write the
result to global variables for verification.

The user space program will launch 32 threads to make sure not only
umetadata, but thread-specific udata and udata_start are handled
correctly. It is verified by writing values in user space, reading
them the bpf program and checking that they match. Also make sure
the data are indeed thread-specific. Finally, a large task local data is
declared to see if the declaration API prevents it from spanning across
two pages.

Signed-off-by: Amery Hung <ameryhung@gmail.com>
---
 .../bpf/prog_tests/test_task_local_data.c     | 156 ++++++++++++++++++
 .../bpf/progs/test_task_local_data_basic.c    |  78 +++++++++
 .../selftests/bpf/task_local_data_common.h    |   8 +
 3 files changed, 242 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/test_task_local_data.c
 create mode 100644 tools/testing/selftests/bpf/progs/test_task_local_data_basic.c

diff --git a/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c b/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c
new file mode 100644
index 000000000000..5754687026f3
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/test_task_local_data.c
@@ -0,0 +1,156 @@
+#include <pthread.h>
+
+#include <test_progs.h>
+#include <bpf/btf.h>
+
+#include "task_local_data.h"
+#include "test_task_local_data_basic.skel.h"
+
+#define TEST_THREAD_NUM 32
+
+/* Used to declare a large tasl local data below to see if bpf_tld_type_var() prevents
+ * a value from crossing the page boundary
+ */
+struct dummy {
+	char data[1000];
+};
+
+/* Declare task local data */
+bpf_tld_type_var(int, value1);
+bpf_tld_type_var(struct test_struct, value2);
+bpf_tld_type_var(struct dummy, dummy);
+bpf_tld_key_type_var("test_basic_value3", int, value3);
+bpf_tld_key_type_var("test_basic_value4", struct test_struct, value4);
+
+/* Serialize access to bpf program's global variables */
+static pthread_mutex_t global_mutex;
+
+static void run_prog_init(struct test_task_local_data_basic *skel, int tid)
+{
+	skel->bss->target_tid = tid;
+	(void)syscall(__NR_getuid);
+	skel->bss->target_tid = -1;
+}
+
+static void run_prog_main(struct test_task_local_data_basic *skel, int tid)
+{
+	skel->bss->target_tid = tid;
+	(void)syscall(__NR_gettid);
+	skel->bss->target_tid = -1;
+}
+
+void *test_task_local_data_basic_thread(void *arg)
+{
+	struct test_task_local_data_basic *skel = (struct test_task_local_data_basic *)arg;
+	int err, tid;
+
+	tid = gettid();
+
+	err = bpf_tld_thread_init();
+	if (!ASSERT_OK(err, "bpf_tld_thread_init"))
+		return NULL;
+
+	value1 = tid + 0;
+	value2.a = tid + 1;
+	value2.b = tid + 2;
+	value2.c = tid + 3;
+	value2.d = tid + 4;
+	value3 = tid + 5;
+	value4.a = tid + 6;
+	value4.b = tid + 7;
+	value4.c = tid + 8;
+	value4.d = tid + 9;
+
+	pthread_mutex_lock(&global_mutex);
+	/* Simulate an initialization bpf prog that runs once for every new task.
+	 * The program caches data offsets for subsequent bpf programs
+	 */
+	run_prog_init(skel, tid);
+	/* Run main prog that lookup task local data and save to global variables */
+	run_prog_main(skel, tid);
+	ASSERT_EQ(skel->bss->test_value1, tid + 0, "bpf_tld_lookup value1");
+	ASSERT_EQ(skel->bss->test_value2.a, tid + 1, "bpf_tld_lookup value2.a");
+	ASSERT_EQ(skel->bss->test_value2.b, tid + 2, "bpf_tld_lookup value2.b");
+	ASSERT_EQ(skel->bss->test_value2.c, tid + 3, "bpf_tld_lookup value2.c");
+	ASSERT_EQ(skel->bss->test_value2.d, tid + 4, "bpf_tld_lookup value2.d");
+	ASSERT_EQ(skel->bss->test_value3, tid + 5, "bpf_tld_lookup value3");
+	ASSERT_EQ(skel->bss->test_value4.a, tid + 6, "bpf_tld_lookup value4.a");
+	ASSERT_EQ(skel->bss->test_value4.b, tid + 7, "bpf_tld_lookup value4.b");
+	ASSERT_EQ(skel->bss->test_value4.c, tid + 8, "bpf_tld_lookup value4.c");
+	ASSERT_EQ(skel->bss->test_value4.d, tid + 9, "bpf_tld_lookup value4.d");
+	pthread_mutex_unlock(&global_mutex);
+
+	/* Make sure valueX are indeed local to threads */
+	ASSERT_EQ(value1, tid + 0, "value1");
+	ASSERT_EQ(value2.a, tid + 1, "value2.a");
+	ASSERT_EQ(value2.b, tid + 2, "value2.b");
+	ASSERT_EQ(value2.c, tid + 3, "value2.c");
+	ASSERT_EQ(value2.d, tid + 4, "value2.d");
+	ASSERT_EQ(value3, tid + 5, "value3");
+	ASSERT_EQ(value4.a, tid + 6, "value4.a");
+	ASSERT_EQ(value4.b, tid + 7, "value4.b");
+	ASSERT_EQ(value4.c, tid + 8, "value4.c");
+	ASSERT_EQ(value4.d, tid + 9, "value4.d");
+
+	value1 = tid + 9;
+	value2.a = tid + 8;
+	value2.b = tid + 7;
+	value2.c = tid + 6;
+	value2.d = tid + 5;
+	value3 = tid + 4;
+	value4.a = tid + 3;
+	value4.b = tid + 2;
+	value4.c = tid + 1;
+	value4.d = tid + 0;
+
+	/* Run main prog again */
+	pthread_mutex_lock(&global_mutex);
+	run_prog_main(skel, tid);
+	ASSERT_EQ(skel->bss->test_value1, tid + 9, "bpf_tld_lookup value1");
+	ASSERT_EQ(skel->bss->test_value2.a, tid + 8, "bpf_tld_lookup value2.a");
+	ASSERT_EQ(skel->bss->test_value2.b, tid + 7, "bpf_tld_lookup value2.b");
+	ASSERT_EQ(skel->bss->test_value2.c, tid + 6, "bpf_tld_lookup value2.c");
+	ASSERT_EQ(skel->bss->test_value2.d, tid + 5, "bpf_tld_lookup value2.d");
+	ASSERT_EQ(skel->bss->test_value3, tid + 4, "bpf_tld_lookup value3");
+	ASSERT_EQ(skel->bss->test_value4.a, tid + 3, "bpf_tld_lookup value4.a");
+	ASSERT_EQ(skel->bss->test_value4.b, tid + 2, "bpf_tld_lookup value4.b");
+	ASSERT_EQ(skel->bss->test_value4.c, tid + 1, "bpf_tld_lookup value4.c");
+	ASSERT_EQ(skel->bss->test_value4.d, tid + 0, "bpf_tld_lookup value4.d");
+	pthread_mutex_unlock(&global_mutex);
+
+	pthread_exit(NULL);
+}
+
+static void test_task_local_data_basic(void)
+{
+	struct test_task_local_data_basic *skel;
+	pthread_t thread[TEST_THREAD_NUM];
+	int i, err;
+
+	ASSERT_OK(pthread_mutex_init(&global_mutex, NULL), "pthread_mutex_init");
+
+	skel = test_task_local_data_basic__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "skel_open_and_load"))
+		return;
+
+	err = test_task_local_data_basic__attach(skel);
+	if (!ASSERT_OK(err, "skel_attach"))
+		goto out;
+
+	for (i = 0; i < TEST_THREAD_NUM; i++) {
+		err = pthread_create(&thread[i], NULL, test_task_local_data_basic_thread, skel);
+		if (!ASSERT_OK(err, "pthread_create"))
+			goto out;
+	}
+
+	for (i = 0; i < TEST_THREAD_NUM; i++)
+		pthread_join(thread[i], NULL);
+out:
+	unlink(TASK_LOCAL_DATA_MAP_PIN_PATH);
+}
+
+void test_task_local_data(void)
+{
+	if (test__start_subtest("task_local_data_basic"))
+		test_task_local_data_basic();
+}
diff --git a/tools/testing/selftests/bpf/progs/test_task_local_data_basic.c b/tools/testing/selftests/bpf/progs/test_task_local_data_basic.c
new file mode 100644
index 000000000000..345d7c6e37de
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/test_task_local_data_basic.c
@@ -0,0 +1,78 @@
+#include <vmlinux.h>
+#include <bpf/bpf_helpers.h>
+
+#include "task_local_data.h"
+
+struct task_local_data_offsets {
+	short value1;
+	short value2;
+	short test_basic_value3;
+	short test_basic_value4;
+};
+
+pid_t target_tid = 0;
+int test_value1 = 0;
+struct test_struct test_value2;
+int test_value3 = 0;
+struct test_struct test_value4;
+
+SEC("tp/syscalls/sys_enter_getuid")
+int prog_init(void *ctx)
+{
+	struct bpf_task_local_data tld;
+	struct task_struct *task;
+	int err;
+
+	task = bpf_get_current_task_btf();
+	if (task->pid != target_tid)
+		return 0;
+
+	err = bpf_tld_init(task, &tld);
+	if (err)
+		return 0;
+
+	bpf_tld_init_var(&tld, value1);
+	bpf_tld_init_var(&tld, value2);
+	bpf_tld_init_var(&tld, test_basic_value3);
+	bpf_tld_init_var(&tld, test_basic_value4);
+
+	return 0;
+}
+
+SEC("tp/syscalls/sys_enter_gettid")
+int prog_main(void *ctx)
+{
+	struct bpf_task_local_data tld;
+	struct test_struct *struct_p;
+	struct task_struct *task;
+	int err, *int_p;
+
+	task = bpf_get_current_task_btf();
+	if (task->pid != target_tid)
+		return 0;
+
+	err = bpf_tld_init(task, &tld);
+	if (err)
+		return 0;
+
+	int_p = bpf_tld_lookup(&tld, value1, sizeof(int));
+	if (int_p)
+		test_value1 = *int_p;
+
+	struct_p = bpf_tld_lookup(&tld, value2, sizeof(struct test_struct));
+	if (struct_p)
+		test_value2 = *struct_p;
+
+	int_p = bpf_tld_lookup(&tld, test_basic_value3, sizeof(int));
+	if (int_p)
+		test_value3 = *int_p;
+
+	struct_p = bpf_tld_lookup(&tld, test_basic_value4, sizeof(struct test_struct));
+	if (struct_p)
+		test_value4 = *struct_p;
+
+	return 0;
+}
+
+char _license[] SEC("license") = "GPL";
+
diff --git a/tools/testing/selftests/bpf/task_local_data_common.h b/tools/testing/selftests/bpf/task_local_data_common.h
index 2a0bb724c77c..ad99c66d3305 100644
--- a/tools/testing/selftests/bpf/task_local_data_common.h
+++ b/tools/testing/selftests/bpf/task_local_data_common.h
@@ -38,4 +38,12 @@ struct task_local_data_map_value {
 	short udata_start;
 };
 
+/* test specific */
+struct test_struct {
+	unsigned long a;
+	unsigned long b;
+	unsigned long c;
+	unsigned long d;
+};
+
 #endif
-- 
2.47.1