Signed-off-by: Paul Gofman pgofman@codeweavers.com --- dlls/ntoskrnl.exe/ntoskrnl.c | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/dlls/ntoskrnl.exe/ntoskrnl.c b/dlls/ntoskrnl.exe/ntoskrnl.c index 6d9ae44808..c10c830738 100644 --- a/dlls/ntoskrnl.exe/ntoskrnl.c +++ b/dlls/ntoskrnl.exe/ntoskrnl.c @@ -2379,10 +2379,10 @@ LONG WINAPI KeInsertQueue(PRKQUEUE Queue, PLIST_ENTRY Entry) */ KAFFINITY WINAPI KeQueryActiveProcessors( void ) { - DWORD_PTR AffinityMask; + DWORD_PTR affinity_mask;
- GetProcessAffinityMask( GetCurrentProcess(), &AffinityMask, NULL); - return AffinityMask; + GetProcessAffinityMask( GetCurrentProcess(), NULL, &affinity_mask); + return affinity_mask; }
ULONG WINAPI KeQueryActiveProcessorCountEx(USHORT group_number)
Signed-off-by: Paul Gofman pgofman@codeweavers.com --- dlls/ntoskrnl.exe/ntoskrnl.c | 19 +++++++++++++++++++ dlls/ntoskrnl.exe/ntoskrnl.exe.spec | 1 + dlls/ntoskrnl.exe/tests/driver.c | 19 +++++++++++++++++-- include/ddk/wdm.h | 1 + 4 files changed, 38 insertions(+), 2 deletions(-)
diff --git a/dlls/ntoskrnl.exe/ntoskrnl.c b/dlls/ntoskrnl.exe/ntoskrnl.c index c10c830738..dc6bbe866c 100644 --- a/dlls/ntoskrnl.exe/ntoskrnl.c +++ b/dlls/ntoskrnl.exe/ntoskrnl.c @@ -2455,6 +2455,25 @@ VOID WINAPI KeSetSystemAffinityThread(KAFFINITY Affinity) FIXME("(%lx) stub\n", Affinity); }
+KAFFINITY WINAPI KeSetSystemAffinityThreadEx(KAFFINITY affinity) +{ + DWORD_PTR system_affinity = KeQueryActiveProcessors(); + GROUP_AFFINITY old, new; + + TRACE("affinity %#lx.\n", affinity); + + affinity &= system_affinity; + + NtQueryInformationThread(GetCurrentThread(), ThreadGroupInformation, + &old, sizeof(old), NULL); + + memset(&new, 0, sizeof(new)); + new.Mask = affinity; + + return NtSetInformationThread(GetCurrentThread(), ThreadGroupInformation, &new, sizeof(new)) + ? 0 : old.Mask; +} +
/*********************************************************************** * KeRevertToUserAffinityThread (NTOSKRNL.EXE.@) diff --git a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec index f25ab9c5e0..e6e641d16c 100644 --- a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec +++ b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec @@ -634,6 +634,7 @@ @ stdcall KeSetPriorityThread(ptr long) @ stub KeSetProfileIrql @ stdcall KeSetSystemAffinityThread(long) +@ stdcall KeSetSystemAffinityThreadEx(long) @ stdcall KeSetTargetProcessorDpc(ptr long) @ stub KeSetTimeIncrement @ stub KeSetTimer diff --git a/dlls/ntoskrnl.exe/tests/driver.c b/dlls/ntoskrnl.exe/tests/driver.c index 65be4a8d35..da6b9c1d63 100644 --- a/dlls/ntoskrnl.exe/tests/driver.c +++ b/dlls/ntoskrnl.exe/tests/driver.c @@ -1712,10 +1712,11 @@ static void test_executable_pool(void)
static void test_affinity(void) { + KAFFINITY (WINAPI *pKeSetSystemAffinityThreadEx)(KAFFINITY affinity); ULONG (WINAPI *pKeQueryActiveProcessorCountEx)(USHORT); KAFFINITY (WINAPI *pKeQueryActiveProcessors)(void); + KAFFINITY mask, mask_all_cpus; ULONG cpu_count, count; - KAFFINITY mask;
pKeQueryActiveProcessorCountEx = get_proc_address("KeQueryActiveProcessorCountEx"); if (!pKeQueryActiveProcessorCountEx) @@ -1727,6 +1728,9 @@ static void test_affinity(void) pKeQueryActiveProcessors = get_proc_address("KeQueryActiveProcessors"); ok(!!pKeQueryActiveProcessors, "KeQueryActiveProcessors is not available.\n");
+ pKeSetSystemAffinityThreadEx = get_proc_address("KeSetSystemAffinityThreadEx"); + ok(!!pKeSetSystemAffinityThreadEx, "KeSetSystemAffinityThreadEx is not available.\n"); + count = pKeQueryActiveProcessorCountEx(1); todo_wine ok(!count, "Got unexpected count %u.\n", count);
@@ -1736,8 +1740,19 @@ static void test_affinity(void) count = pKeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS); ok(count == cpu_count, "Got unexpected count %u.\n", count);
+ mask_all_cpus = ~((~0u) << cpu_count); + mask = pKeQueryActiveProcessors(); - ok(mask == ~((~0u) << cpu_count), "Got unexpected mask %#lx.\n", mask); + ok(mask == mask_all_cpus, "Got unexpected mask %#lx.\n", mask); + + mask = pKeSetSystemAffinityThreadEx(0); + ok(!mask, "Got unexpected mask %#lx.\n", mask); + + mask = pKeSetSystemAffinityThreadEx(0x1); + ok(mask == mask_all_cpus, "Got unexpected mask %#lx.\n", mask); + + mask = pKeSetSystemAffinityThreadEx(~0ull); + ok(mask == 0x1, "Got unexpected mask %#lx.\n", mask); }
static NTSTATUS main_test(DEVICE_OBJECT *device, IRP *irp, IO_STACK_LOCATION *stack) diff --git a/include/ddk/wdm.h b/include/ddk/wdm.h index 1f80a5af29..42ff0c72b1 100644 --- a/include/ddk/wdm.h +++ b/include/ddk/wdm.h @@ -1709,6 +1709,7 @@ void WINAPI KeRevertToUserAffinityThread(void); LONG WINAPI KeSetEvent(PRKEVENT,KPRIORITY,BOOLEAN); KPRIORITY WINAPI KeSetPriorityThread(PKTHREAD,KPRIORITY); void WINAPI KeSetSystemAffinityThread(KAFFINITY); +KAFFINITY WINAPI KeSetSystemAffinityThreadEx(KAFFINITY affinity); BOOLEAN WINAPI KeSetTimerEx(KTIMER*,LARGE_INTEGER,LONG,KDPC*); NTSTATUS WINAPI KeWaitForMultipleObjects(ULONG,void*[],WAIT_TYPE,KWAIT_REASON,KPROCESSOR_MODE,BOOLEAN,LARGE_INTEGER*,KWAIT_BLOCK*); NTSTATUS WINAPI KeWaitForSingleObject(void*,KWAIT_REASON,KPROCESSOR_MODE,BOOLEAN,LARGE_INTEGER*);
Signed-off-by: Paul Gofman pgofman@codeweavers.com --- dlls/ntoskrnl.exe/ntoskrnl.c | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/dlls/ntoskrnl.exe/ntoskrnl.c b/dlls/ntoskrnl.exe/ntoskrnl.c index dc6bbe866c..a0657b52c5 100644 --- a/dlls/ntoskrnl.exe/ntoskrnl.c +++ b/dlls/ntoskrnl.exe/ntoskrnl.c @@ -2450,9 +2450,9 @@ KPRIORITY WINAPI KeSetPriorityThread( PKTHREAD Thread, KPRIORITY Priority ) /*********************************************************************** * KeSetSystemAffinityThread (NTOSKRNL.EXE.@) */ -VOID WINAPI KeSetSystemAffinityThread(KAFFINITY Affinity) +VOID WINAPI KeSetSystemAffinityThread(KAFFINITY affinity) { - FIXME("(%lx) stub\n", Affinity); + KeSetSystemAffinityThreadEx(affinity); }
KAFFINITY WINAPI KeSetSystemAffinityThreadEx(KAFFINITY affinity)
Signed-off-by: Paul Gofman pgofman@codeweavers.com --- dlls/ntoskrnl.exe/ntoskrnl.c | 24 +++++++++++++++++++++++- dlls/ntoskrnl.exe/ntoskrnl.exe.spec | 1 + dlls/ntoskrnl.exe/ntoskrnl_private.h | 1 + dlls/ntoskrnl.exe/tests/driver.c | 16 +++++++++++++++- include/ddk/wdm.h | 1 + 5 files changed, 41 insertions(+), 2 deletions(-)
diff --git a/dlls/ntoskrnl.exe/ntoskrnl.c b/dlls/ntoskrnl.exe/ntoskrnl.c index a0657b52c5..fc7828cf15 100644 --- a/dlls/ntoskrnl.exe/ntoskrnl.c +++ b/dlls/ntoskrnl.exe/ntoskrnl.c @@ -2267,6 +2267,7 @@ static void *create_thread_object( HANDLE handle )
thread->header.Type = 6; thread->header.WaitListHead.Blink = INVALID_HANDLE_VALUE; /* mark as kernel object */ + thread->user_affinity = 0;
if (!NtQueryInformationThread( handle, ThreadBasicInformation, &info, sizeof(info), NULL )) { @@ -2458,6 +2459,7 @@ VOID WINAPI KeSetSystemAffinityThread(KAFFINITY affinity) KAFFINITY WINAPI KeSetSystemAffinityThreadEx(KAFFINITY affinity) { DWORD_PTR system_affinity = KeQueryActiveProcessors(); + PKTHREAD thread = KeGetCurrentThread(); GROUP_AFFINITY old, new;
TRACE("affinity %#lx.\n", affinity); @@ -2467,11 +2469,14 @@ KAFFINITY WINAPI KeSetSystemAffinityThreadEx(KAFFINITY affinity) NtQueryInformationThread(GetCurrentThread(), ThreadGroupInformation, &old, sizeof(old), NULL);
+ if (old.Mask != system_affinity) + thread->user_affinity = old.Mask; + memset(&new, 0, sizeof(new)); new.Mask = affinity;
return NtSetInformationThread(GetCurrentThread(), ThreadGroupInformation, &new, sizeof(new)) - ? 0 : old.Mask; + ? 0 : thread->user_affinity; }
@@ -2483,6 +2488,23 @@ void WINAPI KeRevertToUserAffinityThread(void) FIXME("() stub\n"); }
+void WINAPI KeRevertToUserAffinityThreadEx(KAFFINITY affinity) +{ + DWORD_PTR system_affinity = KeQueryActiveProcessors(); + PRKTHREAD thread = KeGetCurrentThread(); + GROUP_AFFINITY new; + + TRACE("affinity %#lx.\n", affinity); + + affinity &= system_affinity; + + memset(&new, 0, sizeof(new)); + new.Mask = affinity ? affinity + : (thread->user_affinity ? thread->user_affinity : system_affinity); + + NtSetInformationThread(GetCurrentThread(), ThreadGroupInformation, &new, sizeof(new)); + thread->user_affinity = affinity; +}
/*********************************************************************** * IoRegisterFileSystem (NTOSKRNL.EXE.@) diff --git a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec index e6e641d16c..f9c07bcf79 100644 --- a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec +++ b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec @@ -619,6 +619,7 @@ @ stdcall KeResetEvent(ptr) @ stub KeRestoreFloatingPointState @ stdcall KeRevertToUserAffinityThread() +@ stdcall KeRevertToUserAffinityThreadEx(long) @ stub KeRundownQueue @ stub KeSaveFloatingPointState @ stub KeSaveStateForHibernate diff --git a/dlls/ntoskrnl.exe/ntoskrnl_private.h b/dlls/ntoskrnl.exe/ntoskrnl_private.h index 23f03c336d..a1e1b892e8 100644 --- a/dlls/ntoskrnl.exe/ntoskrnl_private.h +++ b/dlls/ntoskrnl.exe/ntoskrnl_private.h @@ -49,6 +49,7 @@ struct _KTHREAD PEPROCESS process; CLIENT_ID id; unsigned int critical_region; + KAFFINITY user_affinity; };
struct _ETHREAD diff --git a/dlls/ntoskrnl.exe/tests/driver.c b/dlls/ntoskrnl.exe/tests/driver.c index da6b9c1d63..6b9b55732c 100644 --- a/dlls/ntoskrnl.exe/tests/driver.c +++ b/dlls/ntoskrnl.exe/tests/driver.c @@ -1713,6 +1713,7 @@ static void test_executable_pool(void) static void test_affinity(void) { KAFFINITY (WINAPI *pKeSetSystemAffinityThreadEx)(KAFFINITY affinity); + void (WINAPI *pKeRevertToUserAffinityThreadEx)(KAFFINITY affinity); ULONG (WINAPI *pKeQueryActiveProcessorCountEx)(USHORT); KAFFINITY (WINAPI *pKeQueryActiveProcessors)(void); KAFFINITY mask, mask_all_cpus; @@ -1731,6 +1732,9 @@ static void test_affinity(void) pKeSetSystemAffinityThreadEx = get_proc_address("KeSetSystemAffinityThreadEx"); ok(!!pKeSetSystemAffinityThreadEx, "KeSetSystemAffinityThreadEx is not available.\n");
+ pKeRevertToUserAffinityThreadEx = get_proc_address("KeRevertToUserAffinityThreadEx"); + ok(!!pKeRevertToUserAffinityThreadEx, "KeRevertToUserAffinityThreadEx is not available.\n"); + count = pKeQueryActiveProcessorCountEx(1); todo_wine ok(!count, "Got unexpected count %u.\n", count);
@@ -1745,14 +1749,24 @@ static void test_affinity(void) mask = pKeQueryActiveProcessors(); ok(mask == mask_all_cpus, "Got unexpected mask %#lx.\n", mask);
+ pKeRevertToUserAffinityThreadEx(0x2); + mask = pKeSetSystemAffinityThreadEx(0); ok(!mask, "Got unexpected mask %#lx.\n", mask);
+ pKeRevertToUserAffinityThreadEx(0x2); + mask = pKeSetSystemAffinityThreadEx(0x1); - ok(mask == mask_all_cpus, "Got unexpected mask %#lx.\n", mask); + ok(mask == 0x2, "Got unexpected mask %#lx.\n", mask);
mask = pKeSetSystemAffinityThreadEx(~0ull); ok(mask == 0x1, "Got unexpected mask %#lx.\n", mask); + + pKeRevertToUserAffinityThreadEx(~0ull); + mask = pKeSetSystemAffinityThreadEx(0x1); + ok(mask == mask_all_cpus, "Got unexpected mask %#lx.\n", mask); + + pKeRevertToUserAffinityThreadEx(mask_all_cpus); }
static NTSTATUS main_test(DEVICE_OBJECT *device, IRP *irp, IO_STACK_LOCATION *stack) diff --git a/include/ddk/wdm.h b/include/ddk/wdm.h index 42ff0c72b1..b0ad0e70d3 100644 --- a/include/ddk/wdm.h +++ b/include/ddk/wdm.h @@ -1706,6 +1706,7 @@ void WINAPI KeReleaseSpinLock(KSPIN_LOCK*,KIRQL); void WINAPI KeReleaseSpinLockFromDpcLevel(KSPIN_LOCK*); LONG WINAPI KeResetEvent(PRKEVENT); void WINAPI KeRevertToUserAffinityThread(void); +void WINAPI KeRevertToUserAffinityThreadEx(KAFFINITY affinity); LONG WINAPI KeSetEvent(PRKEVENT,KPRIORITY,BOOLEAN); KPRIORITY WINAPI KeSetPriorityThread(PKTHREAD,KPRIORITY); void WINAPI KeSetSystemAffinityThread(KAFFINITY);
Signed-off-by: Paul Gofman pgofman@codeweavers.com --- dlls/ntoskrnl.exe/ntoskrnl.c | 2 +- dlls/ntoskrnl.exe/tests/driver.c | 12 +++++++++++- 2 files changed, 12 insertions(+), 2 deletions(-)
diff --git a/dlls/ntoskrnl.exe/ntoskrnl.c b/dlls/ntoskrnl.exe/ntoskrnl.c index fc7828cf15..c4f2180c4c 100644 --- a/dlls/ntoskrnl.exe/ntoskrnl.c +++ b/dlls/ntoskrnl.exe/ntoskrnl.c @@ -2485,7 +2485,7 @@ KAFFINITY WINAPI KeSetSystemAffinityThreadEx(KAFFINITY affinity) */ void WINAPI KeRevertToUserAffinityThread(void) { - FIXME("() stub\n"); + KeRevertToUserAffinityThreadEx(0); }
void WINAPI KeRevertToUserAffinityThreadEx(KAFFINITY affinity) diff --git a/dlls/ntoskrnl.exe/tests/driver.c b/dlls/ntoskrnl.exe/tests/driver.c index 6b9b55732c..0c6a6d9561 100644 --- a/dlls/ntoskrnl.exe/tests/driver.c +++ b/dlls/ntoskrnl.exe/tests/driver.c @@ -1766,7 +1766,17 @@ static void test_affinity(void) mask = pKeSetSystemAffinityThreadEx(0x1); ok(mask == mask_all_cpus, "Got unexpected mask %#lx.\n", mask);
- pKeRevertToUserAffinityThreadEx(mask_all_cpus); + pKeRevertToUserAffinityThreadEx(0); + + mask = pKeSetSystemAffinityThreadEx(0x1); + ok(!mask, "Got unexpected mask %#lx.\n", mask); + + KeRevertToUserAffinityThread(); + + mask = pKeSetSystemAffinityThreadEx(0x1); + ok(!mask, "Got unexpected mask %#lx.\n", mask); + + KeRevertToUserAffinityThread(); }
static NTSTATUS main_test(DEVICE_OBJECT *device, IRP *irp, IO_STACK_LOCATION *stack)
-
Paul Gofman