Commit 7d4f4783 authored by Paul Gofman's avatar Paul Gofman Committed by Alexandre Julliard

ntoskrnl.exe: Implement KeRevertToUserAffinityThreadEx() function.

parent 63122756
......@@ -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.@)
......
......@@ -619,6 +619,7 @@
@ stdcall KeResetEvent(ptr)
@ stub KeRestoreFloatingPointState
@ stdcall KeRevertToUserAffinityThread()
@ stdcall KeRevertToUserAffinityThreadEx(long)
@ stub KeRundownQueue
@ stub KeSaveFloatingPointState
@ stub KeSaveStateForHibernate
......
......@@ -49,6 +49,7 @@ struct _KTHREAD
PEPROCESS process;
CLIENT_ID id;
unsigned int critical_region;
KAFFINITY user_affinity;
};
struct _ETHREAD
......
......@@ -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(~(KAFFINITY)0);
ok(mask == 0x1, "Got unexpected mask %#lx.\n", mask);
pKeRevertToUserAffinityThreadEx(~(KAFFINITY)0);
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)
......
......@@ -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);
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment