[PATCH bpf-next 6/6] selftests/bpf: add tests for the workqueue BPF iterators

[Imran Khan <imran.f.khan@oracle.com>]

Cover both the open-coded and seq_file forms of the workqueue and
worker_pool iterators, plus the pending-work seq iterator:

 - count_workqueues / count_worker_pools drive the open-coded iterators
   (inside bpf_rcu_read_lock()); userspace checks the counts are non-zero.
 - dump_workqueues / dump_worker_pools attach the seq forms and drain
   them, exercising the RCU-per-chunk seq path that reuses the open-coded
   next().
 - dump_pending attaches the workqueue_pending_work seq iterator and
   drains it, exercising the per-pool snapshot path.
 - wq_iter_fail.c checks the verifier rejects the RCU-protected
   open-coded iterators when used outside an RCU critical section.

Also declare the open-coded iterator kfuncs in bpf_experimental.h.

Signed-off-by: Imran Khan <imran.f.khan@oracle.com>
---
 .../testing/selftests/bpf/bpf_experimental.h  | 10 ++
 .../selftests/bpf/prog_tests/wq_iter.c        | 80 ++++++++++++++++
 tools/testing/selftests/bpf/progs/wq_iter.c   | 95 +++++++++++++++++++
 .../selftests/bpf/progs/wq_iter_fail.c        | 37 ++++++++
 4 files changed, 222 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/wq_iter.c
 create mode 100644 tools/testing/selftests/bpf/progs/wq_iter.c
 create mode 100644 tools/testing/selftests/bpf/progs/wq_iter_fail.c

diff --git a/tools/testing/selftests/bpf/bpf_experimental.h b/tools/testing/selftests/bpf/bpf_experimental.h
index 67ff7882299e..6da93c486650 100644
--- a/tools/testing/selftests/bpf/bpf_experimental.h
+++ b/tools/testing/selftests/bpf/bpf_experimental.h
@@ -361,6 +361,16 @@ extern int bpf_iter_dmabuf_new(struct bpf_iter_dmabuf *it) __weak __ksym;
 extern struct dma_buf *bpf_iter_dmabuf_next(struct bpf_iter_dmabuf *it) __weak __ksym;
 extern void bpf_iter_dmabuf_destroy(struct bpf_iter_dmabuf *it) __weak __ksym;
 
+struct bpf_iter_workqueue;
+extern int bpf_iter_workqueue_new(struct bpf_iter_workqueue *it) __weak __ksym;
+extern struct workqueue_struct *bpf_iter_workqueue_next(struct bpf_iter_workqueue *it) __weak __ksym;
+extern void bpf_iter_workqueue_destroy(struct bpf_iter_workqueue *it) __weak __ksym;
+
+struct bpf_iter_worker_pool;
+extern int bpf_iter_worker_pool_new(struct bpf_iter_worker_pool *it) __weak __ksym;
+extern struct worker_pool *bpf_iter_worker_pool_next(struct bpf_iter_worker_pool *it) __weak __ksym;
+extern void bpf_iter_worker_pool_destroy(struct bpf_iter_worker_pool *it) __weak __ksym;
+
 extern int bpf_cgroup_read_xattr(struct cgroup *cgroup, const char *name__str,
 				 struct bpf_dynptr *value_p) __weak __ksym;
 
diff --git a/tools/testing/selftests/bpf/prog_tests/wq_iter.c b/tools/testing/selftests/bpf/prog_tests/wq_iter.c
new file mode 100644
index 000000000000..73383c263bee
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/wq_iter.c
@@ -0,0 +1,80 @@
+// SPDX-License-Identifier: GPL-2.0
+#include <test_progs.h>
+#include "wq_iter.skel.h"
+#include "wq_iter_fail.skel.h"
+
+static void subtest_open_coded(struct wq_iter *skel)
+{
+	LIBBPF_OPTS(bpf_test_run_opts, opts);
+	int err;
+
+	err = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.count_workqueues),
+				     &opts);
+	if (!ASSERT_OK(err, "run count_workqueues"))
+		return;
+	/* There are always several system workqueues (events, ...). */
+	ASSERT_GT(skel->bss->nr_workqueues, 0, "nr_workqueues");
+
+	err = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.count_worker_pools),
+				     &opts);
+	if (!ASSERT_OK(err, "run count_worker_pools"))
+		return;
+	/* At least the per-CPU normal/highpri pools exist. */
+	ASSERT_GT(skel->bss->nr_worker_pools, 0, "nr_worker_pools");
+	ASSERT_GT(skel->bss->nr_percpu_pools, 0, "nr_percpu_pools");
+}
+
+static void drain_iter(struct bpf_program *prog, const char *name)
+{
+	struct bpf_link *link;
+	char buf[512];
+	int iter_fd;
+	ssize_t n;
+
+	link = bpf_program__attach_iter(prog, NULL);
+	if (!ASSERT_OK_PTR(link, name))
+		return;
+	iter_fd = bpf_iter_create(bpf_link__fd(link));
+	if (!ASSERT_GE(iter_fd, 0, "iter_create"))
+		goto out;
+	while ((n = read(iter_fd, buf, sizeof(buf))) > 0)
+		;
+	ASSERT_GE(n, 0, "read iter");
+	close(iter_fd);
+out:
+	bpf_link__destroy(link);
+}
+
+static void subtest_seq(struct wq_iter *skel)
+{
+	/* seq forms are backed by the same open-coded next(). */
+	drain_iter(skel->progs.dump_workqueues, "attach workqueue iter");
+	ASSERT_GT(skel->bss->nr_wq_seq, 0, "nr_wq_seq");
+
+	drain_iter(skel->progs.dump_worker_pools, "attach worker_pool iter");
+	ASSERT_GT(skel->bss->nr_pool_seq, 0, "nr_pool_seq");
+
+	/*
+	 * worklists are usually empty on an idle system; this drives the whole
+	 * per-pool snapshot path and verifies it drains cleanly.
+	 */
+	drain_iter(skel->progs.dump_pending, "attach pending_work iter");
+}
+
+void test_wq_iter(void)
+{
+	struct wq_iter *skel;
+
+	skel = wq_iter__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "wq_iter__open_and_load"))
+		return;
+
+	if (test__start_subtest("open_coded"))
+		subtest_open_coded(skel);
+	if (test__start_subtest("seq"))
+		subtest_seq(skel);
+
+	wq_iter__destroy(skel);
+
+	RUN_TESTS(wq_iter_fail);
+}
diff --git a/tools/testing/selftests/bpf/progs/wq_iter.c b/tools/testing/selftests/bpf/progs/wq_iter.c
new file mode 100644
index 000000000000..15f081bc634b
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/wq_iter.c
@@ -0,0 +1,95 @@
+// SPDX-License-Identifier: GPL-2.0
+#include "vmlinux.h"
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+#include "bpf_experimental.h"
+
+char _license[] SEC("license") = "GPL";
+
+void bpf_rcu_read_lock(void) __ksym;
+void bpf_rcu_read_unlock(void) __ksym;
+
+/* Results, checked by userspace. */
+int nr_workqueues;
+int nr_worker_pools;
+int nr_percpu_pools;
+int nr_wq_seq;
+int nr_pool_seq;
+int nr_pending;
+
+/* --- open-coded iterators (KF_RCU_PROTECTED) --- */
+
+SEC("syscall")
+int count_workqueues(const void *ctx)
+{
+	struct workqueue_struct *wq;
+	int n = 0;
+
+	bpf_rcu_read_lock();
+	bpf_for_each(workqueue, wq)
+		n++;
+	bpf_rcu_read_unlock();
+
+	nr_workqueues = n;
+	return 0;
+}
+
+SEC("syscall")
+int count_worker_pools(const void *ctx)
+{
+	struct worker_pool *pool;
+	int n = 0, percpu = 0;
+
+	bpf_rcu_read_lock();
+	bpf_for_each(worker_pool, pool) {
+		n++;
+		if (pool->cpu >= 0)		/* per-CPU pool */
+			percpu++;
+	}
+	bpf_rcu_read_unlock();
+
+	nr_worker_pools = n;
+	nr_percpu_pools = percpu;
+	return 0;
+}
+
+/* --- seq_file iterators --- */
+
+SEC("iter/workqueue")
+int dump_workqueues(struct bpf_iter__workqueue *ctx)
+{
+	struct seq_file *seq = ctx->meta->seq;
+	struct workqueue_struct *wq = ctx->wq;
+
+	if (!wq)
+		return 0;
+	nr_wq_seq++;
+	BPF_SEQ_PRINTF(seq, "%s\n", wq->name);
+	return 0;
+}
+
+SEC("iter/worker_pool")
+int dump_worker_pools(struct bpf_iter__worker_pool *ctx)
+{
+	struct worker_pool *pool = ctx->pool;
+
+	if (!pool)
+		return 0;
+	nr_pool_seq++;
+	return 0;
+}
+
+SEC("iter/workqueue_pending_work")
+int dump_pending(struct bpf_iter__workqueue_pending_work *ctx)
+{
+	struct seq_file *seq = ctx->meta->seq;
+	struct wq_pending_work_info *info = ctx->info;
+
+	if (!info)			/* final call */
+		return 0;
+
+	nr_pending++;
+	BPF_SEQ_PRINTF(seq, "pool %llu work 0x%llx func 0x%llx\n",
+		       info->pool_id, info->work, info->func);
+	return 0;
+}
diff --git a/tools/testing/selftests/bpf/progs/wq_iter_fail.c b/tools/testing/selftests/bpf/progs/wq_iter_fail.c
new file mode 100644
index 000000000000..fb2dcf9847e2
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/wq_iter_fail.c
@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: GPL-2.0
+#include "vmlinux.h"
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+#include "bpf_misc.h"
+#include "bpf_experimental.h"
+
+char _license[] SEC("license") = "GPL";
+
+/*
+ * The workqueue and worker_pool open-coded iterators are KF_RCU_PROTECTED:
+ * using them outside a bpf_rcu_read_lock() region must be rejected.
+ */
+
+SEC("?syscall")
+__failure __msg("kernel func bpf_iter_workqueue_new requires RCU critical section protection")
+int workqueue_iter_no_rcu(const void *ctx)
+{
+	struct workqueue_struct *wq;
+	int n = 0;
+
+	bpf_for_each(workqueue, wq)
+		n++;
+	return n;
+}
+
+SEC("?syscall")
+__failure __msg("kernel func bpf_iter_worker_pool_new requires RCU critical section protection")
+int worker_pool_iter_no_rcu(const void *ctx)
+{
+	struct worker_pool *pool;
+	int n = 0;
+
+	bpf_for_each(worker_pool, pool)
+		n++;
+	return n;
+}
-- 
2.43.0