[bpf_map_lookup_batch] Add test case to retrieve data in batches

tests/unit/libbpf_test.cpp

@@ -3233,14 +3233,66 @@ TEST_CASE("bind_tail_call_max_exceed", "[libbpf]")
     usersim_trace_logging_set_enabled(false, 0, 0);
 }
 
-TEST_CASE("libbpf map batch", "[libbpf]")
+void
+_test_batch_iteration_maps(fd_t& map_fd, uint32_t batch_size, bpf_map_batch_opts* opts)
+{
+    // Retrieve in batches.
+    const uint8_t batch_size_divisor = 10;
+    uint32_t* in_batch = nullptr;
+    uint32_t out_batch = 0;
+    std::vector<uint32_t> returned_keys(0);
+    std::vector<uint64_t> returned_values(0);
+    uint32_t requested_batch_size = (batch_size / batch_size_divisor) - 2;
+    uint32_t batch_size_count = 0;
+    int result = 0;
+
+    for (;;) {
+        std::vector<uint32_t> batch_keys(requested_batch_size);
+        std::vector<uint64_t> batch_values(requested_batch_size);
+
+        batch_size_count = static_cast<uint32_t>(batch_keys.size());
+        result = bpf_map_lookup_batch(
+            map_fd, in_batch, &out_batch, batch_keys.data(), batch_values.data(), &batch_size_count, opts);
+        if (result == -ENOENT) {
+            printf("No more entries. End of map reached.\n");
+            break;
+        }
+
+        REQUIRE(result == 0);
+        // Number of entries retrieved (batch_size_count) should be less than or equal to requested_batch_size.
+        REQUIRE(batch_size_count <= static_cast<uint32_t>(batch_keys.size()));
+        REQUIRE(batch_size_count <= static_cast<uint32_t>(batch_values.size()));
+
+        batch_keys.resize(batch_size_count);
+        batch_values.resize(batch_size_count);
+        returned_keys.insert(returned_keys.end(), batch_keys.begin(), batch_keys.end());
+        returned_values.insert(returned_values.end(), batch_values.begin(), batch_values.end());
+
+        in_batch = &out_batch;
+    }
+
+    REQUIRE(returned_keys.size() == batch_size);
+    REQUIRE(returned_values.size() == batch_size);
+
+    std::sort(returned_keys.begin(), returned_keys.end());
+    std::sort(returned_values.begin(), returned_values.end());
+
+    // Verify the returned keys and values.
+    for (uint32_t i = 0; i < batch_size; i++) {
+        REQUIRE(returned_keys[i] == i);
+        REQUIRE(returned_values[i] == static_cast<uint64_t>(i) * 2ul);
+    }
+}
+
+void
+_test_maps_batch(bpf_map_type map_type)
 {
     _test_helper_end_to_end test_helper;
     test_helper.initialize();
 
     // Create a hash map.
     union bpf_attr attr = {};
-    attr.map_type = BPF_MAP_TYPE_HASH;
+    attr.map_type = map_type;
     attr.key_size = sizeof(uint32_t);
     attr.value_size = sizeof(uint64_t);
     attr.max_entries = 1024 * 1024;
@@ -3297,6 +3349,9 @@ TEST_CASE("libbpf map batch", "[libbpf]")
             &large_fetched_batch_size,
             &opts) == -ENOENT);
 
+    // Fetch all keys in batches.
+    _test_batch_iteration_maps(map_fd, batch_size, &opts);
+
     // Delete the batch.
     uint32_t delete_batch_size = batch_size;
     opts.elem_flags = 0;
@@ -3379,6 +3434,10 @@ TEST_CASE("libbpf map batch", "[libbpf]")
     REQUIRE(bpf_map_delete_batch(invalid_map_fd, keys.data(), &delete_batch_size, &opts) == -EBADF);
 }
 
+TEST_CASE("libbpf hash map batch", "[libbpf]") { _test_maps_batch(BPF_MAP_TYPE_HASH); }
+
+TEST_CASE("libbpf lru hash map batch", "[libbpf]") { _test_maps_batch(BPF_MAP_TYPE_LRU_HASH); }
+
 void
 _hash_of_map_initial_value_test(ebpf_execution_type_t execution_type)
 {