Signed-off-by: Paul Gofman pgofman@codeweavers.com --- v2: - Use GetActiveProcessorCount() function.
dlls/ntoskrnl.exe/ntoskrnl.c | 6 ++++++ 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, 41 insertions(+)
diff --git a/dlls/ntoskrnl.exe/ntoskrnl.c b/dlls/ntoskrnl.exe/ntoskrnl.c index aa71b7c48c..62e28d2ce5 100644 --- a/dlls/ntoskrnl.exe/ntoskrnl.c +++ b/dlls/ntoskrnl.exe/ntoskrnl.c @@ -2385,6 +2385,12 @@ KAFFINITY WINAPI KeQueryActiveProcessors( void ) return AffinityMask; }
+ULONG WINAPI KeQueryActiveProcessorCountEx(USHORT group_number) +{ + TRACE("group_number %u.\n", group_number); + + return GetActiveProcessorCount(group_number); +}
/********************************************************************** * 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..65be4a8d35 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); + todo_wine 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;