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

Paul Gofman pgofman at codeweavers.com
Fri May 22 09:24:45 CDT 2020


Signed-off-by: Paul Gofman <pgofman at codeweavers.com>
---
 dlls/ntoskrnl.exe/ntoskrnl.c        | 12 +++++++++++
 dlls/ntoskrnl.exe/ntoskrnl.exe.spec |  1 +
 dlls/ntoskrnl.exe/tests/driver.c    | 31 +++++++++++++++++++++++++++++
 include/ddk/wdm.h                   |  1 +
 include/winnt.h                     |  2 ++
 5 files changed, 47 insertions(+)

diff --git a/dlls/ntoskrnl.exe/ntoskrnl.c b/dlls/ntoskrnl.exe/ntoskrnl.c
index aa71b7c48c..830608ce50 100644
--- a/dlls/ntoskrnl.exe/ntoskrnl.c
+++ b/dlls/ntoskrnl.exe/ntoskrnl.c
@@ -2385,6 +2385,18 @@ KAFFINITY WINAPI KeQueryActiveProcessors( void )
     return AffinityMask;
 }
 
+ULONG WINAPI KeQueryActiveProcessorCountEx(USHORT group_number)
+{
+    TRACE("group_number %u.\n", group_number);
+
+    if (group_number && group_number != ALL_PROCESSOR_GROUPS)
+    {
+        FIXME("group_number %u not supported.\n", group_number);
+        return 0;
+    }
+
+    return NtCurrentTeb()->Peb->NumberOfProcessors;
+}
 
 /**********************************************************************
  *           KeQueryInterruptTime   (NTOSKRNL.EXE.@)
diff --git a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
index 0490dcd691..f25ab9c5e0 100644
--- a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
+++ b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
@@ -587,6 +587,7 @@
 @ stub KeProfileInterruptWithSource
 @ stub KePulseEvent
 @ stdcall KeQueryActiveProcessors()
+@ stdcall KeQueryActiveProcessorCountEx(long)
 @ stdcall KeQueryInterruptTime()
 @ stub KeQueryPriorityThread
 @ stub KeQueryRuntimeThread
diff --git a/dlls/ntoskrnl.exe/tests/driver.c b/dlls/ntoskrnl.exe/tests/driver.c
index 10f3865515..6304a5012f 100644
--- a/dlls/ntoskrnl.exe/tests/driver.c
+++ b/dlls/ntoskrnl.exe/tests/driver.c
@@ -1710,6 +1710,36 @@ static void test_executable_pool(void)
 }
 #endif
 
+static void test_affinity(void)
+{
+    ULONG (WINAPI *pKeQueryActiveProcessorCountEx)(USHORT);
+    KAFFINITY (WINAPI *pKeQueryActiveProcessors)(void);
+    ULONG cpu_count, count;
+    KAFFINITY mask;
+
+    pKeQueryActiveProcessorCountEx = get_proc_address("KeQueryActiveProcessorCountEx");
+    if (!pKeQueryActiveProcessorCountEx)
+    {
+        win_skip("KeQueryActiveProcessorCountEx is not available.\n");
+        return;
+    }
+
+    pKeQueryActiveProcessors = get_proc_address("KeQueryActiveProcessors");
+    ok(!!pKeQueryActiveProcessors, "KeQueryActiveProcessors is not available.\n");
+
+    count = pKeQueryActiveProcessorCountEx(1);
+    ok(!count, "Got unexpected count %u.\n", count);
+
+    cpu_count = pKeQueryActiveProcessorCountEx(0);
+    ok(cpu_count, "Got unexpected cpu_count %u.\n", cpu_count);
+
+    count = pKeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS);
+    ok(count == cpu_count, "Got unexpected count %u.\n", count);
+
+    mask = pKeQueryActiveProcessors();
+    ok(mask == ~((~0u) << cpu_count), "Got unexpected mask %#lx.\n", mask);
+}
+
 static NTSTATUS main_test(DEVICE_OBJECT *device, IRP *irp, IO_STACK_LOCATION *stack)
 {
     ULONG length = stack->Parameters.DeviceIoControl.OutputBufferLength;
@@ -1763,6 +1793,7 @@ static NTSTATUS main_test(DEVICE_OBJECT *device, IRP *irp, IO_STACK_LOCATION *st
 #if defined(__i386__) || defined(__x86_64__)
     test_executable_pool();
 #endif
+    test_affinity();
 
     if (main_test_work_item) return STATUS_UNEXPECTED_IO_ERROR;
 
diff --git a/include/ddk/wdm.h b/include/ddk/wdm.h
index 5b1f19e4f1..790d1a1247 100644
--- a/include/ddk/wdm.h
+++ b/include/ddk/wdm.h
@@ -1696,6 +1696,7 @@ void      WINAPI KeInitializeSpinLock(KSPIN_LOCK*);
 void      WINAPI KeInitializeTimerEx(PKTIMER,TIMER_TYPE);
 void      WINAPI KeInitializeTimer(KTIMER*);
 void      WINAPI KeLeaveCriticalRegion(void);
+ULONG     WINAPI KeQueryActiveProcessorCountEx(USHORT);
 KAFFINITY WINAPI KeQueryActiveProcessors(void);
 void      WINAPI KeQuerySystemTime(LARGE_INTEGER*);
 void      WINAPI KeQueryTickCount(LARGE_INTEGER*);
diff --git a/include/winnt.h b/include/winnt.h
index 1f0a7e344e..3e3ac338bc 100644
--- a/include/winnt.h
+++ b/include/winnt.h
@@ -6501,6 +6501,8 @@ typedef struct _GROUP_AFFINITY
     WORD Reserved[3];
 } GROUP_AFFINITY, *PGROUP_AFFINITY;
 
+#define ALL_PROCESSOR_GROUPS 0xffff
+
 typedef struct _PROCESSOR_NUMBER
 {
     WORD Group;
-- 
2.26.2




More information about the wine-devel mailing list