Fix bpf_map_lookup_and_delete_batch API (#3688)

* Initial commit

* merge main

* Initial commit

* Initial commit

* Initial commit

* Initial commit

* Fixed count_returned is 0 condition handling in the base code

* Changed the log trace back to EBPF_LOG_MESSAGE_UINT64

* Add validation of fetched_keys and fetched_values retrieved from lookup_and_delete_batch function

* Addressed PR comment

* Addressed PR comment

libs/api/ebpf_api.cpp

@@ -601,17 +601,28 @@ _ebpf_map_lookup_element_batch_helper(
         }
         count_returned += entries_returned;
 
-        // Point previous_key to the last key in the batch.
-        previous_key = (uint8_t*)keys + (count_returned - 1) * key_size;
+        // Set the previous_key for the next batch.
+        if (find_and_delete || count_returned == 0) {
+            // If find_and_delete is set to true, the entries are already deleted.
+            // If no more entries are returned, there is no need to fetch more.
+            // Hence the previous_key is set to null.
+            previous_key = nullptr;
+        } else {
+            // Point previous_key to the last key in the batch.
+            previous_key = (uint8_t*)keys + (count_returned - 1) * key_size;
+        }
 
         // Partial return signals last no more entries.
         if (entries_returned != entries_to_fetch) {
             break;
         }
     }
 
-    // Copy previous key into out_batch.
-    std::copy(previous_key, previous_key + key_size, (uint8_t*)out_batch);
+    memset((uint8_t*)out_batch, 0, key_size);
+    // Copy previous key into out_batch, if find_and_delete is false.
+    if (!find_and_delete && (previous_key != nullptr)) {
+        std::copy(previous_key, previous_key + key_size, (uint8_t*)out_batch);
+    }
 
     *count = static_cast<uint32_t>(count_returned);
 

libs/execution_context/ebpf_maps.c

@@ -496,7 +496,7 @@ _find_array_map_entry_with_reference(
         return result;
     }
 
-    ebpf_id_t *id = (ebpf_id_t*)(*data);
+    ebpf_id_t* id = (ebpf_id_t*)(*data);
     if (id != NULL && *id == 0) {
         // Turn zero ID into EBPF_OBJECT_NOT_FOUND.
         *data = NULL;
@@ -2922,10 +2922,12 @@ ebpf_map_get_next_key_and_value_batch(
 
         memcpy(key_and_value + output_length + key_size, next_value, value_size);
 
-        if (flags & EBPF_MAP_FIND_FLAG_DELETE) {
-            // If the caller requested deletion, delete the entry.
+        if ((flags & EBPF_MAP_FIND_FLAG_DELETE) && (previous_key != NULL)) {
+            // If the caller requested deletion, delete the previous entry.
             result = ebpf_map_metadata_tables[map->ebpf_map_definition.type].delete_entry(map, previous_key);
             if (result != EBPF_SUCCESS) {
+                EBPF_LOG_MESSAGE_UINT64(
+                    EBPF_TRACELOG_LEVEL_ERROR, EBPF_TRACELOG_KEYWORD_MAP, "Failed to delete entry", result);
                 break;
             }
         }
@@ -2943,5 +2945,16 @@ ebpf_map_get_next_key_and_value_batch(
 
     *key_and_value_length = output_length;
 
+    if ((flags & EBPF_MAP_FIND_FLAG_DELETE) && (previous_key != NULL) && (output_length != 0)) {
+        // If the caller requested deletion, delete the last entry.
+        ebpf_result_t delete_result =
+            ebpf_map_metadata_tables[map->ebpf_map_definition.type].delete_entry(map, previous_key);
+        if (delete_result != EBPF_SUCCESS) {
+            EBPF_LOG_MESSAGE_UINT64(
+                EBPF_TRACELOG_LEVEL_ERROR, EBPF_TRACELOG_KEYWORD_MAP, "Failed to delete last entry", delete_result);
+            result = delete_result;
+        }
+    }
+
     return result;
 }

tests/unit/libbpf_test.cpp

@@ -3632,14 +3632,15 @@ _test_maps_batch(bpf_map_type map_type)
         bpf_map_lookup_batch(
             map_fd, nullptr, &next_key, fetched_keys.data(), fetched_values.data(), &fetched_batch_size, &opts) == 0);
     REQUIRE(fetched_batch_size == batch_size);
+    _test_batch_iteration_maps(map_fd, batch_size, &opts, value_size, num_of_cpus);
 
     // Request more keys than present.
     uint32_t large_fetched_batch_size = fetched_batch_size * 2;
     REQUIRE(
         bpf_map_lookup_batch(
             map_fd, nullptr, &next_key, fetched_keys.data(), fetched_values.data(), &large_fetched_batch_size, &opts) ==
         0);
-    REQUIRE(fetched_batch_size == batch_size);
+    REQUIRE(large_fetched_batch_size == batch_size);
 
     // Search at end of map.
     REQUIRE(
@@ -3652,18 +3653,56 @@ _test_maps_batch(bpf_map_type map_type)
             &large_fetched_batch_size,
             &opts) == -ENOENT);
 
-    // Fetch all keys in batches.
+    // Verify all keys and values in batches.
     _test_batch_iteration_maps(map_fd, batch_size, &opts, value_size, num_of_cpus);
 
-    // Delete the batch.
+    // Delete all keys in one batch.
     uint32_t delete_batch_size = batch_size;
     opts.elem_flags = 0;
-
-    // Delete all keys in one batch.
     REQUIRE(bpf_map_delete_batch(map_fd, keys.data(), &delete_batch_size, &opts) == 0);
     REQUIRE(delete_batch_size == batch_size);
 
-    // Fetch all keys in one batch.
+    // Verify there are no entries, after deletion.
+    REQUIRE(
+        bpf_map_lookup_batch(
+            map_fd, nullptr, &next_key, fetched_keys.data(), fetched_values.data(), &fetched_batch_size, &opts) ==
+        -ENOENT);
+
+    // Lookup and Delete batch operation.
+    opts.elem_flags = BPF_NOEXIST;
+    update_batch_size = batch_size;
+
+    // Populate the map with the keys and values.
+    REQUIRE(bpf_map_update_batch(map_fd, keys.data(), values.data(), &update_batch_size, &opts) == 0);
+    REQUIRE(update_batch_size == batch_size);
+
+    next_key = 0;
+    opts.elem_flags = BPF_ANY;
+    fetched_batch_size = batch_size;
+    REQUIRE(
+        bpf_map_lookup_batch(
+            map_fd, nullptr, &next_key, fetched_keys.data(), fetched_values.data(), &fetched_batch_size, &opts) == 0);
+    REQUIRE(fetched_batch_size == batch_size);
+    _test_batch_iteration_maps(map_fd, batch_size, &opts, value_size, num_of_cpus);
+    // Verify the fetched_keys and fetched_values are returned correctly.
+    std::sort(fetched_keys.begin(), fetched_keys.end());
+    REQUIRE(fetched_keys == keys);
+    std::sort(fetched_values.begin(), fetched_values.end());
+    REQUIRE(fetched_values == values);
+
+    next_key = 0;
+    REQUIRE(
+        bpf_map_lookup_and_delete_batch(
+            map_fd, nullptr, &next_key, fetched_keys.data(), fetched_values.data(), &fetched_batch_size, &opts) == 0);
+    REQUIRE(fetched_batch_size == batch_size);
+    REQUIRE(next_key == 0);
+    // Verify the fetched_keys and fetched_values are returned correctly.
+    std::sort(fetched_keys.begin(), fetched_keys.end());
+    REQUIRE(fetched_keys == keys);
+    std::sort(fetched_values.begin(), fetched_values.end());
+    REQUIRE(fetched_values == values);
+
+    // Verify there are no entries, after deletion.
     REQUIRE(
         bpf_map_lookup_batch(
             map_fd, nullptr, &next_key, fetched_keys.data(), fetched_values.data(), &fetched_batch_size, &opts) ==