Zero initialize batch context state (#3058)

* Zero initialize batch context state

Signed-off-by: Alan Jowett <alan.jowett@microsoft.com>

* Document batch semantics

Signed-off-by: Alan Jowett <alan.jowett@microsoft.com>

* Fix grammar

Signed-off-by: Alan Jowett <alan.jowett@microsoft.com>

* Update ebpf_extension.h

* Update eBpfExtensions.md

* Update docs/eBpfExtensions.md

Co-authored-by: Dave Thaler <dthaler1968@gmail.com>

* Update docs/eBpfExtensions.md

Co-authored-by: Anurag Saxena <43585259+saxena-anurag@users.noreply.github.com>

* Update docs/eBpfExtensions.md

Co-authored-by: Anurag Saxena <43585259+saxena-anurag@users.noreply.github.com>

* Update docs/eBpfExtensions.md

Co-authored-by: Dave Thaler <dthaler1968@gmail.com>

* Update include/ebpf_extension.h

Co-authored-by: Dave Thaler <dthaler1968@gmail.com>

* Update include/ebpf_extension.h

Co-authored-by: Dave Thaler <dthaler1968@gmail.com>

* PR feedback

Signed-off-by: Alan Jowett <alan.jowett@microsoft.com>

* PR feedback

Signed-off-by: Alan Jowett <alan.jowett@microsoft.com>

---------

Signed-off-by: Alan Jowett <alan.jowett@microsoft.com>
Co-authored-by: Alan Jowett <alan.jowett@microsoft.com>
Co-authored-by: Dave Thaler <dthaler1968@gmail.com>
Co-authored-by: Anurag Saxena <43585259+saxena-anurag@users.noreply.github.com>

docs/eBpfExtensions.md

@@ -248,16 +248,66 @@ the provider must free the per-client context passed in via `ProviderBindingCont
 
 ### 2.5 Invoking an eBPF program from Hook NPI Provider
 To invoke an eBPF program, the extension uses the dispatch table supplied by the Hook NPI client during attaching.
-There is only one function in the client dispatch table, which is of the following type:
+The client dispatch table contains the functions, with the following type prototypes:
 
 ```
 /**
- *  @brief This is the only mandatory function in the eBPF Hook NPI client dispatch table.
+ * @brief Invoke the eBPF program.
+ *
+ * @param[in] extension_client_binding_context The context provided by the extension client when the binding was created.
+ * @param[in,out] program_context The context for this invocation of the eBPF program.
+ * @param[out] result The result of the eBPF program.
+ *
+ * @retval EBPF_SUCCESS if successful or an appropriate error code.
+ * @retval EBPF_NO_MEMORY if memory allocation fails.
+ * @retval EBPF_EXTENSION_FAILED_TO_LOAD if required extension is not loaded.
  */
 typedef ebpf_result_t (*ebpf_program_invoke_function_t)(
-    _In_ const void* client_binding_context, _In_ const void* context, _Out_ uint32_t* result);
+    _In_ const void* extension_client_binding_context, _Inout_ void* program_context, _Out_ uint32_t* result);
 
+/**
+ * @brief Prepare the eBPF program for batch invocation.
+ *
+ * @param[in] extension_client_binding_context The context provided by the extension client when the binding was created.
+ * @param[in] state_size The size of the state to be allocated, which should be greater than or equal to
+ * sizeof(ebpf_execution_context_state_t).
+ * @param[out] state The state to be used for batch invocation.
+ *
+ * @retval EBPF_SUCCESS if successful or an appropriate error code.
+ * @retval EBPF_NO_MEMORY if memory allocation fails.
+ * @retval EBPF_EXTENSION_FAILED_TO_LOAD if required extension is not loaded.
+ */
+typedef ebpf_result_t (*ebpf_program_batch_begin_invoke_function_t)(
+    _In_ const void* extension_client_binding_context, size_t state_size, _Out_writes_(state_size) void* state);
+
+/**
+ * @brief Invoke the eBPF program in batch mode.
+ *
+ * @param[in] extension_client_binding_context The context provided by the extension client when the binding was created.
+ * @param[in,out] program_context The context for this invocation of the eBPF program.
+ * @param[out] result The result of the eBPF program.
+ * @param[in] state The state to be used for batch invocation.
+ *
+ * @retval EBPF_SUCCESS.
+ */
+typedef ebpf_result_t (*ebpf_program_batch_invoke_function_t)(
+    _In_ const void* extension_client_binding_context,
+    _Inout_ void* program_context,
+    _Out_ uint32_t* result,
+    _In_ const void* state);
+
+/**
+ * @brief Clean up the eBPF program after batch invocation.
+ *
+ * @param[in] extension_client_binding_context The context provided by the extension client when the binding was created.
+ * @param[in,out] state The state to be used for batch invocation.
+ *
+ * @retval EBPF_SUCCESS.
+ */
+typedef ebpf_result_t (*ebpf_program_batch_end_invoke_function_t)(
+    _In_ const void* extension_client_binding_context, _Inout_ void* state);
 ```
+
 The function pointer can be obtained from the client dispatch table as follows:
 ```
 invoke_program = (ebpf_program_invoke_function_t)client_dispatch_table->function[0];
@@ -270,6 +320,16 @@ the context descriptor (using the `ebpf_context_descriptor_t` type) to the eBPF
 client hosted by the Execution Context. The `result` output parameter holds the return value from the eBPF program
 post execution.
 
+In cases where the same eBPF program will be invoked sequentially with different context data (aka batch invocation),
+the caller can reduce the overhead by using the batch invocation APIs. Prior to the first invocation, the batch
+begin API is called, which caches state used by the eBPF program and prevents the program from being unloaded. The
+caller is responsible for providing storage large enough to store an instance of ebpf_execution_context_state_t and
+ensuring that it remain valid until calling the batch end API. Between the begin and end calls, the caller may call
+the batch invoke API multiple times to invoke the BPF program with minimal overhead. Callers must limit the length
+of time a batch is open and must not change IRQL between calling batch begin and end. Batch end cost may scale with
+the number of times the program has been invoked, so callers should limit the number of calls within a batch to
+prevent long delays in batch end.
+
 ### 2.6 Authoring Helper Functions
 An extension can provide an implementation of helper functions that can be invoked by the eBPF programs. The helper
 functions can be of two types:

include/ebpf_extension.h

@@ -15,18 +15,63 @@ typedef struct _ebpf_extension_dispatch_table
     _Field_size_(count) _ebpf_extension_dispatch_function function[1];
 } ebpf_extension_dispatch_table_t;
 
+/**
+ * @brief Invoke the eBPF program.
+ *
+ * @param[in] extension_client_binding_context The context provided by the extension client when the binding was
+ * created.
+ * @param[in,out] program_context The context for this invocation of the eBPF program.
+ * @param[out] result The result of the eBPF program.
+ *
+ * @retval EBPF_SUCCESS The operation was successful.
+ * @retval EBPF_NO_MEMORY The operation failed due to lack of memory.
+ * @retval EBPF_EXTENSION_FAILED_TO_LOAD The required extension is not loaded.
+ */
 typedef ebpf_result_t (*ebpf_program_invoke_function_t)(
     _In_ const void* extension_client_binding_context, _Inout_ void* program_context, _Out_ uint32_t* result);
 
+/**
+ * @brief Prepare the eBPF program for batch invocation.
+ *
+ * @param[in] extension_client_binding_context The context provided by the extension client when the binding was
+ * created.
+ * @param[in] state_size The size of the state to be allocated, which should be greater than or equal to
+ * sizeof(ebpf_execution_context_state_t).
+ * @param[out] state The state to be used for batch invocation.
+ *
+ * @retval EBPF_SUCCESS The operation was successful.
+ * @retval EBPF_NO_MEMORY The operation failed due to lack of memory.
+ * @retval EBPF_EXTENSION_FAILED_TO_LOAD The required extension is not loaded.
+ */
 typedef ebpf_result_t (*ebpf_program_batch_begin_invoke_function_t)(
     _In_ const void* extension_client_binding_context, size_t state_size, _Out_writes_(state_size) void* state);
 
+/**
+ * @brief Invoke the eBPF program in batch mode.
+ *
+ * @param[in] extension_client_binding_context The context provided by the extension client when the binding was
+ * created.
+ * @param[in,out] program_context The context for this invocation of the eBPF program.
+ * @param[out] result The result of the eBPF program.
+ * @param[in] state The state to be used for batch invocation.
+ *
+ * @retval EBPF_SUCCESS The operation was successful.
+ */
 typedef ebpf_result_t (*ebpf_program_batch_invoke_function_t)(
     _In_ const void* extension_client_binding_context,
     _Inout_ void* program_context,
     _Out_ uint32_t* result,
     _In_ const void* state);
 
+/**
+ * @brief Clean up the eBPF program after batch invocation.
+ *
+ * @param[in] extension_client_binding_context The context provided by the extension client when the binding was
+ * created.
+ * @param[in,out] state The state to be used for batch invocation.
+ *
+ * @retval EBPF_SUCCESS The operation was successful.
+ */
 typedef ebpf_result_t (*ebpf_program_batch_end_invoke_function_t)(
     _In_ const void* extension_client_binding_context, _Inout_ void* state);
 

libs/execution_context/ebpf_link.c

@@ -485,6 +485,8 @@ _ebpf_link_instance_invoke_batch_begin(
         goto Done;
     }
 
+    memset(execution_context_state, 0, sizeof(ebpf_execution_context_state_t));
+
     ebpf_get_execution_context_state(execution_context_state);
     return_value = ebpf_state_store(ebpf_program_get_state_index(), (uintptr_t)state, execution_context_state);
     if (return_value != EBPF_SUCCESS) {