[PATCH 3/5] ntoskrnl.exe: Add KeGenericCallDpc() function.

[Paul Gofman]

Signed-off-by: Paul Gofman <pgofman at codeweavers.com>
---
 dlls/ntoskrnl.exe/ntoskrnl.c        | 113 ++++++++++++++++++++++++++++
 dlls/ntoskrnl.exe/ntoskrnl.exe.spec |   1 +
 include/ddk/wdm.h                   |   7 ++
 3 files changed, 121 insertions(+)

diff --git a/dlls/ntoskrnl.exe/ntoskrnl.c b/dlls/ntoskrnl.exe/ntoskrnl.c
index 2291842a67..1a0cf1c5c2 100644
--- a/dlls/ntoskrnl.exe/ntoskrnl.c
+++ b/dlls/ntoskrnl.exe/ntoskrnl.c
@@ -72,6 +72,10 @@ KSERVICE_TABLE_DESCRIPTOR KeServiceDescriptorTable[4] = { { 0 } };
 
 #define MAX_SERVICE_NAME 260
 
+static TP_POOL *dpc_call_tp;
+static TP_CALLBACK_ENVIRON dpc_call_tpe;
+DECLARE_CRITICAL_SECTION(dpc_call_cs);
+
 /* tid of the thread running client request */
 static DWORD request_thread;
 
@@ -3160,6 +3164,10 @@ BOOL WINAPI DllMain( HINSTANCE inst, DWORD reason, LPVOID reserved )
         break;
     case DLL_PROCESS_DETACH:
         if (reserved) break;
+
+        if (dpc_call_tp)
+            CloseThreadpool(dpc_call_tp);
+
         HeapDestroy( ntoskrnl_heap );
         RtlRemoveVectoredExceptionHandler( handler );
         break;
@@ -4012,6 +4020,111 @@ BOOLEAN WINAPI KdRefreshDebuggerNotPresent(void)
     return !KdDebuggerEnabled;
 }
 
+struct generic_call_dpc_context
+{
+    DEFERRED_REVERSE_BARRIER *reverse_barrier;
+    PKDEFERRED_ROUTINE routine;
+    ULONG *cpu_count_barrier;
+    void *context;
+    ULONG cpu_index;
+};
+
+static void WINAPI generic_call_dpc_callback(TP_CALLBACK_INSTANCE *instance, void *context)
+{
+    struct generic_call_dpc_context *c = context;
+    GROUP_AFFINITY old, new;
+
+    TRACE("instance %p, context %p.\n", instance, context);
+
+    NtQueryInformationThread(GetCurrentThread(), ThreadGroupInformation,
+            &old, sizeof(old), NULL);
+
+    memset(&new, 0, sizeof(new));
+
+    new.Mask = 1 << c->cpu_index;
+    NtSetInformationThread(GetCurrentThread(), ThreadGroupInformation, &new, sizeof(new));
+
+    c->routine((PKDPC)0xdeadbeef, c->context, c->cpu_count_barrier, c->reverse_barrier);
+    NtSetInformationThread(GetCurrentThread(), ThreadGroupInformation, &old, sizeof(old));
+}
+
+void WINAPI KeGenericCallDpc(PKDEFERRED_ROUTINE routine, void *context)
+{
+    ULONG cpu_count = KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS);
+    static struct generic_call_dpc_context *contexts;
+    DEFERRED_REVERSE_BARRIER reverse_barrier;
+    static ULONG last_cpu_count;
+    ULONG cpu_count_barrier;
+    ULONG i;
+
+    TRACE("routine %p, context %p.\n", routine, context);
+
+    EnterCriticalSection(&dpc_call_cs);
+
+    if (!dpc_call_tp)
+    {
+        if (!(dpc_call_tp = CreateThreadpool(NULL)))
+        {
+            ERR("Could not create thread pool.\n");
+            LeaveCriticalSection(&dpc_call_cs);
+            return;
+        }
+
+        SetThreadpoolThreadMinimum(dpc_call_tp, cpu_count);
+        SetThreadpoolThreadMaximum(dpc_call_tp, cpu_count);
+
+        memset(&dpc_call_tpe, 0, sizeof(dpc_call_tpe));
+        dpc_call_tpe.Version = 1;
+        dpc_call_tpe.Pool = dpc_call_tp;
+    }
+
+    reverse_barrier.Barrier = cpu_count;
+    reverse_barrier.TotalProcessors = cpu_count;
+    cpu_count_barrier = cpu_count;
+
+    if (contexts)
+    {
+        if (last_cpu_count < cpu_count)
+        {
+            static struct generic_call_dpc_context *new_contexts;
+            if (!(new_contexts = heap_realloc(contexts, sizeof(*contexts) * cpu_count)))
+            {
+                ERR("No memory.\n");
+                LeaveCriticalSection(&dpc_call_cs);
+                return;
+            }
+            contexts = new_contexts;
+            SetThreadpoolThreadMinimum(dpc_call_tp, cpu_count);
+            SetThreadpoolThreadMaximum(dpc_call_tp, cpu_count);
+        }
+    }
+    else if (!(contexts = heap_alloc(sizeof(*contexts) * cpu_count)))
+    {
+        ERR("No memory.\n");
+        LeaveCriticalSection(&dpc_call_cs);
+        return;
+    }
+
+    memset(contexts, 0, sizeof(*contexts) * cpu_count);
+    last_cpu_count = cpu_count;
+
+    for (i = 0; i < cpu_count; ++i)
+    {
+        contexts[i].reverse_barrier = &reverse_barrier;
+        contexts[i].cpu_count_barrier = &cpu_count_barrier;
+        contexts[i].routine = routine;
+        contexts[i].context = context;
+        contexts[i].cpu_index = i;
+
+        TrySubmitThreadpoolCallback(generic_call_dpc_callback, &contexts[i], &dpc_call_tpe);
+    }
+
+    while (InterlockedCompareExchange((LONG *)&cpu_count_barrier, 0, 0))
+        SwitchToThread();
+
+    LeaveCriticalSection(&dpc_call_cs);
+}
+
 void WINAPI KeSignalCallDpcDone(void *barrier)
 {
     InterlockedDecrement((LONG *)barrier);
diff --git a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
index 28f8d5e4cd..923ba22529 100644
--- a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
+++ b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
@@ -59,6 +59,7 @@
 @ stub KeSetTimeUpdateNotifyRoutine
 @ stub KefAcquireSpinLockAtDpcLevel
 @ stub KefReleaseSpinLockFromDpcLevel
+@ stdcall KeGenericCallDpc(ptr ptr)
 @ stdcall KeSignalCallDpcDone(ptr)
 @ stub KiAcquireSpinLock
 @ stub KiReleaseSpinLock
diff --git a/include/ddk/wdm.h b/include/ddk/wdm.h
index b11f4d239d..1b001d7424 100644
--- a/include/ddk/wdm.h
+++ b/include/ddk/wdm.h
@@ -122,6 +122,12 @@ typedef struct _KMUTANT {
     UCHAR ApcDisable;
 } KMUTANT, *PKMUTANT, *RESTRICTED_POINTER PRKMUTANT, KMUTEX, *PKMUTEX, *RESTRICTED_POINTER PRKMUTEX;
 
+typedef struct _DEFERRED_REVERSE_BARRIER
+{
+    ULONG Barrier;
+    ULONG TotalProcessors;
+} DEFERRED_REVERSE_BARRIER;
+
 typedef enum _KWAIT_REASON
 {
     Executive,
@@ -1686,6 +1692,7 @@ BOOLEAN   WINAPI KeCancelTimer(KTIMER*);
 void      WINAPI KeClearEvent(PRKEVENT);
 NTSTATUS  WINAPI KeDelayExecutionThread(KPROCESSOR_MODE,BOOLEAN,LARGE_INTEGER*);
 void      WINAPI KeEnterCriticalRegion(void);
+void      WINAPI KeGenericCallDpc(PKDEFERRED_ROUTINE,PVOID);
 ULONG     WINAPI KeGetCurrentProcessorNumber(void);
 PKTHREAD  WINAPI KeGetCurrentThread(void);
 void      WINAPI KeInitializeEvent(PRKEVENT,EVENT_TYPE,BOOLEAN);
-- 
2.26.2