Random crash fix - preemptible workitems (#1042)

libs/execution_context/ebpf_native.c

@@ -17,9 +17,6 @@ static const GUID GUID_NULL = {0, 0, 0, {0, 0, 0, 0, 0, 0, 0, 0}};
 
 typedef uint64_t (*helper_function_address)(uint64_t r1, uint64_t r2, uint64_t r3, uint64_t r4, uint64_t r5);
 
-static void
-_ebpf_native_unload_workitem(_In_opt_ const void* module_id);
-
 typedef struct _ebpf_native_map
 {
     map_entry_t* entry;
@@ -82,6 +79,15 @@ ebpf_native_load_driver(_In_z_ const wchar_t* service_name);
 void
 ebpf_native_unload_driver(_In_z_ const wchar_t* service_name);
 
+static void
+_ebpf_native_unload_work_item(_In_opt_ const void* service)
+{
+    // Do not free "service" here. It is freed by platform.
+    if (service != NULL) {
+        ebpf_native_unload_driver((const wchar_t*)service);
+    }
+}
+
 static inline bool
 _ebpf_native_is_map_in_map(_In_ const ebpf_native_map_t* map)
 {
@@ -170,12 +176,11 @@ ebpf_native_release_reference(_In_opt_ ebpf_native_module_t* module)
     if (new_ref_count == 1) {
         // Check if all the program references have been released. If that
         // is the case, explicitly unload the driver, if it is safe to do so.
-        if (ebpf_is_preemptible_work_item_supported()) {
         bool unload = false;
         module_lock_state = ebpf_lock_lock(&module->lock);
         lock_acquired = true;
-            if (!module->detaching && !module->unloading) {
+        if (!module->detaching) {
-                // If the module is not yet unloading or detaching, and reference
+            // If the module is not yet marked as detaching, and reference
             // count is 1, it means all the program references have been
             // released.
             EBPF_LOG_MESSAGE_GUID(
@@ -194,13 +199,8 @@ ebpf_native_release_reference(_In_opt_ ebpf_native_module_t* module)
         ebpf_lock_unlock(&module->lock, module_lock_state);
         lock_acquired = false;
         if (unload) {
-                ebpf_preemptible_work_item_t* work_item = NULL;
+            ebpf_native_unload(module_id);
-                result = ebpf_allocate_preemptible_work_item(&work_item, _ebpf_native_unload_workitem, module_id);
+            ebpf_free(module_id);
-                if (result != EBPF_SUCCESS) {
-                    goto Done;
-                }
-                ebpf_queue_preemptible_work_item(work_item);
-            }
         }
     } else if (new_ref_count == 0) {
         ebpf_lock_state_t state = ebpf_lock_lock(&_ebpf_native_client_table_lock);
@@ -1306,7 +1306,8 @@ ebpf_native_unload(_In_ const GUID* module_id)
     ebpf_native_module_t* module = NULL;
     wchar_t* service_name = NULL;
     size_t service_name_length;
-    bool unload_module = false;
+    bool queue_work_item = false;
+    ebpf_preemptible_work_item_t* work_item = NULL;
 
     // Find the native entry in hash table.
     state = ebpf_lock_lock(&_ebpf_native_client_table_lock);
@@ -1334,8 +1335,6 @@ ebpf_native_unload(_In_ const GUID* module_id)
         result = EBPF_SUCCESS;
         goto Done;
     }
-    module->unloading = true;
-    unload_module = true;
 
     // It is possible that the module is also detaching at the same time and
     // the module memory can be freed immediately after the hash table lock is
@@ -1349,13 +1348,26 @@ ebpf_native_unload(_In_ const GUID* module_id)
 
     memcpy(service_name, module->service_name, service_name_length);
 
+    // Create a work item if we are running at DISPATCH.
+    if (!ebpf_is_preemptible()) {
+        result = ebpf_allocate_preemptible_work_item(&work_item, _ebpf_native_unload_work_item, service_name);
+        if (result != EBPF_SUCCESS) {
+            goto Done;
+        }
+        queue_work_item = true;
+    }
+    module->unloading = true;
+
     ebpf_lock_unlock(&module->lock, module_state);
     module_lock_acquired = false;
     ebpf_lock_unlock(&_ebpf_native_client_table_lock, state);
     lock_acquired = false;
 
-    if (unload_module)
+    if (queue_work_item) {
+        ebpf_queue_preemptible_work_item(work_item);
+    } else {
         ebpf_native_unload_driver(service_name);
+    }
 
 Done:
     if (module_lock_acquired) {
@@ -1367,15 +1379,9 @@ Done:
         lock_acquired = false;
     }
 
+    if (!queue_work_item) {
         ebpf_free(service_name);
+    }
 
     EBPF_RETURN_RESULT(result);
 }
-
-static void
-_ebpf_native_unload_workitem(_In_opt_ const void* module_id)
-{
-    if (module_id != NULL) {
-        ebpf_native_unload((GUID*)module_id);
-    }
-}

libs/platform/ebpf_platform.h

@@ -332,9 +332,7 @@ extern "C"
      * @returns - The previous lock_state required for unlock.
      */
     _Requires_lock_not_held_(*lock) _Acquires_lock_(*lock) _IRQL_requires_max_(DISPATCH_LEVEL) _IRQL_saves_
-        _IRQL_raises_(DISPATCH_LEVEL)
+        _IRQL_raises_(DISPATCH_LEVEL) ebpf_lock_state_t ebpf_lock_lock(_In_ ebpf_lock_t* lock);
-    ebpf_lock_state_t
-    ebpf_lock_lock(_In_ ebpf_lock_t* lock);
 
     /**
      * @brief Release exclusive access to the lock.
@@ -422,16 +420,6 @@ extern "C"
     bool
     ebpf_queue_non_preemptible_work_item(_In_ ebpf_non_preemptible_work_item_t* work_item, _In_opt_ void* parameter_1);
 
-    /**
-     * @brief Query the platform to determine if preemptible work items are
-     *   supported.
-     *
-     * @retval true Preemptible work items are supported.
-     * @retval false Preemptible work items are not supported.
-     */
-    bool
-    ebpf_is_preemptible_work_item_supported();
-
     /**
      * @brief Create a preemptible work item.
      *

libs/platform/kernel/ebpf_platform_kernel.c

@@ -413,12 +413,6 @@ ebpf_is_non_preemptible_work_item_supported()
     return true;
 }
 
-bool
-ebpf_is_preemptible_work_item_supported()
-{
-    return true;
-}
-
 uint32_t
 ebpf_get_current_cpu()
 {

libs/platform/user/ebpf_platform_user.cpp

@@ -507,12 +507,6 @@ ebpf_is_non_preemptible_work_item_supported()
     return false;
 }
 
-bool
-ebpf_is_preemptible_work_item_supported()
-{
-    return false;
-}
-
 uint32_t
 ebpf_get_current_cpu()
 {