[PATCH bpf-next 0/2] bpf: Add a generic bits iterator

[PATCH bpf-next 0/2] bpf: Add a generic bits iterator

[Yafang Shao]

Introducing three new kfuncs, namely bpf_iter_bits_{new,next,destroy}, to
support the newly added bits iter functionality. These functions enable
seamless iteration of bits from a specified memory area.

- bpf_iter_bits_new
  Initializes a new bits iterator for a given memory area. Notably, due to
  limitations within bpf memalloc, the maximum number of bits that can be
  iterated over is constrained to (4096 * 8).
- bpf_iter_bits_next
  Get the next bit in a bpf_iter_bits
- bpf_iter_bits_destroy
  Destroy a bpf_iter_bits

The bits iterator can be used in any context and on any address.

In our specific use case, we leverage the cgroup iterator to traverse
percpu data, subsequently exposing it to userspace through a seq file.
For detailed usage examples, please refer to the example in patch #2.

Changes:
- bpf: Add bpf_iter_cpumask
  https://lwn.net/Articles/961104/
- bpf: Add new bpf helper bpf_for_each_cpu
  https://lwn.net/Articles/939939/

Yafang Shao (2):
  bpf: Add bits iterator
  selftests/bpf: Add selftest for bits iter

 kernel/bpf/helpers.c                          | 100 ++++++++++
 tools/testing/selftests/bpf/config            |   1 +
 .../selftests/bpf/prog_tests/bits_iter.c      | 180 ++++++++++++++++++
 .../bpf/progs/test_bits_iter_failure.c        |  53 ++++++
 .../bpf/progs/test_bits_iter_success.c        | 146 ++++++++++++++
 5 files changed, 480 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/bits_iter.c
 create mode 100644 tools/testing/selftests/bpf/progs/test_bits_iter_failure.c
 create mode 100644 tools/testing/selftests/bpf/progs/test_bits_iter_success.c

-- 
2.39.1

[PATCH bpf-next 1/2] bpf: Add bits iterator

[Yafang Shao]

Add three new kfuncs for the bits iterator:
- bpf_iter_bits_new
  Initialize a new bits iterator for a given memory area. Due to the
  limitation of bpf memalloc, the max number of bits that can be iterated
  over is limited to (4096 * 8).
- bpf_iter_bits_next
  Get the next bit in a bpf_iter_bits
- bpf_iter_bits_destroy
  Destroy a bpf_iter_bits

The bits iterator facilitates the iteration of the bits of a memory area,
such as cpumask. It can be used in any context and on any address.

Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
---
 kernel/bpf/helpers.c | 100 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 100 insertions(+)

diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
index 93edf730d288..052f63891834 100644
--- a/kernel/bpf/helpers.c
+++ b/kernel/bpf/helpers.c
@@ -2542,6 +2542,103 @@ __bpf_kfunc void bpf_throw(u64 cookie)
 	WARN(1, "A call to BPF exception callback should never return\n");
 }
 
+struct bpf_iter_bits {
+	__u64 __opaque[2];
+} __aligned(8);
+
+struct bpf_iter_bits_kern {
+	unsigned long *bits;
+	u32 nr_bits;
+	int bit;
+} __aligned(8);
+
+/**
+ * bpf_iter_bits_new() - Initialize a new bits iterator for a given memory area
+ * @it: The new bpf_iter_bits to be created
+ * @unsafe_ptr__ign: A ponter pointing to a memory area to be iterated over
+ * @nr_bits: The number of bits to be iterated over. Due to the limitation of
+ * memalloc, it can't greater than (4096 * 8).
+ *
+ * This function initializes a new bpf_iter_bits structure for iterating over
+ * a memory area which is specified by the @unsafe_ptr__ign and @nr_bits. It
+ * copy the data of the memory area to the newly created bpf_iter_bits @it for
+ * subsequent iteration operations.
+ *
+ * On success, 0 is returned. On failure, ERR is returned.
+ */
+__bpf_kfunc int
+bpf_iter_bits_new(struct bpf_iter_bits *it, const void *unsafe_ptr__ign, u32 nr_bits)
+{
+	struct bpf_iter_bits_kern *kit = (void *)it;
+	u32 size = BITS_TO_BYTES(nr_bits);
+	int err;
+
+	BUILD_BUG_ON(sizeof(struct bpf_iter_bits_kern) != sizeof(struct bpf_iter_bits));
+	BUILD_BUG_ON(__alignof__(struct bpf_iter_bits_kern) !=
+		     __alignof__(struct bpf_iter_bits));
+
+	if (!unsafe_ptr__ign || !nr_bits) {
+		kit->bits = NULL;
+		return -EINVAL;
+	}
+
+	kit->bits = bpf_mem_alloc(&bpf_global_ma, size);
+	if (!kit->bits)
+		return -ENOMEM;
+
+	err = bpf_probe_read_kernel_common(kit->bits, size, unsafe_ptr__ign);
+	if (err) {
+		bpf_mem_free(&bpf_global_ma, kit->bits);
+		kit->bits = NULL;
+		return err;
+	}
+
+	kit->nr_bits = nr_bits;
+	kit->bit = -1;
+	return 0;
+}
+
+/**
+ * bpf_iter_bits_next() - Get the next bit in a bpf_iter_bits
+ * @it: The bpf_iter_bits to be checked
+ *
+ * This function returns a pointer to a number representing the value of the
+ * next bit in the bits.
+ *
+ * If there are no further bit available, it returns NULL.
+ */
+__bpf_kfunc int *bpf_iter_bits_next(struct bpf_iter_bits *it)
+{
+	struct bpf_iter_bits_kern *kit = (void *)it;
+	const unsigned long *bits = kit->bits;
+	int bit;
+
+	if (!bits)
+		return NULL;
+
+	bit = find_next_bit(bits, kit->nr_bits, kit->bit + 1);
+	if (bit >= kit->nr_bits)
+		return NULL;
+
+	kit->bit = bit;
+	return &kit->bit;
+}
+
+/**
+ * bpf_iter_bits_destroy() - Destroy a bpf_iter_bits
+ * @it: The bpf_iter_bits to be destroyed
+ *
+ * Destroy the resource associated with the bpf_iter_bits.
+ */
+__bpf_kfunc void bpf_iter_bits_destroy(struct bpf_iter_bits *it)
+{
+	struct bpf_iter_bits_kern *kit = (void *)it;
+
+	if (!kit->bits)
+		return;
+	bpf_mem_free(&bpf_global_ma, kit->bits);
+}
+
 __bpf_kfunc_end_defs();
 
 BTF_KFUNCS_START(generic_btf_ids)
@@ -2618,6 +2715,9 @@ BTF_ID_FLAGS(func, bpf_dynptr_is_null)
 BTF_ID_FLAGS(func, bpf_dynptr_is_rdonly)
 BTF_ID_FLAGS(func, bpf_dynptr_size)
 BTF_ID_FLAGS(func, bpf_dynptr_clone)
+BTF_ID_FLAGS(func, bpf_iter_bits_new, KF_ITER_NEW)
+BTF_ID_FLAGS(func, bpf_iter_bits_next, KF_ITER_NEXT | KF_RET_NULL)
+BTF_ID_FLAGS(func, bpf_iter_bits_destroy, KF_ITER_DESTROY)
 BTF_KFUNCS_END(common_btf_ids)
 
 static const struct btf_kfunc_id_set common_kfunc_set = {
-- 
2.39.1

RE: [PATCH bpf-next 1/2] bpf: Add bits iterator

[John Fastabend]

Yafang Shao wrote:
> Add three new kfuncs for the bits iterator:
> - bpf_iter_bits_new
>   Initialize a new bits iterator for a given memory area. Due to the
>   limitation of bpf memalloc, the max number of bits that can be iterated
>   over is limited to (4096 * 8).
> - bpf_iter_bits_next
>   Get the next bit in a bpf_iter_bits
> - bpf_iter_bits_destroy
>   Destroy a bpf_iter_bits
> 
> The bits iterator facilitates the iteration of the bits of a memory area,
> such as cpumask. It can be used in any context and on any address.

Just curious as I see more and a more kfuncs. Did you try to implement
this with existing BPF? The main trick looks to be to get an implementation
of FIND_NEXT_BIT? Without trying seems doable with one of the bpf loop
iterators?

Also this requires a bpf_iter_bits_new across every iteration of the
BPF program or anytime we need to pick up the changes. Any reason
we can't just read the memory directly?

Thanks,
John

> 
> Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
> ---
>  kernel/bpf/helpers.c | 100 +++++++++++++++++++++++++++++++++++++++++++
>  1 file changed, 100 insertions(+)
> 
> diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
> index 93edf730d288..052f63891834 100644
> --- a/kernel/bpf/helpers.c
> +++ b/kernel/bpf/helpers.c
> @@ -2542,6 +2542,103 @@ __bpf_kfunc void bpf_throw(u64 cookie)
>  	WARN(1, "A call to BPF exception callback should never return\n");
>  }
>  
> +struct bpf_iter_bits {
> +	__u64 __opaque[2];
> +} __aligned(8);
> +
> +struct bpf_iter_bits_kern {
> +	unsigned long *bits;
> +	u32 nr_bits;
> +	int bit;
> +} __aligned(8);
> +
> +/**
> + * bpf_iter_bits_new() - Initialize a new bits iterator for a given memory area
> + * @it: The new bpf_iter_bits to be created
> + * @unsafe_ptr__ign: A ponter pointing to a memory area to be iterated over
> + * @nr_bits: The number of bits to be iterated over. Due to the limitation of
> + * memalloc, it can't greater than (4096 * 8).
> + *
> + * This function initializes a new bpf_iter_bits structure for iterating over
> + * a memory area which is specified by the @unsafe_ptr__ign and @nr_bits. It
> + * copy the data of the memory area to the newly created bpf_iter_bits @it for
> + * subsequent iteration operations.
> + *
> + * On success, 0 is returned. On failure, ERR is returned.
> + */
> +__bpf_kfunc int
> +bpf_iter_bits_new(struct bpf_iter_bits *it, const void *unsafe_ptr__ign, u32 nr_bits)
> +{
> +	struct bpf_iter_bits_kern *kit = (void *)it;
> +	u32 size = BITS_TO_BYTES(nr_bits);
> +	int err;
> +
> +	BUILD_BUG_ON(sizeof(struct bpf_iter_bits_kern) != sizeof(struct bpf_iter_bits));
> +	BUILD_BUG_ON(__alignof__(struct bpf_iter_bits_kern) !=
> +		     __alignof__(struct bpf_iter_bits));
> +
> +	if (!unsafe_ptr__ign || !nr_bits) {
> +		kit->bits = NULL;
> +		return -EINVAL;
> +	}
> +
> +	kit->bits = bpf_mem_alloc(&bpf_global_ma, size);
> +	if (!kit->bits)
> +		return -ENOMEM;
> +
> +	err = bpf_probe_read_kernel_common(kit->bits, size, unsafe_ptr__ign);

Specifically, this why can't we iterate over unsafe_ptr__ign?

> +	if (err) {
> +		bpf_mem_free(&bpf_global_ma, kit->bits);
> +		kit->bits = NULL;
> +		return err;
> +	}
> +
> +	kit->nr_bits = nr_bits;
> +	kit->bit = -1;
> +	return 0;
> +}
> +
> +/**
> + * bpf_iter_bits_next() - Get the next bit in a bpf_iter_bits
> + * @it: The bpf_iter_bits to be checked
> + *
> + * This function returns a pointer to a number representing the value of the
> + * next bit in the bits.
> + *
> + * If there are no further bit available, it returns NULL.
> + */
> +__bpf_kfunc int *bpf_iter_bits_next(struct bpf_iter_bits *it)
> +{
> +	struct bpf_iter_bits_kern *kit = (void *)it;
> +	const unsigned long *bits = kit->bits;
> +	int bit;
> +
> +	if (!bits)
> +		return NULL;
> +
> +	bit = find_next_bit(bits, kit->nr_bits, kit->bit + 1);

Seems like this should be ok over unsafe memory as long as find_next_bit
is bounded?

> +	if (bit >= kit->nr_bits)
> +		return NULL;
> +
> +	kit->bit = bit;
> +	return &kit->bit;
> +}


Thanks for working on this looks useful to me.

Re: [PATCH bpf-next 1/2] bpf: Add bits iterator

[Alexei Starovoitov]

On Mon, Feb 26, 2024 at 3:21 PM John Fastabend <john.fastabend@gmail.com> wrote:
> > +__bpf_kfunc int *bpf_iter_bits_next(struct bpf_iter_bits *it)
> > +{
> > +     struct bpf_iter_bits_kern *kit = (void *)it;
> > +     const unsigned long *bits = kit->bits;
> > +     int bit;
> > +
> > +     if (!bits)
> > +             return NULL;
> > +
> > +     bit = find_next_bit(bits, kit->nr_bits, kit->bit + 1);
>
> Seems like this should be ok over unsafe memory as long as find_next_bit
> is bounded?

Are you proposing to add find_next_bit() as a kfunc instead?

With the bpf_can_loop() proposal these two can be combined and
it will probably achieve the same result.
But imo this iterator is small enough to get in now and
delete later when there is a better way.
Ideally we'd need to add new instructions to operate with bits efficiently.

Re: [PATCH bpf-next 1/2] bpf: Add bits iterator

[John Fastabend]

Alexei Starovoitov wrote:
> On Mon, Feb 26, 2024 at 3:21 PM John Fastabend <john.fastabend@gmail.com> wrote:
> > > +__bpf_kfunc int *bpf_iter_bits_next(struct bpf_iter_bits *it)
> > > +{
> > > +     struct bpf_iter_bits_kern *kit = (void *)it;
> > > +     const unsigned long *bits = kit->bits;
> > > +     int bit;
> > > +
> > > +     if (!bits)
> > > +             return NULL;
> > > +
> > > +     bit = find_next_bit(bits, kit->nr_bits, kit->bit + 1);
> >
> > Seems like this should be ok over unsafe memory as long as find_next_bit
> > is bounded?
> 
> Are you proposing to add find_next_bit() as a kfunc instead?

I was suggesting you could likely implement find_next_bit() in
BPF directly no need for a kfunc.

> 
> With the bpf_can_loop() proposal these two can be combined and
> it will probably achieve the same result.
> But imo this iterator is small enough to get in now and
> delete later when there is a better way.

Agree its fine to go in as is IMO. Mostly just curious if anyone
tried to implement it in BPF.

> Ideally we'd need to add new instructions to operate with bits efficiently.

+1.

[PATCH bpf-next 2/2] selftests/bpf: Add selftest for bits iter

[Yafang Shao]

Add selftests for the newly added bits iter.
- bits_iter_success
  - The number of CPUs should be expected when iterating over a cpumask
  - percpu data extracted from the percpu struct should be expected
  - RCU lock is not required
  - It is fine without calling bpf_iter_cpumask_next()
  - It can work as expected when invalid arguments are passed

- bits_iter_failure
  - bpf_iter_bits_destroy() is required after calling
    bpf_iter_bits_new()
  - bpf_iter_bits_destroy() can only destroy an initialized iter
  - bpf_iter_bits_next() must use an initialized iter

Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
---
 tools/testing/selftests/bpf/config            |   1 +
 .../selftests/bpf/prog_tests/bits_iter.c      | 180 ++++++++++++++++++
 .../bpf/progs/test_bits_iter_failure.c        |  53 ++++++
 .../bpf/progs/test_bits_iter_success.c        | 146 ++++++++++++++
 4 files changed, 380 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/bits_iter.c
 create mode 100644 tools/testing/selftests/bpf/progs/test_bits_iter_failure.c
 create mode 100644 tools/testing/selftests/bpf/progs/test_bits_iter_success.c

diff --git a/tools/testing/selftests/bpf/config b/tools/testing/selftests/bpf/config
index 01f241ea2c67..dd4b0935e35f 100644
--- a/tools/testing/selftests/bpf/config
+++ b/tools/testing/selftests/bpf/config
@@ -78,6 +78,7 @@ CONFIG_NF_CONNTRACK_MARK=y
 CONFIG_NF_DEFRAG_IPV4=y
 CONFIG_NF_DEFRAG_IPV6=y
 CONFIG_NF_NAT=y
+CONFIG_PSI=y
 CONFIG_RC_CORE=y
 CONFIG_SECURITY=y
 CONFIG_SECURITYFS=y
diff --git a/tools/testing/selftests/bpf/prog_tests/bits_iter.c b/tools/testing/selftests/bpf/prog_tests/bits_iter.c
new file mode 100644
index 000000000000..778a7c942dba
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/bits_iter.c
@@ -0,0 +1,180 @@
+// SPDX-License-Identifier: GPL-2.0
+/* Copyright (c) 2024 Yafang Shao <laoar.shao@gmail.com> */
+
+#define _GNU_SOURCE
+#include <sched.h>
+
+#include <test_progs.h>
+#include "test_bits_iter_success.skel.h"
+#include "test_bits_iter_failure.skel.h"
+#include "cgroup_helpers.h"
+
+static const char * const positive_testcases[] = {
+	"cpumask_iter",
+};
+
+static const char * const negative_testcases[] = {
+	"null_pointer",
+	"zero_bit",
+	"no_mem",
+	"invalid_bits"
+};
+
+static int read_percpu_data(struct bpf_link *link, int nr_cpu_exp, int nr_running_exp)
+{
+	int iter_fd, len, item, nr_running, psi_running, nr_cpus, err = -1;
+	char buf[128];
+	size_t left;
+	char *p;
+
+	iter_fd = bpf_iter_create(bpf_link__fd(link));
+	if (!ASSERT_GE(iter_fd, 0, "iter_fd"))
+		return -1;
+
+	memset(buf, 0, sizeof(buf));
+	left = ARRAY_SIZE(buf);
+	p = buf;
+	while ((len = read(iter_fd, p, left)) > 0) {
+		p += len;
+		left -= len;
+	}
+
+	item = sscanf(buf, "nr_running %u nr_cpus %u psi_running %u\n",
+		      &nr_running, &nr_cpus, &psi_running);
+	if (!ASSERT_EQ(item, 3, "seq_format"))
+		goto out;
+	if (!ASSERT_EQ(nr_cpus, nr_cpu_exp, "nr_cpus"))
+		goto out;
+	if (!ASSERT_GE(nr_running, nr_running_exp, "nr_running"))
+		goto out;
+	if (!ASSERT_GE(psi_running, nr_running_exp, "psi_running"))
+		goto out;
+
+	err = 0;
+out:
+	close(iter_fd);
+	return err;
+}
+
+static void verify_iter_success(const char *prog_name, bool negtive)
+{
+	DECLARE_LIBBPF_OPTS(bpf_iter_attach_opts, opts);
+	int cgrp_fd, nr_cpus, err, i, chosen = 0;
+	struct test_bits_iter_success *skel;
+	union bpf_iter_link_info linfo;
+	struct bpf_program *prog;
+	struct bpf_link *link;
+	cpu_set_t set;
+
+	if (setup_cgroup_environment())
+		return;
+
+	/* Utilize the cgroup iter */
+	cgrp_fd = get_root_cgroup();
+	if (!ASSERT_GE(cgrp_fd, 0, "create_cgrp"))
+		goto cleanup;
+
+	skel = test_bits_iter_success__open();
+	if (!ASSERT_OK_PTR(skel, "cpumask_iter_success__open"))
+		goto close_fd;
+
+	skel->bss->pid = getpid();
+	nr_cpus = libbpf_num_possible_cpus();
+	skel->bss->total_nr_cpus = nr_cpus;
+
+	err = test_bits_iter_success__load(skel);
+	if (!ASSERT_OK(err, "cpumask_iter_success__load"))
+		goto destroy;
+
+	prog = bpf_object__find_program_by_name(skel->obj, prog_name);
+	if (!ASSERT_OK_PTR(prog, "bpf_object__find_program_by_name"))
+		goto destroy;
+
+	memset(&linfo, 0, sizeof(linfo));
+	linfo.cgroup.cgroup_fd = cgrp_fd;
+	linfo.cgroup.order = BPF_CGROUP_ITER_SELF_ONLY;
+	opts.link_info = &linfo;
+	opts.link_info_len = sizeof(linfo);
+	link = bpf_program__attach_iter(prog, &opts);
+	if (!ASSERT_OK_PTR(link, "bpf_program__attach"))
+		goto destroy;
+
+	if (negtive)
+		goto negtive;
+
+	/* Case 1): Enable all possible CPUs */
+	CPU_ZERO(&set);
+	for (i = 0; i < nr_cpus; i++)
+		CPU_SET(i, &set);
+	err = sched_setaffinity(skel->bss->pid, sizeof(set), &set);
+	if (!ASSERT_OK(err, "setaffinity_all_cpus"))
+		goto free_link;
+	err = read_percpu_data(link, nr_cpus, 1);
+	if (!ASSERT_OK(err, "read_percpu_data"))
+		goto free_link;
+	if (!ASSERT_OK(skel->bss->err, "null_rq"))
+		goto free_link;
+
+	/* Case 2): CPU0 only */
+	CPU_ZERO(&set);
+	CPU_SET(0, &set);
+	err = sched_setaffinity(skel->bss->pid, sizeof(set), &set);
+	if (!ASSERT_OK(err, "setaffinity_cpu0"))
+		goto free_link;
+	err = read_percpu_data(link, 1, 1);
+	if (!ASSERT_OK(err, "read_percpu_data"))
+		goto free_link;
+	if (!ASSERT_OK(skel->bss->err, "null_rq_psi"))
+		goto free_link;
+
+	/* Case 3): Partial CPUs */
+	CPU_ZERO(&set);
+	for (i = 0; i < nr_cpus; i++) {
+		if (i < 4 && i & 0x1)
+			continue;
+		if (i > 8 && i & 0x2)
+			continue;
+		CPU_SET(i, &set);
+		chosen++;
+	}
+	err = sched_setaffinity(skel->bss->pid, sizeof(set), &set);
+	if (!ASSERT_OK(err, "setaffinity_partial_cpus"))
+		goto free_link;
+	err = read_percpu_data(link, chosen, 1);
+	if (!ASSERT_OK(err, "read_percpu_data"))
+		goto free_link;
+
+negtive:
+	ASSERT_OK(skel->bss->err, "null_rq_psi");
+
+free_link:
+	bpf_link__destroy(link);
+destroy:
+	test_bits_iter_success__destroy(skel);
+close_fd:
+	close(cgrp_fd);
+cleanup:
+	cleanup_cgroup_environment();
+}
+
+void test_bits_iter(void)
+{
+	int i;
+
+	for (i = 0; i < ARRAY_SIZE(positive_testcases); i++) {
+		if (!test__start_subtest(positive_testcases[i]))
+			continue;
+
+		verify_iter_success(positive_testcases[i], false);
+	}
+
+	for (i = 0; i < ARRAY_SIZE(negative_testcases); i++) {
+		if (!test__start_subtest(negative_testcases[i]))
+			continue;
+
+		verify_iter_success(negative_testcases[i], true);
+	}
+
+	RUN_TESTS(test_bits_iter_success);
+	RUN_TESTS(test_bits_iter_failure);
+}
diff --git a/tools/testing/selftests/bpf/progs/test_bits_iter_failure.c b/tools/testing/selftests/bpf/progs/test_bits_iter_failure.c
new file mode 100644
index 000000000000..c51f18f4f334
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/test_bits_iter_failure.c
@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: GPL-2.0-only
+/* Copyright (c) 2024 Yafang Shao <laoar.shao@gmail.com> */
+
+#include "vmlinux.h"
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+
+#include "bpf_misc.h"
+#include "task_kfunc_common.h"
+
+char _license[] SEC("license") = "GPL";
+
+int bpf_iter_bits_new(struct bpf_iter_bits *it, const void *unsafe_ptr, u32 nr_bits) __ksym __weak;
+int *bpf_iter_bits_next(struct bpf_iter_bits *it) __ksym __weak;
+void bpf_iter_bits_destroy(struct bpf_iter_bits *it) __ksym __weak;
+
+SEC("iter.s/cgroup")
+__failure __msg("Unreleased reference id=3 alloc_insn=10")
+int BPF_PROG(no_destroy, struct bpf_iter_meta *meta, struct cgroup *cgrp)
+{
+	struct bpf_iter_bits it;
+	struct task_struct *p;
+
+	p = bpf_task_from_pid(1);
+	if (!p)
+		return 1;
+
+	bpf_iter_bits_new(&it, p->cpus_ptr, 8192);
+
+	bpf_iter_bits_next(&it);
+	bpf_task_release(p);
+	return 0;
+}
+
+SEC("iter/cgroup")
+__failure __msg("expected an initialized iter_bits as arg #1")
+int BPF_PROG(next_uninit, struct bpf_iter_meta *meta, struct cgroup *cgrp)
+{
+	struct bpf_iter_bits *it = NULL;
+
+	bpf_iter_bits_next(it);
+	return 0;
+}
+
+SEC("iter/cgroup")
+__failure __msg("expected an initialized iter_bits as arg #1")
+int BPF_PROG(destroy_uninit, struct bpf_iter_meta *meta, struct cgroup *cgrp)
+{
+	struct bpf_iter_bits it = {};
+
+	bpf_iter_bits_destroy(&it);
+	return 0;
+}
diff --git a/tools/testing/selftests/bpf/progs/test_bits_iter_success.c b/tools/testing/selftests/bpf/progs/test_bits_iter_success.c
new file mode 100644
index 000000000000..6e5f12ad17ce
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/test_bits_iter_success.c
@@ -0,0 +1,146 @@
+// SPDX-License-Identifier: GPL-2.0-only
+/* Copyright (c) 2024 Yafang Shao <laoar.shao@gmail.com> */
+
+#include "vmlinux.h"
+#include <linux/const.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+
+#include "task_kfunc_common.h"
+
+char _license[] SEC("license") = "GPL";
+
+extern const struct psi_group_cpu system_group_pcpu __ksym __weak;
+extern const struct rq runqueues __ksym __weak;
+
+int bpf_iter_bits_new(struct bpf_iter_bits *it, const void *unsafe_ptr, u32 nr_bits) __ksym __weak;
+int *bpf_iter_bits_next(struct bpf_iter_bits *it) __ksym __weak;
+void bpf_iter_bits_destroy(struct bpf_iter_bits *it) __ksym __weak;
+
+int pid, err, total_nr_cpus;
+
+SEC("iter.s/cgroup")
+int BPF_PROG(cpumask_iter, struct bpf_iter_meta *meta, struct cgroup *cgrp)
+{
+	u32 nr_running = 0, psi_nr_running = 0, nr_cpus = 0;
+	struct psi_group_cpu *groupc;
+	struct task_struct *p;
+	struct rq *rq;
+	int *cpu;
+
+	/* epilogue */
+	if (!cgrp)
+		return 0;
+
+	p = bpf_task_from_pid(pid);
+	if (!p)
+		return 1;
+
+	bpf_for_each(bits, cpu, p->cpus_ptr, total_nr_cpus) {
+		rq = (struct rq *)bpf_per_cpu_ptr(&runqueues, *cpu);
+		/* Each valid CPU must have a runqueue, even if it is offline. */
+		if (!rq) {
+			err++;
+			continue;
+		}
+
+		nr_running += rq->nr_running;
+		nr_cpus++;
+
+		groupc = (struct psi_group_cpu *)bpf_per_cpu_ptr(&system_group_pcpu, *cpu);
+		if (!groupc) {
+			err++;
+			continue;
+		}
+		psi_nr_running += groupc->tasks[NR_RUNNING];
+	}
+	BPF_SEQ_PRINTF(meta->seq, "nr_running %u nr_cpus %u psi_running %u\n",
+		       nr_running, nr_cpus, psi_nr_running);
+	bpf_task_release(p);
+	return 0;
+}
+
+SEC("iter.s/cgroup")
+int BPF_PROG(null_pointer, struct bpf_iter_meta *meta, struct cgroup *cgrp)
+{
+	int *cpu;
+
+	bpf_for_each(bits, cpu, NULL, total_nr_cpus)
+		err++;
+	return 0;
+}
+
+SEC("iter.s/cgroup")
+int BPF_PROG(zero_bit, struct bpf_iter_meta *meta, struct cgroup *cgrp)
+{
+	struct task_struct *p;
+	int *cpu;
+
+	p = bpf_task_from_pid(pid);
+	if (!p)
+		return 1;
+
+	bpf_for_each(bits, cpu, p->cpus_ptr, 0)
+		err++;
+	bpf_task_release(p);
+	return 0;
+}
+
+SEC("iter.s/cgroup")
+int BPF_PROG(no_mem, struct bpf_iter_meta *meta, struct cgroup *cgrp)
+{
+	struct task_struct *p;
+	int *cpu;
+
+	p = bpf_task_from_pid(pid);
+	if (!p)
+		return 1;
+
+	/* The max size of memalloc is 4096, so it will fail to allocate (8192 * 8) */
+	bpf_for_each(bits, cpu, p->cpus_ptr, 8192 * 8)
+		err++;
+	bpf_task_release(p);
+	return 0;
+}
+
+SEC("iter/cgroup")
+int BPF_PROG(invalid_bits, struct bpf_iter_meta *meta, struct cgroup *cgrp)
+{
+	struct task_struct *p;
+	struct rq *rq;
+	int *cpu;
+
+	p = bpf_task_from_pid(pid);
+	if (!p)
+		return 1;
+
+	bpf_for_each(bits, cpu, p->cpus_ptr, 8192) {
+		rq = (struct rq *)bpf_per_cpu_ptr(&runqueues, *cpu);
+		/* For invalid CPU IDs, the rq must be NULL. */
+		if (!rq)
+			err++;
+	}
+	if (err)
+		err -= 8192 - total_nr_cpus;
+	bpf_task_release(p);
+	return 0;
+}
+
+SEC("iter.s/cgroup")
+int BPF_PROG(no_next, struct bpf_iter_meta *meta, struct cgroup *cgrp)
+{
+	struct bpf_iter_bits it;
+	struct task_struct *p;
+
+	p = bpf_task_from_pid(1);
+	if (!p)
+		return 1;
+
+	bpf_iter_bits_new(&it, p->cpus_ptr, 8192);
+
+	/* It is fine without calling bpf_iter_bits_next(). */
+
+	bpf_iter_bits_destroy(&it);
+	bpf_task_release(p);
+	return 0;
+}
-- 
2.39.1

Re: [PATCH bpf-next 2/2] selftests/bpf: Add selftest for bits iter

[Alexei Starovoitov]

On Sun, Feb 18, 2024 at 3:49 AM Yafang Shao <laoar.shao@gmail.com> wrote:
>
> Add selftests for the newly added bits iter.
> - bits_iter_success
>   - The number of CPUs should be expected when iterating over a cpumask
>   - percpu data extracted from the percpu struct should be expected
>   - RCU lock is not required
>   - It is fine without calling bpf_iter_cpumask_next()
>   - It can work as expected when invalid arguments are passed
>
> - bits_iter_failure
>   - bpf_iter_bits_destroy() is required after calling
>     bpf_iter_bits_new()
>   - bpf_iter_bits_destroy() can only destroy an initialized iter
>   - bpf_iter_bits_next() must use an initialized iter
>
> Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
> ---
>  tools/testing/selftests/bpf/config            |   1 +
>  .../selftests/bpf/prog_tests/bits_iter.c      | 180 ++++++++++++++++++
>  .../bpf/progs/test_bits_iter_failure.c        |  53 ++++++
>  .../bpf/progs/test_bits_iter_success.c        | 146 ++++++++++++++
>  4 files changed, 380 insertions(+)
>  create mode 100644 tools/testing/selftests/bpf/prog_tests/bits_iter.c
>  create mode 100644 tools/testing/selftests/bpf/progs/test_bits_iter_failure.c
>  create mode 100644 tools/testing/selftests/bpf/progs/test_bits_iter_success.c
>
> diff --git a/tools/testing/selftests/bpf/config b/tools/testing/selftests/bpf/config
> index 01f241ea2c67..dd4b0935e35f 100644
> --- a/tools/testing/selftests/bpf/config
> +++ b/tools/testing/selftests/bpf/config
> @@ -78,6 +78,7 @@ CONFIG_NF_CONNTRACK_MARK=y
>  CONFIG_NF_DEFRAG_IPV4=y
>  CONFIG_NF_DEFRAG_IPV6=y
>  CONFIG_NF_NAT=y
> +CONFIG_PSI=y
>  CONFIG_RC_CORE=y
>  CONFIG_SECURITY=y
>  CONFIG_SECURITYFS=y
> diff --git a/tools/testing/selftests/bpf/prog_tests/bits_iter.c b/tools/testing/selftests/bpf/prog_tests/bits_iter.c
> new file mode 100644
> index 000000000000..778a7c942dba
> --- /dev/null
> +++ b/tools/testing/selftests/bpf/prog_tests/bits_iter.c
> @@ -0,0 +1,180 @@
> +// SPDX-License-Identifier: GPL-2.0
> +/* Copyright (c) 2024 Yafang Shao <laoar.shao@gmail.com> */
> +
> +#define _GNU_SOURCE
> +#include <sched.h>
> +
> +#include <test_progs.h>
> +#include "test_bits_iter_success.skel.h"
> +#include "test_bits_iter_failure.skel.h"
> +#include "cgroup_helpers.h"
> +
> +static const char * const positive_testcases[] = {
> +       "cpumask_iter",
> +};
> +
> +static const char * const negative_testcases[] = {
> +       "null_pointer",
> +       "zero_bit",
> +       "no_mem",
> +       "invalid_bits"
> +};
> +
> +static int read_percpu_data(struct bpf_link *link, int nr_cpu_exp, int nr_running_exp)
> +{
> +       int iter_fd, len, item, nr_running, psi_running, nr_cpus, err = -1;
> +       char buf[128];
> +       size_t left;
> +       char *p;
> +
> +       iter_fd = bpf_iter_create(bpf_link__fd(link));
> +       if (!ASSERT_GE(iter_fd, 0, "iter_fd"))
> +               return -1;
> +
> +       memset(buf, 0, sizeof(buf));
> +       left = ARRAY_SIZE(buf);
> +       p = buf;
> +       while ((len = read(iter_fd, p, left)) > 0) {
> +               p += len;
> +               left -= len;
> +       }
> +
> +       item = sscanf(buf, "nr_running %u nr_cpus %u psi_running %u\n",
> +                     &nr_running, &nr_cpus, &psi_running);
> +       if (!ASSERT_EQ(item, 3, "seq_format"))
> +               goto out;
> +       if (!ASSERT_EQ(nr_cpus, nr_cpu_exp, "nr_cpus"))
> +               goto out;
> +       if (!ASSERT_GE(nr_running, nr_running_exp, "nr_running"))
> +               goto out;
> +       if (!ASSERT_GE(psi_running, nr_running_exp, "psi_running"))
> +               goto out;
> +
> +       err = 0;
> +out:
> +       close(iter_fd);
> +       return err;
> +}

..
> +
> +       /* Case 1): Enable all possible CPUs */
> +       CPU_ZERO(&set);
> +       for (i = 0; i < nr_cpus; i++)
> +               CPU_SET(i, &set);
> +       err = sched_setaffinity(skel->bss->pid, sizeof(set), &set);
> +       if (!ASSERT_OK(err, "setaffinity_all_cpus"))
> +               goto free_link;
> +       err = read_percpu_data(link, nr_cpus, 1);
> +       if (!ASSERT_OK(err, "read_percpu_data"))
> +               goto free_link;

The patch 1 looks good, but this test fails on s390.

read_percpu_data:FAIL:nr_cpus unexpected nr_cpus: actual 0 != expected 2
verify_iter_success:FAIL:read_percpu_data unexpected error: -1 (errno 95)

Please see CI.

So either add it to DENYLIST.s390x in the same commit or make it work.

pw-bot: cr

Re: [PATCH bpf-next 2/2] selftests/bpf: Add selftest for bits iter

[Yafang Shao]

On Fri, Feb 23, 2024 at 1:36 AM Alexei Starovoitov
<alexei.starovoitov@gmail.com> wrote:
>
> On Sun, Feb 18, 2024 at 3:49 AM Yafang Shao <laoar.shao@gmail.com> wrote:
> >
> > Add selftests for the newly added bits iter.
> > - bits_iter_success
> >   - The number of CPUs should be expected when iterating over a cpumask
> >   - percpu data extracted from the percpu struct should be expected
> >   - RCU lock is not required
> >   - It is fine without calling bpf_iter_cpumask_next()
> >   - It can work as expected when invalid arguments are passed
> >
> > - bits_iter_failure
> >   - bpf_iter_bits_destroy() is required after calling
> >     bpf_iter_bits_new()
> >   - bpf_iter_bits_destroy() can only destroy an initialized iter
> >   - bpf_iter_bits_next() must use an initialized iter
> >
> > Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
> > ---
> >  tools/testing/selftests/bpf/config            |   1 +
> >  .../selftests/bpf/prog_tests/bits_iter.c      | 180 ++++++++++++++++++
> >  .../bpf/progs/test_bits_iter_failure.c        |  53 ++++++
> >  .../bpf/progs/test_bits_iter_success.c        | 146 ++++++++++++++
> >  4 files changed, 380 insertions(+)
> >  create mode 100644 tools/testing/selftests/bpf/prog_tests/bits_iter.c
> >  create mode 100644 tools/testing/selftests/bpf/progs/test_bits_iter_failure.c
> >  create mode 100644 tools/testing/selftests/bpf/progs/test_bits_iter_success.c
> >
> > diff --git a/tools/testing/selftests/bpf/config b/tools/testing/selftests/bpf/config
> > index 01f241ea2c67..dd4b0935e35f 100644
> > --- a/tools/testing/selftests/bpf/config
> > +++ b/tools/testing/selftests/bpf/config
> > @@ -78,6 +78,7 @@ CONFIG_NF_CONNTRACK_MARK=y
> >  CONFIG_NF_DEFRAG_IPV4=y
> >  CONFIG_NF_DEFRAG_IPV6=y
> >  CONFIG_NF_NAT=y
> > +CONFIG_PSI=y
> >  CONFIG_RC_CORE=y
> >  CONFIG_SECURITY=y
> >  CONFIG_SECURITYFS=y
> > diff --git a/tools/testing/selftests/bpf/prog_tests/bits_iter.c b/tools/testing/selftests/bpf/prog_tests/bits_iter.c
> > new file mode 100644
> > index 000000000000..778a7c942dba
> > --- /dev/null
> > +++ b/tools/testing/selftests/bpf/prog_tests/bits_iter.c
> > @@ -0,0 +1,180 @@
> > +// SPDX-License-Identifier: GPL-2.0
> > +/* Copyright (c) 2024 Yafang Shao <laoar.shao@gmail.com> */
> > +
> > +#define _GNU_SOURCE
> > +#include <sched.h>
> > +
> > +#include <test_progs.h>
> > +#include "test_bits_iter_success.skel.h"
> > +#include "test_bits_iter_failure.skel.h"
> > +#include "cgroup_helpers.h"
> > +
> > +static const char * const positive_testcases[] = {
> > +       "cpumask_iter",
> > +};
> > +
> > +static const char * const negative_testcases[] = {
> > +       "null_pointer",
> > +       "zero_bit",
> > +       "no_mem",
> > +       "invalid_bits"
> > +};
> > +
> > +static int read_percpu_data(struct bpf_link *link, int nr_cpu_exp, int nr_running_exp)
> > +{
> > +       int iter_fd, len, item, nr_running, psi_running, nr_cpus, err = -1;
> > +       char buf[128];
> > +       size_t left;
> > +       char *p;
> > +
> > +       iter_fd = bpf_iter_create(bpf_link__fd(link));
> > +       if (!ASSERT_GE(iter_fd, 0, "iter_fd"))
> > +               return -1;
> > +
> > +       memset(buf, 0, sizeof(buf));
> > +       left = ARRAY_SIZE(buf);
> > +       p = buf;
> > +       while ((len = read(iter_fd, p, left)) > 0) {
> > +               p += len;
> > +               left -= len;
> > +       }
> > +
> > +       item = sscanf(buf, "nr_running %u nr_cpus %u psi_running %u\n",
> > +                     &nr_running, &nr_cpus, &psi_running);
> > +       if (!ASSERT_EQ(item, 3, "seq_format"))
> > +               goto out;
> > +       if (!ASSERT_EQ(nr_cpus, nr_cpu_exp, "nr_cpus"))
> > +               goto out;
> > +       if (!ASSERT_GE(nr_running, nr_running_exp, "nr_running"))
> > +               goto out;
> > +       if (!ASSERT_GE(psi_running, nr_running_exp, "psi_running"))
> > +               goto out;
> > +
> > +       err = 0;
> > +out:
> > +       close(iter_fd);
> > +       return err;
> > +}
>
> ..
> > +
> > +       /* Case 1): Enable all possible CPUs */
> > +       CPU_ZERO(&set);
> > +       for (i = 0; i < nr_cpus; i++)
> > +               CPU_SET(i, &set);
> > +       err = sched_setaffinity(skel->bss->pid, sizeof(set), &set);
> > +       if (!ASSERT_OK(err, "setaffinity_all_cpus"))
> > +               goto free_link;
> > +       err = read_percpu_data(link, nr_cpus, 1);
> > +       if (!ASSERT_OK(err, "read_percpu_data"))
> > +               goto free_link;
>
> The patch 1 looks good, but this test fails on s390.
>
> read_percpu_data:FAIL:nr_cpus unexpected nr_cpus: actual 0 != expected 2
> verify_iter_success:FAIL:read_percpu_data unexpected error: -1 (errno 95)
>
> Please see CI.
>
> So either add it to DENYLIST.s390x in the same commit or make it work.
>
> pw-bot: cr

The reason for the failure on s390x architecture is currently unclear.
One plausible explanation is that total_nr_cpus remains 0 when
executing the following code:

    bpf_for_each(bits, cpu, p->cpus_ptr, total_nr_cpus)

This is despite setting total_nr_cpus to the value obtained from
libbpf_num_possible_cpus():

    skel->bss->total_nr_cpus = libbpf_num_possible_cpus();

A potential workaround could involve using a hardcoded number of CPUs,
such as 8192, instead of relying on total_nr_cpus. This approach might
mitigate the issue temporarily.

-- 
Regards
Yafang

Re: [PATCH bpf-next 2/2] selftests/bpf: Add selftest for bits iter

[Eduard Zingerman]

On Fri, 2024-02-23 at 10:29 +0800, Yafang Shao wrote:

[...]

> > The patch 1 looks good, but this test fails on s390.
> > 
> > read_percpu_data:FAIL:nr_cpus unexpected nr_cpus: actual 0 != expected 2
> > verify_iter_success:FAIL:read_percpu_data unexpected error: -1 (errno 95)
> > 
> > Please see CI.
> > 
> > So either add it to DENYLIST.s390x in the same commit or make it work.
> > 
> > pw-bot: cr
> 
> The reason for the failure on s390x architecture is currently unclear.
> One plausible explanation is that total_nr_cpus remains 0 when
> executing the following code:
> 
>     bpf_for_each(bits, cpu, p->cpus_ptr, total_nr_cpus)
> 
> This is despite setting total_nr_cpus to the value obtained from
> libbpf_num_possible_cpus():
> 
>     skel->bss->total_nr_cpus = libbpf_num_possible_cpus();
> 
> A potential workaround could involve using a hardcoded number of CPUs,
> such as 8192, instead of relying on total_nr_cpus. This approach might
> mitigate the issue temporarily. 

I'm sorry, but is it really necessary to deal with total number of
CPUs in a test for bit iterator?
Tbh, cpumask_iter / verify_iter_success seem to be over-complicated.
Would it be possible to reuse test_loader.c's RUN_TESTS for this feature?
It supports __retval(...) annotation, so it should be possible to:
- create a map (even a constant map) with some known data;
- peek a BPF program type that supports BPF_PROG_TEST_RUN syscall command;
- organize test BPF programs so that they create bit iterators for
  this test data and return some expected quantities (e.g. a sum),
  verified by __retval.

This should limit the amount of code on prog_tests/*.c side
to the bare minimum.

Re: [PATCH bpf-next 2/2] selftests/bpf: Add selftest for bits iter

[Yafang Shao]

On Fri, Feb 23, 2024 at 7:52 PM Eduard Zingerman <eddyz87@gmail.com> wrote:
>
> On Fri, 2024-02-23 at 10:29 +0800, Yafang Shao wrote:
>
> [...]
>
> > > The patch 1 looks good, but this test fails on s390.
> > >
> > > read_percpu_data:FAIL:nr_cpus unexpected nr_cpus: actual 0 != expected 2
> > > verify_iter_success:FAIL:read_percpu_data unexpected error: -1 (errno 95)
> > >
> > > Please see CI.
> > >
> > > So either add it to DENYLIST.s390x in the same commit or make it work.
> > >
> > > pw-bot: cr
> >
> > The reason for the failure on s390x architecture is currently unclear.
> > One plausible explanation is that total_nr_cpus remains 0 when
> > executing the following code:
> >
> >     bpf_for_each(bits, cpu, p->cpus_ptr, total_nr_cpus)
> >
> > This is despite setting total_nr_cpus to the value obtained from
> > libbpf_num_possible_cpus():
> >
> >     skel->bss->total_nr_cpus = libbpf_num_possible_cpus();
> >
> > A potential workaround could involve using a hardcoded number of CPUs,
> > such as 8192, instead of relying on total_nr_cpus. This approach might
> > mitigate the issue temporarily.
>
> I'm sorry, but is it really necessary to deal with total number of
> CPUs in a test for bit iterator?

The CPU number verification is served to validate the functionality of
bpf_iter_bits_next(). However, I believe we can streamline the logic
by removing the surrounding code.

> Tbh, cpumask_iter / verify_iter_success seem to be over-complicated.
> Would it be possible to reuse test_loader.c's RUN_TESTS for this feature?
> It supports __retval(...) annotation, so it should be possible to:
> - create a map (even a constant map) with some known data;
> - peek a BPF program type that supports BPF_PROG_TEST_RUN syscall command;
> - organize test BPF programs so that they create bit iterators for
>   this test data and return some expected quantities (e.g. a sum),
>   verified by __retval.
>
> This should limit the amount of code on prog_tests/*.c side
> to the bare minimum.

Thank you for your suggestion. I will consider it carefully.

-- 
Regards
Yafang

Re: [PATCH bpf-next 2/2] selftests/bpf: Add selftest for bits iter

[kernel test robot]

Hi Yafang,

kernel test robot noticed the following build errors:

[auto build test ERROR on bpf-next/master]

url:    https://github.com/intel-lab-lkp/linux/commits/Yafang-Shao/bpf-Add-bits-iterator/20240218-195123
base:   https://git.kernel.org/pub/scm/linux/kernel/git/bpf/bpf-next.git master
patch link:    https://lore.kernel.org/r/20240218114818.13585-3-laoar.shao%40gmail.com
patch subject: [PATCH bpf-next 2/2] selftests/bpf: Add selftest for bits iter
compiler: gcc-12 (Debian 12.2.0-14) 12.2.0
reproduce (this is a W=1 build): (https://download.01.org/0day-ci/archive/20240226/202402260354.rwSw5NBc-lkp@intel.com/reproduce)

If you fix the issue in a separate patch/commit (i.e. not just a new version of
the same patch/commit), kindly add following tags
| Reported-by: kernel test robot <lkp@intel.com>
| Closes: https://lore.kernel.org/oe-kbuild-all/202402260354.rwSw5NBc-lkp@intel.com/

All errors (new ones prefixed by >>):

>> progs/test_bits_iter_success.c:55:27: error: incomplete definition of type 'struct psi_group_cpu'
      55 |                 psi_nr_running += groupc->tasks[NR_RUNNING];
         |                                   ~~~~~~^
   progs/test_bits_iter_success.c:13:21: note: forward declaration of 'struct psi_group_cpu'
      13 | extern const struct psi_group_cpu system_group_pcpu __ksym __weak;
         |                     ^
>> progs/test_bits_iter_success.c:55:35: error: use of undeclared identifier 'NR_RUNNING'; did you mean 'T_RUNNING'?
      55 |                 psi_nr_running += groupc->tasks[NR_RUNNING];
         |                                                 ^~~~~~~~~~
         |                                                 T_RUNNING
   /tools/include/vmlinux.h:48349:3: note: 'T_RUNNING' declared here
    48349 |                 T_RUNNING = 0,
          |                 ^
   2 errors generated.

-- 
0-DAY CI Kernel Test Service
https://github.com/intel/lkp-tests/wiki