[PATCH bpf-next v2] selftests/bpf: Wait for wq callback before checking ok_sleepable

[Jonathan Bastari Thamrin via B4 Relay <devnull+thamrinjonathan.gmail.com@kernel.org>]

From: Jonathan Bastari Thamrin <thamrinjonathan@gmail.com>

wq test fails with ok_sleepable reading 0 instead of 2 as workqueue
callback hasn't executed before check.

This happens as callback runs later on a kworker and 50us isn't always
enough.

Change from 50us sleep to waiting for callback by checking wq_cb_run.
wq_cb_sleepable() now increments wq_cb_run after setting ok_sleepable
and test loops with sched_yield() until wq_cb_run is set before
checking.

Before (usleep(50)): 3 fails out of 5000 runs
Diagnosing issue (usleep(0)): 23 out of 100 runs
After: 0 out of 2000 runs
Tested using vmtest on x86_64, test pinned to one CPU with `yes` load

Link: https://github.com/kernel-patches/vmtest/issues/455
Signed-off-by: Jonathan Bastari Thamrin <thamrinjonathan@gmail.com>
---
Changes in v2:
- Wait for callback via a wq_cb_run flag instead of polling with a
  timeout, similar to test_struct_ops_assoc.c (Kumar)
- Link to v1: https://lore.kernel.org/r/20261005-wq-poll-fix-v1-1-99ce4bac59f4@gmail.com
---
 tools/testing/selftests/bpf/prog_tests/wq.c |  5 +++--
 tools/testing/selftests/bpf/progs/wq.c      | 10 ++++++----
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/tools/testing/selftests/bpf/prog_tests/wq.c b/tools/testing/selftests/bpf/prog_tests/wq.c
index 97926e35aaf2..95ed4b69c0f5 100644
--- a/tools/testing/selftests/bpf/prog_tests/wq.c
+++ b/tools/testing/selftests/bpf/prog_tests/wq.c
@@ -28,8 +28,9 @@ void serial_test_wq(void)
 	ASSERT_OK(err, "test_run");
 	ASSERT_EQ(topts.retval, 0, "test_run");
 
-	usleep(50); /* 10 usecs should be enough, but give it extra */
-
+	/* wait for wq callback to finish */
+	while (!READ_ONCE(wq_skel->bss->wq_cb_run))
+		sched_yield();
 	ASSERT_EQ(wq_skel->bss->ok_sleepable, (1 << 1), "ok_sleepable");
 clean_up:
 	wq__destroy(wq_skel);
diff --git a/tools/testing/selftests/bpf/progs/wq.c b/tools/testing/selftests/bpf/progs/wq.c
index 25be2cd9d42c..8148db35effe 100644
--- a/tools/testing/selftests/bpf/progs/wq.c
+++ b/tools/testing/selftests/bpf/progs/wq.c
@@ -52,6 +52,7 @@ struct {
 
 __u32 ok;
 __u32 ok_sleepable;
+__u32 wq_cb_run;
 
 static int test_elem_callback(void *map, int *key,
 		int (callback_fn)(void *map, int *key, void *value))
@@ -130,11 +131,12 @@ static int wq_cb_sleepable(void *map, int *key, void *value)
 	struct elem *data = (struct elem *)value;
 	int offset = data->ok_offset;
 
-	if (*key != offset)
-		return 0;
+	if (*key == offset) {
+		bpf_kfunc_call_test_sleepable();
+		ok_sleepable |= (1 << offset);
+	}
 
-	bpf_kfunc_call_test_sleepable();
-	ok_sleepable |= (1 << offset);
+	wq_cb_run++;
 	return 0;
 }
 

---
base-commit: a30c772627b2ef4bc715332b1b04938a914556cd
change-id: 20261005-wq-poll-fix-f46c47dbaf03

Best regards,
-- 
Jonathan Bastari Thamrin <thamrinjonathan@gmail.com>