[PATCH bpf-next v1 1/1] selftests/bpf: Trace bpf_local_storage_update to debug flaky local storage tests

[Amery Hung <ameryhung@gmail.com>]

task_local_storage/sys_enter_exit and cgrp_local_storage/
cgroup_iter_sleepable occasionally fail in CI possibly because
bpf_{task,cgrp}_storage_get() returns NULL. Add a fexit probe on
bpf_local_storage_update() to capture the actual error code when this
happens. It will allow us to tell if it is trylock failure in
kmalloc_nolock(), timeout/deadlock in rqspinlock or something else.

Signed-off-by: Amery Hung <ameryhung@gmail.com>
---
 .../bpf/prog_tests/cgrp_local_storage.c       | 15 +++++++++++++--
 .../bpf/prog_tests/task_local_storage.c       |  1 +
 .../selftests/bpf/progs/cgrp_ls_sleepable.c   | 18 ++++++++++++++++++
 .../selftests/bpf/progs/task_local_storage.c  | 19 +++++++++++++++++++
 4 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/tools/testing/selftests/bpf/prog_tests/cgrp_local_storage.c b/tools/testing/selftests/bpf/prog_tests/cgrp_local_storage.c
index 478a77cb67e6..c4398ccf3493 100644
--- a/tools/testing/selftests/bpf/prog_tests/cgrp_local_storage.c
+++ b/tools/testing/selftests/bpf/prog_tests/cgrp_local_storage.c
@@ -176,7 +176,7 @@ static void test_cgroup_iter_sleepable(int cgroup_fd, __u64 cgroup_id)
 	DECLARE_LIBBPF_OPTS(bpf_iter_attach_opts, opts);
 	union bpf_iter_link_info linfo;
 	struct cgrp_ls_sleepable *skel;
-	struct bpf_link *link;
+	struct bpf_link *link, *fexit_link;
 	int err, iter_fd;
 	char buf[16];
 
@@ -200,16 +200,27 @@ static void test_cgroup_iter_sleepable(int cgroup_fd, __u64 cgroup_id)
 	if (!ASSERT_OK_PTR(link, "attach_iter"))
 		goto out;
 
+	fexit_link = bpf_program__attach(skel->progs.fexit_update);
+	if (!ASSERT_OK_PTR(fexit_link, "attach_fexit"))
+		goto out_link;
+
 	iter_fd = bpf_iter_create(bpf_link__fd(link));
 	if (!ASSERT_GE(iter_fd, 0, "iter_create"))
-		goto out_link;
+		goto out_fexit_link;
+
+	skel->bss->target_pid = sys_gettid();
 
 	/* trigger the program run */
 	(void)read(iter_fd, buf, sizeof(buf));
 
+	skel->bss->target_pid = 0;
+
+	ASSERT_EQ(skel->bss->update_err, 0, "update_err");
 	ASSERT_EQ(skel->bss->cgroup_id, cgroup_id, "cgroup_id");
 
 	close(iter_fd);
+out_fexit_link:
+	bpf_link__destroy(fexit_link);
 out_link:
 	bpf_link__destroy(link);
 out:
diff --git a/tools/testing/selftests/bpf/prog_tests/task_local_storage.c b/tools/testing/selftests/bpf/prog_tests/task_local_storage.c
index 1b26c12f255a..5b2b56cc3a4f 100644
--- a/tools/testing/selftests/bpf/prog_tests/task_local_storage.c
+++ b/tools/testing/selftests/bpf/prog_tests/task_local_storage.c
@@ -47,6 +47,7 @@ static void test_sys_enter_exit(void)
 	skel->bss->target_pid = 0;
 
 	/* 2x gettid syscalls */
+	ASSERT_EQ(skel->bss->update_err, 0, "update_err");
 	ASSERT_EQ(skel->bss->enter_cnt, 2, "enter_cnt");
 	ASSERT_EQ(skel->bss->exit_cnt, 2, "exit_cnt");
 	ASSERT_EQ(skel->bss->mismatch_cnt, 0, "mismatch_cnt");
diff --git a/tools/testing/selftests/bpf/progs/cgrp_ls_sleepable.c b/tools/testing/selftests/bpf/progs/cgrp_ls_sleepable.c
index a2de95f85648..37bd6b03ba01 100644
--- a/tools/testing/selftests/bpf/progs/cgrp_ls_sleepable.c
+++ b/tools/testing/selftests/bpf/progs/cgrp_ls_sleepable.c
@@ -4,6 +4,7 @@
 #include <bpf/bpf_helpers.h>
 #include <bpf/bpf_tracing.h>
 #include "bpf_misc.h"
+#include "err.h"
 
 char _license[] SEC("license") = "GPL";
 
@@ -16,6 +17,7 @@ struct {
 
 __s32 target_pid;
 __u64 cgroup_id;
+long update_err;
 int target_hid;
 bool is_cgroup1;
 
@@ -123,3 +125,19 @@ int yes_rcu_lock(void *ctx)
 	bpf_rcu_read_unlock();
 	return 0;
 }
+
+SEC("fexit/bpf_local_storage_update")
+int BPF_PROG(fexit_update, void *owner, struct bpf_local_storage_map *smap,
+	     void *value, u64 map_flags, bool swap_uptrs,
+	     struct bpf_local_storage_data *ret)
+{
+	struct task_struct *task = bpf_get_current_task_btf();
+
+	if (task->pid != target_pid)
+		return 0;
+
+	if (IS_ERR_VALUE(ret))
+		update_err = PTR_ERR(ret);
+
+	return 0;
+}
diff --git a/tools/testing/selftests/bpf/progs/task_local_storage.c b/tools/testing/selftests/bpf/progs/task_local_storage.c
index 80a0a20db88d..34fa3d6451d2 100644
--- a/tools/testing/selftests/bpf/progs/task_local_storage.c
+++ b/tools/testing/selftests/bpf/progs/task_local_storage.c
@@ -14,12 +14,15 @@ struct {
 	__type(value, long);
 } enter_id SEC(".maps");
 
+#include "err.h"
+
 #define MAGIC_VALUE 0xabcd1234
 
 pid_t target_pid = 0;
 int mismatch_cnt = 0;
 int enter_cnt = 0;
 int exit_cnt = 0;
+long update_err = 0;
 
 SEC("tp_btf/sys_enter")
 int BPF_PROG(on_enter, struct pt_regs *regs, long id)
@@ -62,3 +65,19 @@ int BPF_PROG(on_exit, struct pt_regs *regs, long id)
 		__sync_fetch_and_add(&mismatch_cnt, 1);
 	return 0;
 }
+
+SEC("fexit/bpf_local_storage_update")
+int BPF_PROG(fexit_update, void *owner, struct bpf_local_storage_map *smap,
+	     void *value, u64 map_flags, bool swap_uptrs,
+	     struct bpf_local_storage_data *ret)
+{
+	struct task_struct *task = bpf_get_current_task_btf();
+
+	if (task->pid != target_pid)
+		return 0;
+
+	if (IS_ERR_VALUE(ret))
+		update_err = PTR_ERR(ret);
+
+	return 0;
+}
-- 
2.52.0