Repository: microsoft/ebpf-for-windows (C)

Issue:
eBPF core should expose batched invoke APIs
https://github.com/microsoft/xdp-for-windows invokes an eBPF program in a tight loop for each packet.

There is overhead involved in invoking an eBPF program (acquiring/release rundown protection on the program information provider, joining/leaving the epoch).

It would make more sense to expose an advanced dispatch API that has the following functions:
1) _ebpf_link_instance_invoke_batch_begin
2) _ebpf_link_instance_invoke_batch
3) _ebpf_link_instance_invoke_end

Extension would then have the option of either calling the non-batching API, paying the cost overhead and invoking the BPF program once or calling the _batch_begin, then calling the BPF program for multiple times and finally calling the _batch_end API.


Discussion: @mtfriesen do you think this will resolve your concerns around eBPF invocation overhead?



Fix (patch):
diff --git a/libs/execution_context/ebpf_link.c b/libs/execution_context/ebpf_link.c
index 3d090eb765..3e74c30938 100644
--- a/libs/execution_context/ebpf_link.c
+++ b/libs/execution_context/ebpf_link.c
@@ -45,11 +45,46 @@ static ebpf_result_t
 _ebpf_link_instance_invoke(
     _In_ const void* extension_client_binding_context, _Inout_ void* program_context, _Out_ uint32_t* result);
 
-static struct
+static ebpf_result_t
+_ebpf_link_instance_invoke_batch_begin(
+    _In_ const void* extension_client_binding_context, 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_end(_In_ const void* extension_client_binding_context);
+
+typedef enum _ebpf_link_dispatch_table_version
+{
+    EBPF_LINK_DISPATCH_TABLE_VERSION_1 = 1, ///< Version 1 of the dispatch table. This is the original version. It
+                                            ///< does not support batch processing.
+    EBPF_LINK_DISPATCH_TABLE_VERSION_2 = 2, ///< Version 2 of the dispatch table. This version supports batch
+                                            ///< processing.
+    EBPF_LINK_DISPATCH_TABLE_VERSION = EBPF_LINK_DISPATCH_TABLE_VERSION_2, ///< Current version of the dispatch table.
+} ebpf_link_dispatch_table_version_t;
+
+const typedef struct _ebpf_link_dispatch_table
 {
-    size_t size;
-    _ebpf_extension_dispatch_function function[1];
-} _ebpf_link_dispatch_table = {1, {_ebpf_link_instance_invoke}};
+    ebpf_extension_dispatch_table_t;
+    _ebpf_extension_dispatch_function new_functions[];
+} ebpf_link_dispatch_table_t;
+
+static ebpf_link_dispatch_table_t _ebpf_link_dispatch_table = {
+    EBPF_LINK_DISPATCH_TABLE_VERSION,
+    4, // Count of functions. This should be updated when new functions are added.
+    _ebpf_link_instance_invoke,
+    _ebpf_link_instance_invoke_batch_begin,
+    _ebpf_link_instance_invoke_batch,
+    _ebpf_link_instance_invoke_batch_end,
+};
+
+// Assert that new_functions is aligned with ebpf_extension_dispatch_table_t->function.
+C_ASSERT(sizeof(ebpf_extension_dispatch_table_t) == EBPF_OFFSET_OF(ebpf_link_dispatch_table_t, new_functions));
 
 static void
 _ebpf_link_free(_Frees_ptr_ ebpf_core_object_t* object)
@@ -314,25 +349,95 @@ static ebpf_result_t
 _ebpf_link_instance_invoke(
     _In_ const void* extension_client_binding_context, _Inout_ void* program_context, _Out_ uint32_t* result)
 {
-    // No function entry exit traces as this is a high volume function.
+    ebpf_execution_context_state_t state = {0};
+    ebpf_result_t return_value;
+    return_value = _ebpf_link_instance_invoke_batch_begin(
+        extension_client_binding_context, sizeof(ebpf_execution_context_state_t), &state);
+
+    if (return_value != EBPF_SUCCESS) {
+        goto Done;
+    }
+
+    return_value = _ebpf_link_instance_invoke_batch(extension_client_binding_context, program_context, result, &state);
+    ebpf_assert(return_value == EBPF_SUCCESS);
+
+    return_value = _ebpf_link_instance_invoke_batch_end(extension_client_binding_context);
+    ebpf_assert(return_value == EBPF_SUCCESS);
+
+Done:
+    return return_value;
+}
+
+static ebpf_result_t
+_ebpf_link_instance_invoke_batch_begin(
+    _In_ const void* extension_client_binding_context, size_t state_size, _Out_writes_(state_size) void* state)
+{
+    bool epoch_entered = false;
+    bool provider_reference_held = false;
     ebpf_result_t return_value;
     ebpf_link_t* link = (ebpf_link_t*)ebpf_extension_get_client_context(extension_client_binding_context);
 
+    if (state_size < sizeof(ebpf_execution_context_state_t)) {
+        return_value = EBPF_INVALID_ARGUMENT;
+        goto Done;
+    }
+
     if (link == NULL) {
         GUID npi_id = ebpf_extension_get_provider_guid(extension_client_binding_context);
         EBPF_LOG_MESSAGE_GUID(
             EBPF_TRACELOG_LEVEL_WARNING, EBPF_TRACELOG_KEYWORD_LINK, "Client context is null", npi_id);
         return_value = EBPF_FAILED;
-        goto Exit;
+        goto Done;
     }
 
     return_value = ebpf_epoch_enter();
-    if (return_value != EBPF_SUCCESS)
-        goto Exit;
-    ebpf_program_invoke(link->program, program_context, result);
+    if (return_value != EBPF_SUCCESS) {
+        goto Done;
+    }
+    epoch_entered = true;
+
+    return_value = ebpf_program_reference_providers(link->program);
+    if (return_value != EBPF_SUCCESS) {
+        goto Done;
+    }
+    provider_reference_held = true;
+
+    ebpf_get_execution_context_state((ebpf_execution_context_state_t*)state);
+
+Done:
+    if (return_value != EBPF_SUCCESS && provider_reference_held) {
+        ebpf_program_dereference_providers(link->program);
+    }
+
+    if (return_value != EBPF_SUCCESS && epoch_entered) {
+        ebpf_epoch_exit();
+    }
+
+    return return_value;
+}
+
+static ebpf_result_t
+_ebpf_link_instance_invoke_batch_end(_In_ const void* extension_client_binding_context)
+{
+    ebpf_link_t* link = (ebpf_link_t*)ebpf_extension_get_client_context(extension_client_binding_context);
+    ebpf_program_dereference_providers(link->program);
     ebpf_epoch_exit();
+    return EBPF_SUCCESS;
+}
+
+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)
+{
+    // No function entry exit traces as this is a high volume function.
+    ebpf_result_t return_value = EBPF_SUCCESS;
+    ebpf_link_t* link = (ebpf_link_t*)ebpf_extension_get_client_context(extension_client_binding_context);
+
+    ebpf_program_invoke(link->program, program_context, result, (ebpf_execution_context_state_t*)state);
 
-Exit:
     EBPF_RETURN_RESULT(return_value);
 }
 
diff --git a/libs/execution_context/ebpf_program.c b/libs/execution_context/ebpf_program.c
index 85ae7d8bbc..8ba97d48b9 100644
--- a/libs/execution_context/ebpf_program.c
+++ b/libs/execution_context/ebpf_program.c
@@ -1019,25 +1019,36 @@ ebpf_program_set_tail_call(_In_ const ebpf_program_t* next_program)
     return EBPF_SUCCESS;
 }
 
+_Must_inspect_result_ ebpf_result_t
+ebpf_program_reference_providers(_Inout_ ebpf_program_t* program)
+{
+    if (!program->info_extension_client || !ebpf_extension_reference_provider_data(program->info_extension_client)) {
+        return EBPF_EXTENSION_FAILED_TO_LOAD;
+    }
+    return EBPF_SUCCESS;
+}
+
 void
-ebpf_program_invoke(_In_ const ebpf_program_t* program, _Inout_ void* context, _Out_ uint32_t* result)
+ebpf_program_dereference_providers(_Inout_ ebpf_program_t* program)
+{
+    ebpf_extension_dereference_provider_data(program->info_extension_client);
+}
+
+void
+ebpf_program_invoke(
+    _In_ const ebpf_program_t* program,
+    _Inout_ void* context,
+    _Out_ uint32_t* result,
+    _In_ const ebpf_execution_context_state_t* execution_state)
 {
     // High volume call - Skip entry/exit logging.
     ebpf_program_tail_call_state_t state = {0};
     const ebpf_program_t* current_program = program;
 
-    bool provider_data_referenced = false;
     bool program_state_stored = false;
 
-    if (!program->info_extension_client || !ebpf_extension_reference_provider_data(program->info_extension_client)) {
-        *result = 0;
-        return;
-    }
-
-    provider_data_referenced = true;
-
     state.context = context;
-    if (!ebpf_state_store(_ebpf_program_state_index, (uintptr_t)&state) == EBPF_SUCCESS) {
+    if (!ebpf_state_store(_ebpf_program_state_index, (uintptr_t)&state, execution_state) == EBPF_SUCCESS) {
         *result = 0;
         goto Done;
     }
@@ -1080,10 +1091,7 @@ ebpf_program_invoke(_In_ const ebpf_program_t* program, _Inout_ void* context, _
 
 Done:
     if (program_state_stored) {
-        ebpf_assert_success(ebpf_state_store(_ebpf_program_state_index, 0));
-    }
-    if (provider_data_referenced) {
-        ebpf_extension_dereference_provider_data(program->info_extension_client);
+        ebpf_assert_success(ebpf_state_store(_ebpf_program_state_index, 0, execution_state));
     }
 }
 
@@ -1705,15 +1713,30 @@ _ebpf_program_test_run_work_item(_Inout_opt_ void* work_item_context)
     uint64_t cumulative_time = 0;
     ebpf_result_t result;
     uint32_t return_value = 0;
-    uint8_t old_irql;
+    uint8_t old_irql = 0;
     uintptr_t old_thread_affinity;
+    size_t batch_size = options->batch_size ? options->batch_size : 1024;
+    ebpf_execution_context_state_t execution_context_state = {0};
+    bool epoch_entered = false;
+    bool irql_raised = false;
+    bool thread_affinity_set = false;
 
     result = ebpf_set_current_thread_affinity((uintptr_t)1 << options->cpu, &old_thread_affinity);
     if (result != EBPF_SUCCESS) {
         goto Done;
     }
+    thread_affinity_set = true;
 
     old_irql = ebpf_raise_irql(context->required_irql);
+    irql_raised = true;
+
+    result = ebpf_epoch_enter();
+    if (result != EBPF_SUCCESS) {
+        goto Done;
+    }
+    epoch_entered = true;
+
+    ebpf_get_execution_context_state(&execution_context_state);
 
     uint64_t start_time = ebpf_query_time_since_boot(false);
     for (size_t i = 0; i < options->repeat_count; i++) {
@@ -1721,11 +1744,15 @@ _ebpf_program_test_run_work_item(_Inout_opt_ void* work_item_context)
             result = EBPF_CANCELED;
             break;
         }
-        result = ebpf_epoch_enter();
-        if (result != EBPF_SUCCESS)
-            break;
-        ebpf_program_invoke(context->program, context->context, &return_value);
-        ebpf_epoch_exit();
+        // Start a new epoch every batch_size iterations.
+        if ((i % batch_size == (batch_size - 1))) {
+            ebpf_epoch_exit();
+            result = ebpf_epoch_enter();
+            if (result != EBPF_SUCCESS) {
+                break;
+            }
+        }
+        ebpf_program_invoke(context->program, context->context, &return_value, &execution_context_state);
         if (ebpf_should_yield_processor()) {
             // Compute the elapsed time since the last yield.
             end_time = ebpf_query_time_since_boot(false);
@@ -1751,10 +1778,19 @@ _ebpf_program_test_run_work_item(_Inout_opt_ void* work_item_context)
     options->duration /= options->repeat_count;
     options->return_value = return_value;
 
-    ebpf_lower_irql(old_irql);
-    ebpf_restore_current_thread_affinity(old_thread_affinity);
-
 Done:
+    if (epoch_entered) {
+        ebpf_epoch_exit();
+    }
+
+    if (irql_raised) {
+        ebpf_lower_irql(old_irql);
+    }
+
+    if (thread_affinity_set) {
+        ebpf_restore_current_thread_affinity(old_thread_affinity);
+    }
+
     if (context->program_data && context->program_data->context_destroy != NULL && context->context != NULL) {
         context->program_data->context_destroy(
             context->context,
diff --git a/libs/execution_context/ebpf_program.h b/libs/execution_context/ebpf_program.h
index ec513929d6..a22c515971 100644
--- a/libs/execution_context/ebpf_program.h
+++ b/libs/execution_context/ebpf_program.h
@@ -169,11 +169,16 @@ extern "C"
      * @brief Invoke an ebpf_program_t instance.
      *
      * @param[in] program Program to invoke.
-     * @param[in] context Pointer to eBPF context for this program.
+     * @param[in,out] context Pointer to eBPF context for this program.
      * @param[out] result Output from the program.
+     * @param[in] execution_state Execution context state.
      */
     void
-    ebpf_program_invoke(_In_ const ebpf_program_t* program, _Inout_ void* context, _Out_ uint32_t* result);
+    ebpf_program_invoke(
+        _In_ const ebpf_program_t* program,
+        _Inout_ void* context,
+        _Out_ uint32_t* result,
+        _In_ const ebpf_execution_context_state_t* execution_state);
 
     /**
      * @brief Store the helper function IDs that are used by the eBPF program in an array
@@ -341,6 +346,24 @@ extern "C"
         _In_ ebpf_helper_function_addresses_changed_callback_t callback,
         _In_opt_ void* context);
 
+    /**
+     * @brief Acquire a reference to the program information provider.
+     *
+     * @param[in,out] program Program to acquire a reference to the provider for.
+     * @retval EBPF_SUCCESS The operation was successful.
+     * @retval EBPF_EXTENSION_FAILED_TO_LOAD The provider failed to load.
+     */
+    _Must_inspect_result_ ebpf_result_t
+    ebpf_program_reference_providers(_Inout_ ebpf_program_t* program);
+
+    /**
+     * @brief Release a reference to the program information provider.
+     *
+     * @param[in,out] program Program to release a reference to the provider for.
+     */
+    void
+    ebpf_program_dereference_providers(_Inout_ ebpf_program_t* program);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/libs/platform/ebpf_extension.c b/libs/platform/ebpf_extension.c
index c791cf62e6..c95acb63bf 100644
--- a/libs/platform/ebpf_extension.c
+++ b/libs/platform/ebpf_extension.c
@@ -347,7 +347,7 @@ ebpf_extension_unload(_Frees_ptr_opt_ ebpf_extension_client_t* client_context)
 void*
 ebpf_extension_get_client_context(_In_ const void* extension_client_binding_context)
 {
-    EBPF_LOG_ENTRY();
+    // No function entry/exit logging here because this is called with every eBPF program invocation.
     void* local_extension_client_context = NULL;
     ebpf_extension_client_binding_context_t* local_client_binding_context =
         (ebpf_extension_client_binding_context_t*)extension_client_binding_context;
@@ -355,7 +355,7 @@ ebpf_extension_get_client_context(_In_ const void* extension_client_binding_cont
     if (local_client_context != NULL)
         local_extension_client_context = local_client_context->extension_client_context;
 
-    EBPF_RETURN_POINTER(void*, local_extension_client_context);
+    return local_extension_client_context;
 }
 
 GUID
diff --git a/libs/platform/ebpf_platform.h b/libs/platform/ebpf_platform.h
index b1781dd89a..e41ffe6bcc 100644
--- a/libs/platform/ebpf_platform.h
+++ b/libs/platform/ebpf_platform.h
@@ -81,9 +81,9 @@ extern "C"
     typedef ebpf_result_t (*_ebpf_extension_dispatch_function)();
     typedef struct _ebpf_extension_dispatch_table
     {
-        uint16_t version;
-        uint16_t size;
-        _ebpf_extension_dispatch_function function[1];
+        uint16_t version; ///< Version of the dispatch table.
+        uint16_t count;   ///< Number of entries in the dispatch table.
+        _Field_size_(count) _ebpf_extension_dispatch_function function[1];
     } ebpf_extension_dispatch_table_t;
 
     typedef struct _ebpf_extension_data
@@ -100,6 +100,21 @@ extern "C"
         enum bpf_link_type link_type;
     } ebpf_attach_provider_data_t;
 
+    /***
+     * The state of the execution context when the eBPF program was invoked.
+     * This is used to cache state that won't change during the execution of
+     * the eBPF program and is expensive to query.
+     */
+    typedef struct _ebpf_execution_context_state
+    {
+        union
+        {
+            uint64_t thread;
+            uint32_t cpu;
+        } id;
+        uint8_t current_irql;
+    } ebpf_execution_context_state_t;
+
 #define EBPF_ATTACH_CLIENT_DATA_VERSION 0
 #define EBPF_ATTACH_PROVIDER_DATA_VERSION 1
 
@@ -1241,6 +1256,14 @@ extern "C"
     _IRQL_requires_max_(PASSIVE_LEVEL) _Must_inspect_result_ ebpf_result_t
         ebpf_platform_get_authentication_id(_Out_ uint64_t* authentication_id);
 
+    /**
+     * @brief Query the current execution context state.
+     *
+     * @param[out] state The captured execution context state.
+     */
+    void
+    ebpf_get_execution_context_state(_Out_ ebpf_execution_context_state_t* state);
+
 #define EBPF_TRACELOG_EVENT_SUCCESS "EbpfSuccess"
 #define EBPF_TRACELOG_EVENT_RETURN "EbpfReturn"
 #define EBPF_TRACELOG_EVENT_GENERIC_ERROR "EbpfGenericError"
diff --git a/libs/platform/ebpf_state.c b/libs/platform/ebpf_state.c
index 03a7efadb9..b3d18d5ad5 100644
--- a/libs/platform/ebpf_state.c
+++ b/libs/platform/ebpf_state.c
@@ -85,14 +85,15 @@ ebpf_state_allocate_index(_Out_ size_t* new_index)
 }
 
 static _Must_inspect_result_ ebpf_result_t
-_ebpf_state_get_entry(_Outptr_ ebpf_state_entry_t** entry)
+_ebpf_state_get_entry(
+    _In_ const ebpf_execution_context_state_t* execution_context_state, _Outptr_ ebpf_state_entry_t** entry)
 {
     // High frequency call, don't log entry/exit.
     ebpf_state_entry_t* local_entry = NULL;
 
-    if (!ebpf_is_non_preemptible_work_item_supported() || ebpf_is_preemptible()) {
+    if (execution_context_state->current_irql == PASSIVE_LEVEL) {
         ebpf_result_t return_value;
-        uint64_t current_thread_id = ebpf_get_current_thread_id();
+        uint64_t current_thread_id = execution_context_state->id.thread;
 
         return_value =
             ebpf_hash_table_find(_ebpf_state_thread_table, (const uint8_t*)&current_thread_id, (uint8_t**)&local_entry);
@@ -114,7 +115,7 @@ _ebpf_state_get_entry(_Outptr_ ebpf_state_entry_t** entry)
                 _ebpf_state_thread_table, (const uint8_t*)&current_thread_id, (uint8_t**)&local_entry));
         }
     } else {
-        uint32_t current_cpu = ebpf_get_current_cpu();
+        uint32_t current_cpu = execution_context_state->id.cpu;
         if (current_cpu >= _ebpf_state_cpu_table_size) {
             return EBPF_OPERATION_NOT_SUPPORTED;
         }
@@ -125,13 +126,13 @@ _ebpf_state_get_entry(_Outptr_ ebpf_state_entry_t** entry)
 }
 
 _Must_inspect_result_ ebpf_result_t
-ebpf_state_store(size_t index, uintptr_t value)
+ebpf_state_store(size_t index, uintptr_t value, _In_ const ebpf_execution_context_state_t* execution_context_state)
 {
     // High frequency call, don't log entry/exit.
     ebpf_state_entry_t* entry = NULL;
     ebpf_result_t return_value;
 
-    return_value = _ebpf_state_get_entry(&entry);
+    return_value = _ebpf_state_get_entry(execution_context_state, &entry);
     if (return_value == EBPF_SUCCESS) {
         entry->state[index] = value;
     }
@@ -144,8 +145,10 @@ ebpf_state_load(size_t index, _Out_ uintptr_t* value)
     // High frequency call, don't log entry/exit.
     ebpf_state_entry_t* entry = NULL;
     ebpf_result_t return_value;
+    ebpf_execution_context_state_t execution_context_state = {0};
+    ebpf_get_execution_context_state(&execution_context_state);
 
-    return_value = _ebpf_state_get_entry(&entry);
+    return_value = _ebpf_state_get_entry(&execution_context_state, &entry);
     if (return_value == EBPF_SUCCESS) {
         *value = entry->state[index];
     }
diff --git a/libs/platform/ebpf_state.h b/libs/platform/ebpf_state.h
index 736929bd46..68bbeabac2 100644
--- a/libs/platform/ebpf_state.h
+++ b/libs/platform/ebpf_state.h
@@ -43,12 +43,13 @@ extern "C"
      *
      * @param[in] index Assigned for storing state.
      * @param[in] value Value to be stored.
+     * @param[in] execution_context_state Execution context state.
      * @retval EBPF_SUCCESS The operation was successful.
      * @retval EBPF_NO_MEMORY Unable to allocate resources for this
      *  operation.
      */
     _Must_inspect_result_ ebpf_result_t
-    ebpf_state_store(size_t index, uintptr_t value);
+    ebpf_state_store(size_t index, uintptr_t value, _In_ const ebpf_execution_context_state_t* execution_context_state);
 
     /**
      * @brief Load a value in the state tracker.
diff --git a/libs/platform/kernel/ebpf_platform_kernel.c b/libs/platform/kernel/ebpf_platform_kernel.c
index 432c070696..d720367be4 100644
--- a/libs/platform/kernel/ebpf_platform_kernel.c
+++ b/libs/platform/kernel/ebpf_platform_kernel.c
@@ -807,3 +807,14 @@ ebpf_should_yield_processor()
     // KeShouldYieldProcessor returns TRUE if the current thread should yield the processor.
     return KeShouldYieldProcessor() != FALSE;
 }
+
+void
+ebpf_get_execution_context_state(_Out_ ebpf_execution_context_state_t* state)
+{
+    state->current_irql = KeGetCurrentIrql();
+    if (state->current_irql == DISPATCH_LEVEL) {
+        state->id.cpu = ebpf_get_current_cpu();
+    } else {
+        state->id.thread = ebpf_get_current_thread_id();
+    }
+}
\ No newline at end of file
diff --git a/libs/platform/user/ebpf_platform_user.cpp b/libs/platform/user/ebpf_platform_user.cpp
index 45dba21633..1f0ed4cc24 100644
--- a/libs/platform/user/ebpf_platform_user.cpp
+++ b/libs/platform/user/ebpf_platform_user.cpp
@@ -1197,3 +1197,21 @@ ebpf_should_yield_processor()
 {
     return false;
 }
+
+void
+ebpf_get_execution_context_state(_Out_ ebpf_execution_context_state_t* state)
+{
+    if (ebpf_non_preemptible) {
+        state->current_irql = DISPATCH_LEVEL;
+        state->id.cpu = ebpf_get_current_cpu();
+    } else {
+        state->current_irql = PASSIVE_LEVEL;
+        state->id.thread = GetCurrentThreadId();
+    }
+}
+
+uint8_t
+ebpf_get_current_irql()
+{
+    return ebpf_non_preemptible ? DISPATCH_LEVEL : PASSIVE_LEVEL;
+}
\ No newline at end of file