Implement preemptible work items for user mode (#1025)

* fix

* cr comments

libs/platform/user/ebpf_platform_user.cpp

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: MIT
 
 #include "ebpf_platform.h"
+#include "ebpf_utilities.h"
 #include <intsafe.h>
 #include <map>
 #include <mutex>
@@ -20,15 +21,78 @@ bool _ebpf_platform_is_preemptible = true;
 
 extern "C" bool ebpf_fuzzing_enabled = false;
 
+// Thread pool related globals.
+static TP_CALLBACK_ENVIRON _callback_environment;
+static PTP_POOL _pool = nullptr;
+static PTP_CLEANUP_GROUP _cleanup_group = nullptr;
+
+static ebpf_result_t
+_initialize_thread_pool()
+{
+    ebpf_result_t result = EBPF_SUCCESS;
+    bool cleanup_group_created = false;
+    bool return_value;
+
+    InitializeThreadpoolEnvironment(&_callback_environment);
+    _pool = CreateThreadpool(nullptr);
+    if (_pool == nullptr) {
+        result = win32_error_code_to_ebpf_result(GetLastError());
+        goto Exit;
+    }
+
+    SetThreadpoolThreadMaximum(_pool, 1);
+    return_value = SetThreadpoolThreadMinimum(_pool, 1);
+    if (!return_value) {
+        result = win32_error_code_to_ebpf_result(GetLastError());
+        goto Exit;
+    }
+
+    _cleanup_group = CreateThreadpoolCleanupGroup();
+    if (_cleanup_group == nullptr) {
+        result = win32_error_code_to_ebpf_result(GetLastError());
+        goto Exit;
+    }
+    cleanup_group_created = true;
+
+    SetThreadpoolCallbackPool(&_callback_environment, _pool);
+    SetThreadpoolCallbackCleanupGroup(&_callback_environment, _cleanup_group, nullptr);
+
+Exit:
+    if (result != EBPF_SUCCESS) {
+        if (cleanup_group_created) {
+            CloseThreadpoolCleanupGroup(_cleanup_group);
+        }
+        if (_pool) {
+            CloseThreadpool(_pool);
+            _pool = nullptr;
+        }
+    }
+    return result;
+}
+
+static void
+_clean_up_thread_pool()
+{
+    if (!_pool) {
+        return;
+    }
+
+    CloseThreadpoolCleanupGroupMembers(_cleanup_group, false, nullptr);
+    CloseThreadpoolCleanupGroup(_cleanup_group);
+    CloseThreadpool(_pool);
+}
+
 ebpf_result_t
 ebpf_platform_initiate()
 {
-    return EBPF_SUCCESS;
+    return _initialize_thread_pool();
 }
 
 void
 ebpf_platform_terminate()
-{}
+{
+    _clean_up_thread_pool();
+}
 
 ebpf_result_t
 ebpf_get_code_integrity_state(_Out_ ebpf_code_integrity_state_t* state)
@@ -364,7 +428,9 @@ ebpf_lock_destroy(_In_ ebpf_lock_t* lock)
 }
 
 _Requires_lock_not_held_(*lock) _Acquires_lock_(*lock) _IRQL_requires_max_(DISPATCH_LEVEL) _IRQL_saves_
-    _IRQL_raises_(DISPATCH_LEVEL) ebpf_lock_state_t ebpf_lock_lock(_In_ ebpf_lock_t* lock)
+    _IRQL_raises_(DISPATCH_LEVEL)
+ebpf_lock_state_t
+ebpf_lock_lock(_In_ ebpf_lock_t* lock)
 {
     AcquireSRWLockExclusive(reinterpret_cast<PSRWLOCK>(lock));
     return 0;
@@ -446,9 +512,7 @@ ebpf_is_non_preemptible_work_item_supported()
 bool
 ebpf_is_preemptible_work_item_supported()
 {
-    // TODO: (Issue# 854) Add support for creating threads to queue preemptible work items in user mode.
+    return true;
-    // https://github.com/microsoft/ebpf-for-windows/issues/854
-    return false;
 }
 
 uint32_t
@@ -491,12 +555,34 @@ ebpf_queue_non_preemptible_work_item(_In_ ebpf_non_preemptible_work_item_t* work
     return false;
 }
 
+typedef struct _ebpf_preemptible_work_item
+{
+    PTP_WORK work;
+    void (*work_item_routine)(_In_opt_ const void* work_item_context);
+    void* work_item_context;
+} ebpf_preemptible_work_item_t;
+
+static void
+_ebpf_preemptible_routine(_Inout_ PTP_CALLBACK_INSTANCE instance, _In_opt_ PVOID parameter, _Inout_ PTP_WORK work)
+{
+    UNREFERENCED_PARAMETER(instance);
+    UNREFERENCED_PARAMETER(work);
+
+    if (parameter == nullptr) {
+        return;
+    }
+
+    ebpf_preemptible_work_item_t* work_item = (ebpf_preemptible_work_item_t*)parameter;
+    work_item->work_item_routine(work_item->work_item_context);
+
+    ebpf_free(work_item->work_item_context);
+    ebpf_free(work_item);
+}
+
 void
 ebpf_queue_preemptible_work_item(_In_ ebpf_preemptible_work_item_t* work_item)
 {
-    UNREFERENCED_PARAMETER(work_item);
+    SubmitThreadpoolWork(work_item->work);
-    // TODO: (Issue# 854) Add support for creating threads to queue preemptible work items in user mode.
-    // https://github.com/microsoft/ebpf-for-windows/issues/854
 }
 
 ebpf_result_t
@@ -505,12 +591,26 @@ ebpf_allocate_preemptible_work_item(
     _In_ void (*work_item_routine)(_In_opt_ const void* work_item_context),
     _In_opt_ void* work_item_context)
 {
-    UNREFERENCED_PARAMETER(work_item);
+    ebpf_result_t result = EBPF_SUCCESS;
-    UNREFERENCED_PARAMETER(work_item_routine);
+    *work_item = (ebpf_preemptible_work_item_t*)ebpf_allocate(sizeof(ebpf_preemptible_work_item_t));
-    UNREFERENCED_PARAMETER(work_item_context);
+    if (*work_item == nullptr) {
-    // TODO: (Issue# 854) Add support for creating threads to queue preemptible work items in user mode.
+        return EBPF_NO_MEMORY;
-    // https://github.com/microsoft/ebpf-for-windows/issues/854
+    }
-    return EBPF_OPERATION_NOT_SUPPORTED;
+
+    (*work_item)->work = CreateThreadpoolWork(_ebpf_preemptible_routine, *work_item, &_callback_environment);
+    if ((*work_item)->work == nullptr) {
+        result = win32_error_code_to_ebpf_result(GetLastError());
+        goto Done;
+    }
+    (*work_item)->work_item_routine = work_item_routine;
+    (*work_item)->work_item_context = work_item_context;
+
+Done:
+    if (result != EBPF_SUCCESS) {
+        ebpf_free(*work_item);
+        *work_item = nullptr;
+    }
+    return result;
 }
 
 typedef struct _ebpf_timer_work_item
@@ -538,11 +638,11 @@ ebpf_allocate_timer_work_item(
 {
     *work_item = (ebpf_timer_work_item_t*)ebpf_allocate(sizeof(ebpf_timer_work_item_t));
 
-    if (*work_item == NULL)
+    if (*work_item == nullptr)
         goto Error;
 
-    (*work_item)->threadpool_timer = CreateThreadpoolTimer(_ebpf_timer_callback, *work_item, NULL);
+    (*work_item)->threadpool_timer = CreateThreadpoolTimer(_ebpf_timer_callback, *work_item, nullptr);
-    if ((*work_item)->threadpool_timer == NULL)
+    if ((*work_item)->threadpool_timer == nullptr)
         goto Error;
 
     (*work_item)->work_item_routine = work_item_routine;
@@ -551,8 +651,8 @@ ebpf_allocate_timer_work_item(
     return EBPF_SUCCESS;
 
 Error:
-    if (*work_item != NULL) {
+    if (*work_item != nullptr) {
-        if ((*work_item)->threadpool_timer != NULL)
+        if ((*work_item)->threadpool_timer != nullptr)
             CloseThreadpoolTimer((*work_item)->threadpool_timer);
 
         ebpf_free(*work_item);
@@ -697,7 +797,7 @@ Done:
 uint32_t
 ebpf_result_to_win32_error_code(ebpf_result_t result)
 {
-    static uint32_t (*RtlNtStatusToDosError)(NTSTATUS Status) = NULL;
+    static uint32_t (*RtlNtStatusToDosError)(NTSTATUS Status) = nullptr;
     if (!RtlNtStatusToDosError) {
         HMODULE ntdll = LoadLibrary(L"ntdll.dll");
         if (!ntdll) {