Call KeEnterCriticalRegion during epoch (#2277)

* Call KeEnterCriticalRegion during epoch

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

* PR feedback

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

---------

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

libs/platform/ebpf_epoch.c

@@ -318,6 +318,9 @@ ebpf_epoch_enter()
     current_cpu = ebpf_get_current_cpu();
 
     if (is_preemptible) {
+        // Prevent APCs from running.
+        ebpf_enter_critical_region();
+
         // Block until a thread entry is available.
         ebpf_semaphore_wait(_ebpf_epoch_cpu_table[current_cpu].epoch_table_semaphore);
         // After the semaphore is acquired, then there is at least 1 available epoch entry.
@@ -369,7 +372,11 @@ ebpf_epoch_exit(_In_ ebpf_epoch_state_t* epoch_state)
     }
 
     if (is_preemptible) {
+        // Release the thread entry.
         ebpf_semaphore_release(_ebpf_epoch_cpu_table[current_cpu].epoch_table_semaphore);
+
+        // Allow APCs to run.
+        ebpf_leave_critical_region();
     }
 }
 

libs/platform/ebpf_platform.h

@@ -1313,6 +1313,19 @@ extern "C"
     void
     ebpf_semaphore_release(_In_ ebpf_semaphore_t* semaphore);
 
+    /**
+     * @brief Enter a critical region. This will defer execution of kernel APCs
+     * until ebpf_leave_critical_region is called.
+     */
+    void
+    ebpf_enter_critical_region();
+
+    /**
+     * @brief Leave a critical region. This will resume execution of kernel APCs.
+     */
+    void
+    ebpf_leave_critical_region();
+
 #define EBPF_TRACELOG_EVENT_SUCCESS "EbpfSuccess"
 #define EBPF_TRACELOG_EVENT_RETURN "EbpfReturn"
 #define EBPF_TRACELOG_EVENT_GENERIC_ERROR "EbpfGenericError"

libs/platform/kernel/ebpf_platform_kernel.c

@@ -872,3 +872,15 @@ ebpf_semaphore_release(_In_ ebpf_semaphore_t* semaphore)
 {
     KeReleaseSemaphore(&semaphore->semaphore, 0, 1, FALSE);
 }
+
+void
+ebpf_enter_critical_region()
+{
+    KeEnterCriticalRegion();
+}
+
+void
+ebpf_leave_critical_region()
+{
+    KeLeaveCriticalRegion();
+}

libs/platform/user/ebpf_platform_user.cpp

@@ -1306,3 +1306,15 @@ ebpf_semaphore_release(_In_ ebpf_semaphore_t* semaphore)
 {
     ReleaseSemaphore(semaphore->semaphore, 1, nullptr);
 }
+
+void
+ebpf_enter_critical_region()
+{
+    // This is a no-op for the user mode implementation.
+}
+
+void
+ebpf_leave_critical_region()
+{
+    // This is a no-op for the user mode implementation.
+}