[PATCH] selftests/bpf: Fix map_lookup_percpu_elem on sparse CPU IDs

[Suchit Karunakaran <suchitkarunakaran@gmail.com>]

libbpf_num_possible_cpus() returns the number of possible CPUs, which is
appropriate for sizing packed per-CPU map value buffers. It is not the
upper bound for logical CPU IDs.

For a possible CPU mask such as 0,2-3, the test loops over CPU IDs 0
through 2. This incorrectly visits CPU 1 and misses CPU 3. It also
initializes packed per-CPU slots using their slot indexes rather than
the corresponding logical CPU IDs.

Parse the possible CPU mask and keep the packed slot count separate
from the logical CPU ID range. Populate each dense per-CPU slot with
its logical CPU ID, calculate the corresponding expected sum, and make
the BPF program iterate over the full CPU ID range.

Fixes: 7aa424e02a04bba5ecc84afe9b58b16e9e0b34f8 ("selftests/bpf: Fix some bugs in map_lookup_percpu_elem testcase")

Signed-off-by: Suchit Karunakaran <suchitkarunakaran@gmail.com>
---
 .../bpf/prog_tests/map_lookup_percpu_elem.c   | 34 ++++++++++++++-----
 .../bpf/progs/test_map_lookup_percpu_elem.c   |  8 ++---
 2 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/tools/testing/selftests/bpf/prog_tests/map_lookup_percpu_elem.c b/tools/testing/selftests/bpf/prog_tests/map_lookup_percpu_elem.c
index bfb1bf3fd427..849f33259c00 100644
--- a/tools/testing/selftests/bpf/prog_tests/map_lookup_percpu_elem.c
+++ b/tools/testing/selftests/bpf/prog_tests/map_lookup_percpu_elem.c
@@ -2,29 +2,45 @@
 /* Copyright (c) 2022 Bytedance */
 
 #include <test_progs.h>
+#include "bpf/libbpf_internal.h"
 #include "test_map_lookup_percpu_elem.skel.h"
 
 void test_map_lookup_percpu_elem(void)
 {
 	struct test_map_lookup_percpu_elem *skel;
-	__u64 key = 0, sum;
-	int ret, i, nr_cpus = libbpf_num_possible_cpus();
+	bool *possible = NULL;
+	__u64 key = 0, sum = 0;
+	int cpu, nr_cpu_ids, nr_cpus, ret, slot = 0;
 	__u64 *buf;
 
-	buf = malloc(nr_cpus*sizeof(__u64));
-	if (!ASSERT_OK_PTR(buf, "malloc"))
+	ret = parse_cpu_mask_file("/sys/devices/system/cpu/possible", &possible,
+				  &nr_cpu_ids);
+	if (!ASSERT_OK(ret, "parse possible CPU mask"))
 		return;
 
-	for (i = 0; i < nr_cpus; i++)
-		buf[i] = i;
-	sum = (nr_cpus - 1) * nr_cpus / 2;
+	nr_cpus = libbpf_num_possible_cpus();
+	if (!ASSERT_GT(nr_cpus, 0, "libbpf_num_possible_cpus"))
+		goto free_mask;
+
+	buf = malloc(nr_cpus * sizeof(*buf));
+	if (!ASSERT_OK_PTR(buf, "malloc"))
+		goto free_mask;
+
+	for (cpu = 0; cpu < nr_cpu_ids; cpu++) {
+		if (!possible[cpu])
+			continue;
+		buf[slot++] = cpu;
+		sum += cpu;
+	}
+	if (!ASSERT_EQ(slot, nr_cpus, "possible CPU mask weight"))
+		goto exit;
 
 	skel = test_map_lookup_percpu_elem__open();
 	if (!ASSERT_OK_PTR(skel, "test_map_lookup_percpu_elem__open"))
 		goto exit;
 
 	skel->rodata->my_pid = getpid();
-	skel->rodata->nr_cpus = nr_cpus;
+	skel->rodata->nr_cpu_ids = nr_cpu_ids;
 
 	ret = test_map_lookup_percpu_elem__load(skel);
 	if (!ASSERT_OK(ret, "test_map_lookup_percpu_elem__load"))
@@ -55,4 +71,6 @@ void test_map_lookup_percpu_elem(void)
 	test_map_lookup_percpu_elem__destroy(skel);
 exit:
 	free(buf);
+free_mask:
+	free(possible);
 }
diff --git a/tools/testing/selftests/bpf/progs/test_map_lookup_percpu_elem.c b/tools/testing/selftests/bpf/progs/test_map_lookup_percpu_elem.c
index ca827b1092da..d8da0696b97c 100644
--- a/tools/testing/selftests/bpf/progs/test_map_lookup_percpu_elem.c
+++ b/tools/testing/selftests/bpf/progs/test_map_lookup_percpu_elem.c
@@ -7,7 +7,7 @@
 __u64 percpu_array_elem_sum = 0;
 __u64 percpu_hash_elem_sum = 0;
 __u64 percpu_lru_hash_elem_sum = 0;
-const volatile int nr_cpus;
+const volatile int nr_cpu_ids;
 const volatile int my_pid;
 
 struct {
@@ -57,17 +57,17 @@ int sysenter_getuid(const void *ctx)
 
 	map_ctx.map = &percpu_array_map;
 	map_ctx.sum = 0;
-	bpf_loop(nr_cpus, read_percpu_elem_callback, &map_ctx, 0);
+	bpf_loop(nr_cpu_ids, read_percpu_elem_callback, &map_ctx, 0);
 	percpu_array_elem_sum = map_ctx.sum;
 
 	map_ctx.map = &percpu_hash_map;
 	map_ctx.sum = 0;
-	bpf_loop(nr_cpus, read_percpu_elem_callback, &map_ctx, 0);
+	bpf_loop(nr_cpu_ids, read_percpu_elem_callback, &map_ctx, 0);
 	percpu_hash_elem_sum = map_ctx.sum;
 
 	map_ctx.map = &percpu_lru_hash_map;
 	map_ctx.sum = 0;
-	bpf_loop(nr_cpus, read_percpu_elem_callback, &map_ctx, 0);
+	bpf_loop(nr_cpu_ids, read_percpu_elem_callback, &map_ctx, 0);
 	percpu_lru_hash_elem_sum = map_ctx.sum;
 
 	return 0;
-- 
2.55.0