Use context_header for tail call perf improvement (#3698)

* fix * fix * fix * fix tests * fix tests * change to capabilities, add check for prog_array map * fix analysis build, add tests * update documentation * add context_header_support callback for maps * CR comment * change capabilities to bitfield * Apply suggestions from code review Co-authored-by: Dave Thaler <[email protected]> * cr comments * cr comments * Apply suggestions from code review Co-authored-by: Dave Thaler <[email protected]> * CR comments --------- Co-authored-by: Dave Thaler <[email protected]>
microsoft · Jul 15, 2024 · 00c9138 · 00c9138
1 parent cad5807
commit 00c9138
Show file tree

Hide file tree

Showing 24 changed files with 706 additions and 184 deletions.
diff --git a/docs/eBpfExtensions.md b/docs/eBpfExtensions.md
@@ -100,6 +100,43 @@ structure from provided data and context buffers.
 * `context_destroy`: Pointer to `ebpf_program_context_destroy_t` function that destroys a program type specific
 context structure and populates the returned data and context buffers.
 * `required_irql`: IRQL at which the eBPF program is invoked by bpf_prog_test_run_opts.
+* `capabilities`: 32-bit integer describing the optional capabilities / features supported by the extension.
+   * `supports_context_header`: Flag indicating extension supports adding a context header at the start of each context passed to the eBPF program.
+
+**Capabilities**
+
+`supports_context_header`:
+
+Flag indicating that extension supports adding a context header at the start of each context passed to the eBPF program.
+An extension can choose to opt in to support context header at the start of each program context structure that is
+passed to the eBPF program. To support this feature, the extension can use the macro `EBPF_CONTEXT_HEADER` to include
+the context header at the start of the program context structure. Even when the context header is added, the pointer
+passed to the eBPF program is after the context header.
+
+*Example*
+
+Below is an example of a sample extension where it is now including eBPF context header at the start of the original
+context structure:
+
+```c
+// Original sample extension program context that is passed to the eBPF program.
+typedef struct _sample_program_context
+{
+    uint8_t* data_start;
+    uint8_t* data_end;
+    uint32_t uint32_data;
+    uint16_t uint16_data;
+} sample_program_context_t;
+
+// Program context including the context header.
+typedef struct _sample_program_context_header
+{
+    EBPF_CONTEXT_HEADER;
+    sample_program_context_t context;
+} sample_program_context_header_t;
+```
+The extension passes a pointer to `context` inside `sample_program_context_header_t`, and not a pointer to
+`sample_program_context_header_t`, when invoking the eBPF program.
 
 #### `ebpf_program_info_t` Struct
 The various fields of this structure should be set as follows:

diff --git a/include/ebpf_extension.h b/include/ebpf_extension.h
@@ -125,3 +125,5 @@ typedef struct _ebpf_execution_context_state
         uint32_t count;
     } tail_call_state;
 } ebpf_execution_context_state_t;
+
+#define EBPF_CONTEXT_HEADER uint64_t context_header[8]
diff --git a/include/ebpf_program_types.h b/include/ebpf_program_types.h
@@ -73,6 +73,17 @@ typedef void (*ebpf_program_context_destroy_t)(
     _Out_writes_bytes_to_opt_(*context_size_out, *context_size_out) uint8_t* context_out,
     _Inout_ size_t* context_size_out);
 
+#pragma warning(push)
+#pragma warning(disable : 4201) // nonstandard extension used: nameless struct/union
+typedef union _program_data_capabilities
+{
+    uint32_t value;
+    struct
+    {
+        bool supports_context_header : 1; // Program supports context header.
+    };
+} program_data_capabilities_t;
+
 // This is the type definition for the eBPF program data
 // when version is EBPF_PROGRAM_DATA_CURRENT_VERSION.
 typedef struct _ebpf_program_data
@@ -87,7 +98,12 @@ typedef struct _ebpf_program_data
     ebpf_program_context_create_t context_create;   ///< Pointer to context create function.
     ebpf_program_context_destroy_t context_destroy; ///< Pointer to context destroy function.
     uint8_t required_irql;                          ///< IRQL at which the program is invoked.
+    program_data_capabilities_t capabilities;       ///< Capabilities supported by the program information provider.
 } ebpf_program_data_t;
+#pragma warning(pop)
+
+static_assert(
+    EBPF_FIELD_SIZE(ebpf_program_data_t, capabilities) == sizeof(uint32_t), "Size of capabilities is 32 bits.");
 
 // This is the type definition for the eBPF program section information
 // when version is EBPF_PROGRAM_SECTION_INFORMATION_CURRENT_VERSION.

diff --git a/include/ebpf_windows.h b/include/ebpf_windows.h
@@ -124,7 +124,7 @@ typedef enum _ebpf_helper_function
     }
 
 #define EBPF_PROGRAM_DATA_CURRENT_VERSION 1
-#define EBPF_PROGRAM_DATA_CURRENT_VERSION_SIZE EBPF_SIZE_INCLUDING_FIELD(ebpf_program_data_t, required_irql)
+#define EBPF_PROGRAM_DATA_CURRENT_VERSION_SIZE EBPF_SIZE_INCLUDING_FIELD(ebpf_program_data_t, capabilities)
 #define EBPF_PROGRAM_DATA_CURRENT_VERSION_TOTAL_SIZE sizeof(ebpf_program_data_t)
 #define EBPF_PROGRAM_DATA_HEADER                                                   \
     {                                                                              \

diff --git a/libs/execution_context/ebpf_core.c b/libs/execution_context/ebpf_core.c
@@ -2115,14 +2115,12 @@ _ebpf_core_map_find_and_delete_element(_Inout_ ebpf_map_t* map, _In_ const uint8
 static int64_t
 _ebpf_core_tail_call(void* context, ebpf_map_t* map, uint32_t index)
 {
-    UNREFERENCED_PARAMETER(context);
-
     // Get program from map[index].
     ebpf_program_t* callee = ebpf_map_get_program_from_entry(map, sizeof(index), (uint8_t*)&index);
     if (callee == NULL) {
         return -EBPF_INVALID_ARGUMENT;
     }
-    return -ebpf_program_set_tail_call(callee);
+    return -ebpf_program_set_tail_call(context, callee);
 }
 
 static uint32_t

diff --git a/libs/execution_context/ebpf_link.c b/libs/execution_context/ebpf_link.c
@@ -75,19 +75,37 @@ static ebpf_result_t
 _ebpf_link_instance_invoke(
     _In_ const void* extension_client_binding_context, _Inout_ void* program_context, _Out_ uint32_t* result);
 
+static ebpf_result_t
+_ebpf_link_instance_invoke_with_context_header(
+    _In_ const void* extension_client_binding_context, _Inout_ void* program_context, _Out_ uint32_t* result);
+
 static ebpf_result_t
 _ebpf_link_instance_invoke_batch_begin(size_t state_size, _Out_writes_(state_size) void* state);
 
+static ebpf_result_t
+_ebpf_link_instance_invoke_batch_begin_with_context_header(size_t state_size, _Out_writes_(state_size) void* state);
+
 static ebpf_result_t
 _ebpf_link_instance_invoke_batch(
     _In_ const void* extension_client_binding_context,
     _Inout_ void* program_context,
     _Out_ uint32_t* result,
     _In_ const void* state);
 
+static ebpf_result_t
+_ebpf_link_instance_invoke_batch_with_context_header(
+    _In_ const void* extension_client_binding_context,
+    _Inout_ void* program_context,
+    _Out_ uint32_t* result,
+    _In_ const void* state);
+
 static ebpf_result_t
 _ebpf_link_instance_invoke_batch_end(_Inout_ void* state);
 
+static ebpf_result_t
+_ebpf_link_instance_invoke_batch_end_with_context_header(_Inout_ void* state);
+
+// Dispatch table when program information provider does not support context header.
 static const ebpf_extension_program_dispatch_table_t _ebpf_link_dispatch_table = {
     EBPF_LINK_DISPATCH_TABLE_VERSION_CURRENT,
     EBPF_LINK_DISPATCH_TABLE_FUNCTION_COUNT_CURRENT, // Count of functions. This should be updated when new functions
@@ -98,6 +116,17 @@ static const ebpf_extension_program_dispatch_table_t _ebpf_link_dispatch_table =
     _ebpf_link_instance_invoke_batch_end,
 };
 
+// Dispatch table when program information provider supports context header.
+static const ebpf_extension_program_dispatch_table_t _ebpf_link_dispatch_table_with_context_header = {
+    EBPF_LINK_DISPATCH_TABLE_VERSION_CURRENT,
+    EBPF_LINK_DISPATCH_TABLE_FUNCTION_COUNT_CURRENT, // Count of functions. This should be updated when new functions
+                                                     // are added.
+    _ebpf_link_instance_invoke_with_context_header,
+    _ebpf_link_instance_invoke_batch_begin_with_context_header,
+    _ebpf_link_instance_invoke_batch_with_context_header,
+    _ebpf_link_instance_invoke_batch_end_with_context_header,
+};
+
 // Assert that the invoke function is aligned with ebpf_extension_dispatch_table_t->function.
 C_ASSERT(
     EBPF_OFFSET_OF(ebpf_extension_dispatch_table_t, function) ==
@@ -115,6 +144,7 @@ _ebpf_link_client_attach_provider(
     ebpf_link_t* link = (ebpf_link_t*)client_context;
     void* provider_binding_context;
     void* provider_dispatch;
+    void* client_dispatch_table = NULL;
     const ebpf_attach_provider_data_t* attach_provider_data =
         (const ebpf_attach_provider_data_t*)provider_registration_instance->NpiSpecificCharacteristics;
 
@@ -180,11 +210,17 @@ _ebpf_link_client_attach_provider(
     link->link_type = attach_provider_data->link_type;
     link->bpf_attach_type = attach_provider_data->bpf_attach_type;
 
+    if (ebpf_program_supports_context_header(link->program)) {
+        client_dispatch_table = (void*)&_ebpf_link_dispatch_table_with_context_header;
+    } else {
+        client_dispatch_table = (void*)&_ebpf_link_dispatch_table;
+    }
+
     ebpf_lock_unlock(&link->lock, state);
     lock_held = false;
 
     status = NmrClientAttachProvider(
-        nmr_binding_handle, link, &_ebpf_link_dispatch_table, &provider_binding_context, &provider_dispatch);
+        nmr_binding_handle, link, client_dispatch_table, &provider_binding_context, &provider_dispatch);
 
     if (!NT_SUCCESS(status)) {
         EBPF_LOG_MESSAGE_NTSTATUS(
@@ -279,7 +315,8 @@ ebpf_link_create(
 
     // Note: This must be the last thing done in this function as it inserts the object into the global list.
     // After this point, the object can be accessed by other threads.
-    ebpf_result_t result = EBPF_OBJECT_INITIALIZE(&local_link->object, EBPF_OBJECT_LINK, _ebpf_link_free, NULL, NULL);
+    ebpf_result_t result =
+        EBPF_OBJECT_INITIALIZE(&local_link->object, EBPF_OBJECT_LINK, _ebpf_link_free, NULL, NULL, NULL);
     if (result != EBPF_SUCCESS) {
         EBPF_LOG_MESSAGE_ERROR(
             EBPF_TRACELOG_LEVEL_ERROR, EBPF_TRACELOG_KEYWORD_LINK, "ebpf_object_initialize failed for link", result);
@@ -523,7 +560,78 @@ _ebpf_link_instance_invoke_batch(
     ebpf_result_t return_value;
     ebpf_link_t* link = (ebpf_link_t*)client_binding_context;
 
-    return_value = ebpf_program_invoke(link->program, program_context, result, (ebpf_execution_context_state_t*)state);
+    return_value =
+        ebpf_program_invoke(link->program, false, program_context, result, (ebpf_execution_context_state_t*)state);
+
+    EBPF_RETURN_RESULT(return_value);
+}
+
+static ebpf_result_t
+_ebpf_link_instance_invoke_with_context_header(
+    _In_ const void* extension_client_binding_context, _Inout_ void* program_context, _Out_ uint32_t* result)
+{
+    ebpf_execution_context_state_t state = {0};
+    ebpf_result_t return_value;
+    return_value =
+        _ebpf_link_instance_invoke_batch_begin_with_context_header(sizeof(ebpf_execution_context_state_t), &state);
+
+    if (return_value != EBPF_SUCCESS) {
+        goto Done;
+    }
+
+    return_value = _ebpf_link_instance_invoke_batch_with_context_header(
+        extension_client_binding_context, program_context, result, &state);
+    (void)_ebpf_link_instance_invoke_batch_end(&state);
+
+Done:
+    return return_value;
+}
+
+static ebpf_result_t
+_ebpf_link_instance_invoke_batch_begin_with_context_header(size_t state_size, _Out_writes_(state_size) void* state)
+{
+    ebpf_execution_context_state_t* execution_context_state = (ebpf_execution_context_state_t*)state;
+    bool epoch_entered = false;
+    ebpf_result_t return_value = EBPF_SUCCESS;
+    if (state_size < sizeof(ebpf_execution_context_state_t)) {
+        return_value = EBPF_INVALID_ARGUMENT;
+        goto Done;
+    }
+
+    memset(execution_context_state, 0, sizeof(ebpf_execution_context_state_t));
+
+    ebpf_epoch_enter((ebpf_epoch_state_t*)(execution_context_state->epoch_state));
+    epoch_entered = true;
+
+Done:
+    if (return_value != EBPF_SUCCESS && epoch_entered) {
+        ebpf_epoch_exit((ebpf_epoch_state_t*)(execution_context_state->epoch_state));
+    }
+
+    return return_value;
+}
+
+static ebpf_result_t
+_ebpf_link_instance_invoke_batch_end_with_context_header(_Inout_ void* state)
+{
+    ebpf_execution_context_state_t* execution_context_state = (ebpf_execution_context_state_t*)state;
+    ebpf_epoch_exit((ebpf_epoch_state_t*)(execution_context_state->epoch_state));
+    return EBPF_SUCCESS;
+}
+
+static ebpf_result_t
+_ebpf_link_instance_invoke_batch_with_context_header(
+    _In_ const void* client_binding_context,
+    _Inout_ void* program_context,
+    _Out_ uint32_t* result,
+    _In_ const void* state)
+{
+    // No function entry exit traces as this is a high volume function.
+    ebpf_result_t return_value;
+    ebpf_link_t* link = (ebpf_link_t*)client_binding_context;
+
+    return_value =
+        ebpf_program_invoke(link->program, true, program_context, result, (ebpf_execution_context_state_t*)state);
 
     EBPF_RETURN_RESULT(return_value);
 }

diff --git a/libs/execution_context/ebpf_maps.c b/libs/execution_context/ebpf_maps.c
@@ -29,6 +29,7 @@ typedef struct _ebpf_core_object_map
     ebpf_lock_t lock;
     ebpf_map_definition_in_memory_t inner_template_map_definition;
     bool is_program_type_set;
+    bool supports_context_header;
     ebpf_program_type_t program_type;
 } ebpf_core_object_map_t;
 
@@ -160,7 +161,7 @@ typedef struct _ebpf_core_lru_map
 {
     ebpf_core_map_t core_map; //< Core map structure.
     size_t partition_count;   //< Number of LRU partitions. Limited to a maximum of EBPF_LRU_MAXIMUM_PARTITIONS.
-    uint8_t padding[24];      //< Required to ensure partitions are cache aligned.
+    uint8_t padding[16];      //< Required to ensure partitions are cache aligned.
     ebpf_lru_partition_t
         partitions[1]; //< Array of LRU partitions. Limited to a maximum of EBPF_LRU_MAXIMUM_PARTITIONS.
 } ebpf_core_lru_map_t;
@@ -345,6 +346,13 @@ _get_map_program_type(_In_ const ebpf_core_object_t* object)
     return map->program_type;
 }
 
+static bool
+_ebpf_map_get_program_context_header_support(_In_ const ebpf_core_object_t* object)
+{
+    const ebpf_core_object_map_t* map = (const ebpf_core_object_map_t*)object;
+    return map->supports_context_header;
+}
+
 typedef struct _ebpf_map_metadata_table
 {
     ebpf_map_type_t map_type;
@@ -658,17 +666,21 @@ _associate_program_with_prog_array_map(_Inout_ ebpf_core_map_t* map, _In_ const
     ebpf_assert(map->ebpf_map_definition.type == BPF_MAP_TYPE_PROG_ARRAY);
     ebpf_core_object_map_t* program_array = EBPF_FROM_FIELD(ebpf_core_object_map_t, core_map, map);
 
-    // Validate that the program type is
+    // Validate that the program type and context header support is
     // not in conflict with the map's program type.
     ebpf_program_type_t program_type = ebpf_program_type_uuid(program);
+    bool supports_context_header = ebpf_program_supports_context_header(program);
     ebpf_result_t result = EBPF_SUCCESS;
 
     ebpf_lock_state_t lock_state = ebpf_lock_lock(&program_array->lock);
 
     if (!program_array->is_program_type_set) {
         program_array->is_program_type_set = TRUE;
         program_array->program_type = program_type;
-    } else if (memcmp(&program_array->program_type, &program_type, sizeof(program_type)) != 0) {
+        program_array->supports_context_header = supports_context_header;
+    } else if (
+        memcmp(&program_array->program_type, &program_type, sizeof(program_type)) != 0 ||
+        program_array->supports_context_header != supports_context_header) {
         result = EBPF_INVALID_FD;
     }
 
@@ -708,10 +720,14 @@ static _Requires_lock_held_(object_map->lock) ebpf_result_t _validate_map_value_
     const ebpf_core_map_t* map = &object_map->core_map;
 
     ebpf_program_type_t value_program_type = {0};
+    bool value_supports_context_header = false;
     bool is_program_type_set = false;
 
     if (value_object->get_program_type) {
         value_program_type = value_object->get_program_type(value_object);
+        ebpf_assert(value_object->get_context_header_support != NULL);
+        __analysis_assume(value_object->get_context_header_support != NULL);
+        value_supports_context_header = value_object->get_context_header_support(value_object);
         is_program_type_set = true;
     }
 
@@ -729,7 +745,10 @@ static _Requires_lock_held_(object_map->lock) ebpf_result_t _validate_map_value_
         if (!object_map->is_program_type_set) {
             object_map->is_program_type_set = TRUE;
             object_map->program_type = value_program_type;
-        } else if (memcmp(&object_map->program_type, &value_program_type, sizeof(value_program_type)) != 0) {
+            object_map->supports_context_header = value_supports_context_header;
+        } else if (
+            memcmp(&object_map->program_type, &value_program_type, sizeof(value_program_type)) != 0 ||
+            object_map->supports_context_header != value_supports_context_header) {
             result = EBPF_INVALID_FD;
             goto Error;
         }
@@ -2401,7 +2420,10 @@ ebpf_map_create(
 
     const ebpf_map_metadata_table_t* table = &ebpf_map_metadata_tables[local_map->ebpf_map_definition.type];
     ebpf_object_get_program_type_t get_program_type = (table->get_object_from_entry) ? _get_map_program_type : NULL;
-    result = EBPF_OBJECT_INITIALIZE(&local_map->object, EBPF_OBJECT_MAP, _ebpf_map_delete, NULL, get_program_type);
+    ebpf_object_get_context_header_support_t get_context_header_support =
+        (table->get_object_from_entry) ? _ebpf_map_get_program_context_header_support : NULL;
+    result = EBPF_OBJECT_INITIALIZE(
+        &local_map->object, EBPF_OBJECT_MAP, _ebpf_map_delete, NULL, get_program_type, get_context_header_support);
     if (result != EBPF_SUCCESS) {
         goto Exit;
     }