Add support for cache aligned epoch memory and switch maps to use it (#…

…4008) * Add support for cache aligned epoch memory and switch maps to use it Signed-off-by: Alan Jowett <[email protected]> * PR feedback Signed-off-by: Alan Jowett <[email protected]> * PR feedback Signed-off-by: Alan Jowett <[email protected]> --------- Signed-off-by: Alan Jowett <[email protected]> Co-authored-by: Alan Jowett <[email protected]>
microsoft · Nov 13, 2024 · f1b2c0f · f1b2c0f
1 parent 3bdeb43
commit f1b2c0f
Show file tree

Hide file tree

Showing 5 changed files with 94 additions and 8 deletions.
diff --git a/libs/execution_context/ebpf_maps.c b/libs/execution_context/ebpf_maps.c
@@ -984,7 +984,7 @@ _create_hash_map_internal(
     ebpf_core_map_t* local_map = NULL;
     *map = NULL;
 
-    local_map = ebpf_epoch_allocate_with_tag(map_struct_size, EBPF_POOL_TAG_MAP);
+    local_map = ebpf_epoch_allocate_cache_aligned_with_tag(map_struct_size, EBPF_POOL_TAG_MAP);
     if (local_map == NULL) {
         retval = EBPF_NO_MEMORY;
         goto Done;
@@ -1022,7 +1022,7 @@ _create_hash_map_internal(
         if (local_map && local_map->data) {
             ebpf_hash_table_destroy((ebpf_hash_table_t*)local_map->data);
         }
-        ebpf_epoch_free(local_map);
+        ebpf_epoch_free_cache_aligned(local_map);
         local_map = NULL;
     }
     return retval;
@@ -1044,7 +1044,7 @@ static void
 _delete_hash_map(_In_ _Post_invalid_ ebpf_core_map_t* map)
 {
     ebpf_hash_table_destroy((ebpf_hash_table_t*)map->data);
-    ebpf_epoch_free(map);
+    ebpf_epoch_free_cache_aligned(map);
 }
 
 static void
@@ -1408,7 +1408,7 @@ _create_lru_hash_map(
         if (lru_map && lru_map->core_map.data) {
             ebpf_hash_table_destroy((ebpf_hash_table_t*)lru_map->core_map.data);
         }
-        ebpf_epoch_free(lru_map);
+        ebpf_epoch_free_cache_aligned(lru_map);
         lru_map = NULL;
     }
 
@@ -1420,7 +1420,7 @@ _delete_lru_hash_map(_In_ _Post_invalid_ ebpf_core_map_t* map)
 {
     ebpf_core_lru_map_t* lru_map = EBPF_FROM_FIELD(ebpf_core_lru_map_t, core_map, map);
     ebpf_hash_table_destroy((ebpf_hash_table_t*)lru_map->core_map.data);
-    ebpf_epoch_free(map);
+    ebpf_epoch_free_cache_aligned(map);
 }
 
 static ebpf_result_t

diff --git a/libs/runtime/ebpf_epoch.c b/libs/runtime/ebpf_epoch.c
@@ -162,9 +162,10 @@ static KDPC _ebpf_epoch_timer_dpc;
  */
 typedef enum _ebpf_epoch_allocation_type
 {
-    EBPF_EPOCH_ALLOCATION_MEMORY,          ///< Memory allocation.
-    EBPF_EPOCH_ALLOCATION_WORK_ITEM,       ///< Work item.
-    EBPF_EPOCH_ALLOCATION_SYNCHRONIZATION, ///< Synchronization object.
+    EBPF_EPOCH_ALLOCATION_MEMORY,               ///< Memory allocation.
+    EBPF_EPOCH_ALLOCATION_WORK_ITEM,            ///< Work item.
+    EBPF_EPOCH_ALLOCATION_SYNCHRONIZATION,      ///< Synchronization object.
+    EBPF_EPOCH_ALLOCATION_MEMORY_CACHE_ALIGNED, ///< Memory allocation that is cache aligned.
 } ebpf_epoch_allocation_type_t;
 
 /**
@@ -177,6 +178,9 @@ typedef struct _ebpf_epoch_allocation_header
     ebpf_epoch_allocation_type_t entry_type; ///< Type of entry.
 } ebpf_epoch_allocation_header_t;
 
+static_assert(
+    sizeof(ebpf_epoch_allocation_header_t) < EBPF_CACHE_LINE_SIZE, "Header size must be less than cache line");
+
 /**
  * @brief This structure is used as a place holder when a custom action needs
  * to be performed on epoch end. Typically this is releasing memory that can't
@@ -455,6 +459,21 @@ _Ret_writes_maybenull_(size) void* ebpf_epoch_allocate(size_t size)
     return ebpf_epoch_allocate_with_tag(size, EBPF_POOL_TAG_EPOCH);
 }
 
+_Must_inspect_result_
+_Ret_writes_maybenull_(size) void* ebpf_epoch_allocate_cache_aligned_with_tag(size_t size, uint32_t tag)
+{
+    ebpf_assert(size);
+    ebpf_epoch_allocation_header_t* header;
+
+    size += EBPF_CACHE_LINE_SIZE;
+    header = (ebpf_epoch_allocation_header_t*)ebpf_allocate_cache_aligned_with_tag(size, tag);
+    if (header) {
+        header = (ebpf_epoch_allocation_header_t*)((uint8_t*)header + EBPF_CACHE_LINE_SIZE);
+    }
+
+    return header;
+}
+
 void
 ebpf_epoch_free(_Frees_ptr_opt_ void* memory)
 {
@@ -473,6 +492,24 @@ ebpf_epoch_free(_Frees_ptr_opt_ void* memory)
     _ebpf_epoch_insert_in_free_list(header);
 }
 
+void
+ebpf_epoch_free_cache_aligned(_Frees_ptr_opt_ void* memory)
+{
+    ebpf_epoch_allocation_header_t* header = (ebpf_epoch_allocation_header_t*)memory;
+
+    if (!memory) {
+        return;
+    }
+
+    header = (ebpf_epoch_allocation_header_t*)((uint8_t*)header - EBPF_CACHE_LINE_SIZE);
+
+    // Pool corruption or double free.
+    EBPF_EPOCH_FAIL_FAST(FAST_FAIL_HEAP_METADATA_CORRUPTION, header->freed_epoch == 0);
+    header->entry_type = EBPF_EPOCH_ALLOCATION_MEMORY_CACHE_ALIGNED;
+
+    _ebpf_epoch_insert_in_free_list(header);
+}
+
 ebpf_epoch_work_item_t*
 ebpf_epoch_allocate_work_item(_In_ void* callback_context, _In_ const void (*callback)(_Inout_ void* context))
 {
@@ -591,6 +628,9 @@ _ebpf_epoch_release_free_list(_Inout_ ebpf_epoch_cpu_entry_t* cpu_entry, int64_t
                 KeSetEvent(&synchronization->event, 0, false);
                 break;
             }
+            case EBPF_EPOCH_ALLOCATION_MEMORY_CACHE_ALIGNED:
+                ebpf_free_cache_aligned(header);
+                break;
             default:
                 // Pool corruption or internal error.
                 EBPF_EPOCH_FAIL_FAST(FAST_FAIL_CORRUPT_LIST_ENTRY, !"Invalid entry type");

diff --git a/libs/runtime/ebpf_epoch.h b/libs/runtime/ebpf_epoch.h
@@ -57,6 +57,15 @@ extern "C"
      */
     _Must_inspect_result_ _Ret_writes_maybenull_(size) void* ebpf_epoch_allocate(size_t size);
 
+    /**
+     * @brief Allocate cache aligned memory under epoch control.
+     * @param[in] size Size of memory to allocate
+     * @param[in] tag Pool tag to use.
+     * @returns Pointer to memory block allocated, or null on failure.
+     */
+    _Must_inspect_result_
+        _Ret_writes_maybenull_(size) void* ebpf_epoch_allocate_cache_aligned_with_tag(size_t size, uint32_t tag);
+
     /**
      * @brief Allocate memory under epoch control with tag.
      * @param[in] size Size of memory to allocate
@@ -73,6 +82,13 @@ extern "C"
     void
     ebpf_epoch_free(_Frees_ptr_opt_ void* memory);
 
+    /**
+     * @brief Free memory under epoch control.
+     * @param[in] memory Allocation to be freed once epoch ends.
+     */
+    void
+    ebpf_epoch_free_cache_aligned(_Frees_ptr_opt_ void* memory);
+
     /**
      * @brief Wait for the current epoch to end.
      */

diff --git a/libs/runtime/unit/platform_unit_test.cpp b/libs/runtime/unit/platform_unit_test.cpp
@@ -531,6 +531,23 @@ TEST_CASE("epoch_test_single_epoch", "[platform]")
     ebpf_epoch_synchronize();
 }
 
+TEST_CASE("epoch_test_single_epoch_cache_aligned", "[platform]")
+{
+    _test_helper test_helper;
+    test_helper.initialize();
+
+    ebpf_epoch_scope_t epoch_scope;
+    void* memory = ebpf_epoch_allocate_cache_aligned_with_tag(10, 0);
+    if (memory) {
+        memset(memory, 0, 10);
+    }
+
+    REQUIRE(memory == EBPF_CACHE_ALIGN_POINTER(memory));
+    ebpf_epoch_free_cache_aligned(memory);
+    epoch_scope.exit();
+    ebpf_epoch_synchronize();
+}
+
 TEST_CASE("epoch_test_two_threads", "[platform]")
 {
     _test_helper test_helper;

diff --git a/libs/shared/ebpf_shared_framework.h b/libs/shared/ebpf_shared_framework.h
@@ -74,6 +74,12 @@ ebpf_free(_Frees_ptr_opt_ void* pointer)
     cxplat_free(pointer, CXPLAT_POOL_FLAG_NON_PAGED, 0);
 }
 
+__forceinline void
+ebpf_free_cache_aligned(_Frees_ptr_opt_ void* pointer)
+{
+    cxplat_free(pointer, (cxplat_pool_flags_t)(CXPLAT_POOL_FLAG_NON_PAGED | CXPLAT_POOL_FLAG_CACHE_ALIGNED), 0);
+}
+
 #define ebpf_reallocate cxplat_reallocate
 
 /**
@@ -88,6 +94,13 @@ __forceinline __drv_allocatesMem(Mem) _Must_inspect_result_
     return cxplat_allocate(CXPLAT_POOL_FLAG_NON_PAGED, size, tag);
 }
 
+__forceinline __drv_allocatesMem(Mem) _Must_inspect_result_
+    _Ret_writes_maybenull_(size) void* ebpf_allocate_cache_aligned_with_tag(size_t size, uint32_t tag)
+{
+    return cxplat_allocate(
+        (cxplat_pool_flags_t)(CXPLAT_POOL_FLAG_NON_PAGED | CXPLAT_POOL_FLAG_CACHE_ALIGNED), size, tag);
+}
+
 #define ebpf_safe_size_t_add(augend, addend, result) \
     ebpf_result_from_cxplat_status(cxplat_safe_size_t_add(augend, addend, result))
 #define ebpf_safe_size_t_subtract(minuend, subtrahend, result) \