From 4ed48ca524ced4296ecd31988f4a687705774410 Mon Sep 17 00:00:00 2001 From: Alan Jowett Date: Sun, 3 Nov 2024 11:28:31 -0800 Subject: [PATCH 1/6] Switch to group affinity Signed-off-by: Alan Jowett --- src/ke.cpp | 197 +++++++++++++++++++++++++------------------------ src/platform.h | 6 +- 2 files changed, 104 insertions(+), 99 deletions(-) diff --git a/src/ke.cpp b/src/ke.cpp index fd6f29d..a6de027 100644 --- a/src/ke.cpp +++ b/src/ke.cpp @@ -21,11 +21,14 @@ #pragma region irqls -thread_local KIRQL _usersim_current_irql = PASSIVE_LEVEL; -thread_local GROUP_AFFINITY _usersim_dispatch_previous_affinity; -thread_local bool _usersim_affinity_and_priority_override = false; -thread_local GROUP_AFFINITY _usersim_affinity_before_override; -thread_local int _usersim_thread_priority_before_override; +thread_local KIRQL _usersim_current_irql = PASSIVE_LEVEL; ///< The current threads emulated IRQL. +thread_local GROUP_AFFINITY _usersim_group_affinity_cache = {}; ///< The effective group affinity of the current thread. +thread_local GROUP_AFFINITY + _usersim_group_before_raise_irql; ///< The group affinity of the current thread before raising IRQL. +thread_local int _usersim_thread_priority_cache = + THREAD_PRIORITY_NORMAL; ///< The effective priority of the current thread. +thread_local int + _usersim_thread_priority_before_raise_irql; ///< The priority of the current thread before raising IRQL. static uint32_t _usersim_original_priority_class; static std::vector _usersim_dispatch_locks; @@ -110,6 +113,72 @@ usersim_clean_up_irql() } } +/** + * @brief Change the priority of the current thread and cache the new priority. + * If the new priority is the same as the cached priority, this function is a no-op. + * + * @param[in] priority New priority to set. + * @param[out] old_priority Old priority of the thread. + * @return true Success. + * @return false Failure. + */ +bool +usersim_set_current_thread_priority(int priority, int* old_priority) +{ + if (_usersim_thread_priority_cache == priority) { + if (old_priority != nullptr) { + *old_priority = _usersim_thread_priority_cache; + } + return true; + } else { + bool result = SetThreadPriority(GetCurrentThread(), priority); + if (result) { + _usersim_thread_priority_cache = priority; + if (old_priority != nullptr) { + *old_priority = _usersim_thread_priority_cache; + } + } + return result; + } +} + +/** + * @brief Change the affinity of the current thread and cache the new affinity. + * If the new affinity is the same as the cached affinity, this function is a no-op. + * + * @param[in] new_affinity The new affinity to set. + * @param[out] old_affinity The old affinity of the thread. + * @return true Success. + * @return false Failure. + */ +bool +usersim_set_current_thread_affinity(const GROUP_AFFINITY* new_affinity, GROUP_AFFINITY* old_affinity) +{ + if (memcmp(new_affinity, &_usersim_group_affinity_cache, sizeof(*new_affinity)) == 0) { + if (old_affinity != nullptr) { + *old_affinity = _usersim_group_affinity_cache; + } + return true; + } + + bool result = SetThreadGroupAffinity(GetCurrentThread(), new_affinity, old_affinity); + if (result) { + _usersim_group_affinity_cache = *new_affinity; + if (old_affinity != nullptr) { + *old_affinity = _usersim_group_affinity_cache; + } + } + return result; +} + +void +usersim_restore_current_thread_affinity(GROUP_AFFINITY* old_affinity) +{ + if (old_affinity != nullptr) { + usersim_set_current_thread_affinity(old_affinity, nullptr); + } +} + const int _irql_thread_priority[3] = { /* PASSIVE_LEVEL */ THREAD_PRIORITY_NORMAL, /* APC_LEVEL */ THREAD_PRIORITY_ABOVE_NORMAL, @@ -125,55 +194,12 @@ _get_irql_thread_priority(KIRQL irql) inline BOOL _set_current_thread_priority_by_irql(KIRQL new_irql) { - if (_usersim_affinity_and_priority_override) { - return true; + if (new_irql > DISPATCH_LEVEL) { + return usersim_set_current_thread_priority( + _get_irql_thread_priority(new_irql), &_usersim_thread_priority_before_raise_irql); + } else { + return usersim_set_current_thread_priority(_usersim_thread_priority_before_raise_irql, nullptr); } - return SetThreadPriority(GetCurrentThread(), _get_irql_thread_priority(new_irql)); -} - -/** - * @brief Preset the affinity and priority of the current thread to the specified processor to - * lower the cost of of KeRaiseIrql and LeLowerIrql. - * - * Setting affinity and thread priority is a costly operation, so we want to avoid doing it - * every time we raise or lower the IRQL. Instead, we can set the affinity and priority of the - * current thread to the processor that the thread is currently running on. This way, we can - * avoid the cost of setting the affinity and priority every time we raise or lower the IRQL. - * - * This is useful primarily when running performance tests that involve a lot of calls to - * KeRaiseIrql and KeLowerIrql. - * - * @param processor_index The index of the processor to set the affinity to. - */ -void -usersim_set_affinity_and_priority_override(uint32_t processor_index) -{ - GetThreadGroupAffinity(GetCurrentThread(), &_usersim_affinity_before_override); - _usersim_thread_priority_before_override = GetThreadPriority(GetCurrentThread()); - - _set_current_thread_priority_by_irql(DISPATCH_LEVEL); - PROCESSOR_NUMBER processor; - KeGetProcessorNumberFromIndex(processor_index, &processor); - - GROUP_AFFINITY new_affinity = {0}; - new_affinity.Group = processor.Group; - new_affinity.Mask = (ULONG_PTR)1 << processor.Number; - bool result = SetThreadGroupAffinity(GetCurrentThread(), &new_affinity, nullptr); - ASSERT(result); - - _usersim_affinity_and_priority_override = true; -} - -/** - * @brief Restore the affinity and priority of the current thread to the values before the call to - * usersim_set_affinity_and_priority_override. - */ -void -usersim_clear_affinity_and_priority_override() -{ - SetThreadPriority(GetCurrentThread(), _usersim_thread_priority_before_override); - SetThreadGroupAffinity(GetCurrentThread(), &_usersim_affinity_before_override, nullptr); - _usersim_affinity_and_priority_override = false; } KIRQL @@ -195,13 +221,11 @@ _IRQL_requires_max_(HIGH_LEVEL) _IRQL_raises_(new_irql) _IRQL_saves_ KIRQL KfRai if (new_irql >= DISPATCH_LEVEL && old_irql < DISPATCH_LEVEL) { PROCESSOR_NUMBER processor; uint32_t processor_index = KeGetCurrentProcessorNumberEx(&processor); - if (!_usersim_affinity_and_priority_override) { - GROUP_AFFINITY new_affinity = {0}; - new_affinity.Group = processor.Group; - new_affinity.Mask = (ULONG_PTR)1 << processor.Number; - result = SetThreadGroupAffinity(GetCurrentThread(), &new_affinity, nullptr); - ASSERT(result); - } + GROUP_AFFINITY new_affinity = {0}; + new_affinity.Group = processor.Group; + new_affinity.Mask = (ULONG_PTR)1 << processor.Number; + result = usersim_set_current_thread_affinity(&new_affinity, &_usersim_group_before_raise_irql); + ASSERT(result); _usersim_dispatch_locks[processor_index].lock(); } @@ -223,12 +247,8 @@ KeLowerIrql(_In_ KIRQL new_irql) if (_usersim_current_irql >= DISPATCH_LEVEL && new_irql < DISPATCH_LEVEL) { uint32_t processor_index = KeGetCurrentProcessorNumberEx(nullptr); _usersim_dispatch_locks[processor_index].unlock(); - if (!_usersim_affinity_and_priority_override) { - GROUP_AFFINITY new_affinity; - usersim_get_current_thread_group_affinity(&new_affinity); - result = SetThreadGroupAffinity(GetCurrentThread(), &new_affinity, nullptr); - ASSERT(result); - } + result = usersim_set_current_thread_affinity(&_usersim_group_before_raise_irql, nullptr); + ASSERT(result); } _usersim_current_irql = new_irql; result = _set_current_thread_priority_by_irql(new_irql); @@ -330,21 +350,22 @@ KeQueryActiveProcessorCountEx(_In_ USHORT group_number) { return KeQueryMaximumP KAFFINITY KeSetSystemAffinityThreadEx(KAFFINITY affinity) { + GROUP_AFFINITY new_affinity = {0}; + new_affinity.Mask = affinity; + new_affinity.Group = _usersim_group_affinity_cache.Group; GROUP_AFFINITY old_affinity; - usersim_get_current_thread_group_affinity(&old_affinity); - // Reject affinities that are not valid for the current group. - // Assume that the affinity mask is contiguous. - KAFFINITY valid_affinity_mask = (1 << KeQueryMaximumProcessorCountEx(old_affinity.Group)) - 1; - if ((affinity & valid_affinity_mask) == 0) { - return 0; - } - _usersim_thread_affinity.Group = old_affinity.Group; - _usersim_thread_affinity.Mask = affinity; - if (KeGetCurrentIrql() < DISPATCH_LEVEL && SetThreadAffinityMask(GetCurrentThread(), affinity) == 0) { - _usersim_thread_affinity = old_affinity; + + if (KeGetCurrentIrql() >= DISPATCH_LEVEL) { + _usersim_group_affinity_cache = new_affinity; return 0; + } else { + bool result = usersim_set_current_thread_affinity(&new_affinity, &old_affinity); + if (result) { + return (KAFFINITY)old_affinity.Mask; + } else { + return 0; + } } - return (KAFFINITY)old_affinity.Mask; } _IRQL_requires_min_(PASSIVE_LEVEL) _IRQL_requires_max_(APC_LEVEL) NTKERNELAPI VOID @@ -356,7 +377,7 @@ _IRQL_requires_min_(PASSIVE_LEVEL) _IRQL_requires_max_(APC_LEVEL) NTKERNELAPI VO void KeSetSystemGroupAffinityThread(_In_ const PGROUP_AFFINITY Affinity, _Out_opt_ PGROUP_AFFINITY PreviousAffinity) { - if (!SetThreadGroupAffinity(GetCurrentThread(), Affinity, PreviousAffinity)) { + if (usersim_set_current_thread_affinity(Affinity, PreviousAffinity)) { DWORD error = GetLastError(); #if defined(NDEBUG) UNREFERENCED_PARAMETER(error); @@ -369,29 +390,13 @@ KeSetSystemGroupAffinityThread(_In_ const PGROUP_AFFINITY Affinity, _Out_opt_ PG void KeRevertToUserGroupAffinityThread(PGROUP_AFFINITY PreviousAffinity) { - SetThreadGroupAffinity(GetCurrentThread(), PreviousAffinity, NULL); + (void)usersim_set_current_thread_affinity(PreviousAffinity, nullptr); } PKTHREAD NTAPI KeGetCurrentThread(VOID) { return (PKTHREAD)usersim_get_current_thread_id(); } -void -usersim_get_current_thread_group_affinity(_Out_ GROUP_AFFINITY* affinity) -{ - if (_usersim_thread_affinity.Mask != 0) { - *affinity = _usersim_thread_affinity; - } else { - // The thread's current group affinity has never been explicitly set. Report the - // current affinity is all processors in the current group. - PROCESSOR_NUMBER processor; - KeGetCurrentProcessorNumberEx(&processor); - RtlZeroMemory(affinity, sizeof(*affinity)); - affinity->Group = processor.Group; - affinity->Mask = ((ULONG_PTR)1 << GetMaximumProcessorCount(affinity->Group)) - 1; - } -} - #pragma endregion threads #pragma region semaphores diff --git a/src/platform.h b/src/platform.h index 30b14ce..45d7099 100644 --- a/src/platform.h +++ b/src/platform.h @@ -426,11 +426,11 @@ usersim_random_uint32(); uint64_t usersim_query_time_since_boot(bool include_suspended_time); -_Must_inspect_result_ usersim_result_t -usersim_set_current_thread_affinity(uintptr_t new_thread_affinity_mask, _Out_ uintptr_t* old_thread_affinity_mask); +_Must_inspect_result_ bool +usersim_set_current_thread_affinity(const GROUP_AFFINITY* new_affinity, GROUP_AFFINITY* old_affinity); void -usersim_restore_current_thread_affinity(uintptr_t old_thread_affinity_mask); +usersim_restore_current_thread_affinity(GROUP_AFFINITY* old_affinity); typedef _Return_type_success_(return >= 0) LONG NTSTATUS; From 274d15dd70f9c3757f42ef751acb29ebf9eccf10 Mon Sep 17 00:00:00 2001 From: Alan Jowett Date: Sun, 3 Nov 2024 12:05:37 -0800 Subject: [PATCH 2/6] PR feedback Signed-off-by: Alan Jowett --- inc/usersim/ke.h | 8 ++++---- src/ke.cpp | 3 ++- src/platform.h | 6 ------ tests/ke_test.cpp | 4 ---- 4 files changed, 6 insertions(+), 15 deletions(-) diff --git a/inc/usersim/ke.h b/inc/usersim/ke.h index ea1671c..4413cc8 100644 --- a/inc/usersim/ke.h +++ b/inc/usersim/ke.h @@ -90,12 +90,12 @@ void usersim_clean_up_irql(); USERSIM_API -void -usersim_set_affinity_and_priority_override(uint32_t processor_index); +bool +usersim_set_current_thread_priority(int priority, int* old_priority); USERSIM_API -void -usersim_clear_affinity_and_priority_override(); +bool +usersim_set_current_thread_affinity(const GROUP_AFFINITY* new_affinity, GROUP_AFFINITY* old_affinity); #pragma endregion irqls diff --git a/src/ke.cpp b/src/ke.cpp index a6de027..d6d7ac8 100644 --- a/src/ke.cpp +++ b/src/ke.cpp @@ -151,6 +151,7 @@ usersim_set_current_thread_priority(int priority, int* old_priority) * @return true Success. * @return false Failure. */ +USERSIM_API bool usersim_set_current_thread_affinity(const GROUP_AFFINITY* new_affinity, GROUP_AFFINITY* old_affinity) { @@ -194,7 +195,7 @@ _get_irql_thread_priority(KIRQL irql) inline BOOL _set_current_thread_priority_by_irql(KIRQL new_irql) { - if (new_irql > DISPATCH_LEVEL) { + if (new_irql >= DISPATCH_LEVEL) { return usersim_set_current_thread_priority( _get_irql_thread_priority(new_irql), &_usersim_thread_priority_before_raise_irql); } else { diff --git a/src/platform.h b/src/platform.h index 45d7099..8235e4b 100644 --- a/src/platform.h +++ b/src/platform.h @@ -426,12 +426,6 @@ usersim_random_uint32(); uint64_t usersim_query_time_since_boot(bool include_suspended_time); -_Must_inspect_result_ bool -usersim_set_current_thread_affinity(const GROUP_AFFINITY* new_affinity, GROUP_AFFINITY* old_affinity); - -void -usersim_restore_current_thread_affinity(GROUP_AFFINITY* old_affinity); - typedef _Return_type_success_(return >= 0) LONG NTSTATUS; /** diff --git a/tests/ke_test.cpp b/tests/ke_test.cpp index e812ab7..17b0030 100644 --- a/tests/ke_test.cpp +++ b/tests/ke_test.cpp @@ -40,8 +40,6 @@ TEST_CASE("irql", "[ke]") TEST_CASE("irql_perf_override", "[ke]") { - usersim_set_affinity_and_priority_override(0); - REQUIRE(KeGetCurrentIrql() == PASSIVE_LEVEL); KIRQL old_irql; @@ -65,8 +63,6 @@ TEST_CASE("irql_perf_override", "[ke]") KeLowerIrql(old_irql); REQUIRE(KeGetCurrentIrql() == PASSIVE_LEVEL); - - usersim_clear_affinity_and_priority_override(); } TEST_CASE("KfRaiseIrql", "[ke]") From 5c7fe8ce05c4b5c9f2f6fa53f7e456bfbb792ae9 Mon Sep 17 00:00:00 2001 From: Alan Jowett Date: Tue, 3 Dec 2024 12:13:39 -0800 Subject: [PATCH 3/6] Bug fix Signed-off-by: Alan Jowett --- external/Catch2 | 2 +- src/ke.cpp | 39 ++++++++++++++++++++------------------- 2 files changed, 21 insertions(+), 20 deletions(-) diff --git a/external/Catch2 b/external/Catch2 index 0321d2f..506276c 160000 --- a/external/Catch2 +++ b/external/Catch2 @@ -1 +1 @@ -Subproject commit 0321d2fce328b5e2ad106a8230ff20e0d5bf5501 +Subproject commit 506276c59217429c93abd2fe9507c7f45eb81072 diff --git a/src/ke.cpp b/src/ke.cpp index d6d7ac8..2ee9855 100644 --- a/src/ke.cpp +++ b/src/ke.cpp @@ -125,21 +125,20 @@ usersim_clean_up_irql() bool usersim_set_current_thread_priority(int priority, int* old_priority) { + bool result; if (_usersim_thread_priority_cache == priority) { + result = true; + } else { + result = SetThreadPriority(GetCurrentThread(), priority); + } + + if (result) { if (old_priority != nullptr) { *old_priority = _usersim_thread_priority_cache; } - return true; - } else { - bool result = SetThreadPriority(GetCurrentThread(), priority); - if (result) { - _usersim_thread_priority_cache = priority; - if (old_priority != nullptr) { - *old_priority = _usersim_thread_priority_cache; - } - } - return result; + _usersim_thread_priority_cache = priority; } + return result; } /** @@ -155,19 +154,18 @@ USERSIM_API bool usersim_set_current_thread_affinity(const GROUP_AFFINITY* new_affinity, GROUP_AFFINITY* old_affinity) { + bool result; if (memcmp(new_affinity, &_usersim_group_affinity_cache, sizeof(*new_affinity)) == 0) { - if (old_affinity != nullptr) { - *old_affinity = _usersim_group_affinity_cache; - } - return true; + result = true; + } else { + result = SetThreadGroupAffinity(GetCurrentThread(), new_affinity, old_affinity); } - bool result = SetThreadGroupAffinity(GetCurrentThread(), new_affinity, old_affinity); if (result) { - _usersim_group_affinity_cache = *new_affinity; if (old_affinity != nullptr) { *old_affinity = _usersim_group_affinity_cache; } + _usersim_group_affinity_cache = *new_affinity; } return result; } @@ -195,6 +193,9 @@ _get_irql_thread_priority(KIRQL irql) inline BOOL _set_current_thread_priority_by_irql(KIRQL new_irql) { + if (_usersim_current_irql == new_irql) { + return TRUE; + } if (new_irql >= DISPATCH_LEVEL) { return usersim_set_current_thread_priority( _get_irql_thread_priority(new_irql), &_usersim_thread_priority_before_raise_irql); @@ -378,7 +379,7 @@ _IRQL_requires_min_(PASSIVE_LEVEL) _IRQL_requires_max_(APC_LEVEL) NTKERNELAPI VO void KeSetSystemGroupAffinityThread(_In_ const PGROUP_AFFINITY Affinity, _Out_opt_ PGROUP_AFFINITY PreviousAffinity) { - if (usersim_set_current_thread_affinity(Affinity, PreviousAffinity)) { + if (!usersim_set_current_thread_affinity(Affinity, PreviousAffinity)) { DWORD error = GetLastError(); #if defined(NDEBUG) UNREFERENCED_PARAMETER(error); @@ -653,10 +654,10 @@ class _usersim_emulated_dpc l.unlock(); _usersim_dispatch_locks[i].lock(); - _usersim_current_irql = DISPATCH_LEVEL; + //_usersim_current_irql = DISPATCH_LEVEL; work_item->work_item_routine(work_item, context, parameter_1, parameter_2); _usersim_dispatch_locks[i].unlock(); - _usersim_current_irql = PASSIVE_LEVEL; + //_usersim_current_irql = PASSIVE_LEVEL; l.lock(); } } From 7633aee1a5818b03f58bbdcfca86b5b9146121f0 Mon Sep 17 00:00:00 2001 From: Alan Jowett Date: Tue, 3 Dec 2024 12:58:05 -0800 Subject: [PATCH 4/6] Fix bug in test Signed-off-by: Alan Jowett --- tests/ke_test.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/ke_test.cpp b/tests/ke_test.cpp index 17b0030..1fedaee 100644 --- a/tests/ke_test.cpp +++ b/tests/ke_test.cpp @@ -177,11 +177,12 @@ TEST_CASE("threads", "[ke]") KAFFINITY new_affinity = ((ULONG_PTR)1 << processor_count) - 1; KAFFINITY old_affinity = KeSetSystemAffinityThreadEx(new_affinity); - REQUIRE(old_affinity != 0); + // No old affinity was set. + REQUIRE(old_affinity == 0); KeRevertToUserAffinityThreadEx(old_affinity); - for (ULONG i = 0; i < processor_count; i++) { + for (ULONG i = 0; i < processor_count; i++) { PROCESSOR_NUMBER processor_number = {}; GROUP_AFFINITY old_group_affinity = {}; GROUP_AFFINITY new_group_affinity = {}; From 5100b266033ce9be9d6b2c4b0bc39eff477955f3 Mon Sep 17 00:00:00 2001 From: Alan Jowett Date: Tue, 3 Dec 2024 13:36:02 -0800 Subject: [PATCH 5/6] Revert catch2 change Signed-off-by: Alan Jowett --- external/Catch2 | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/external/Catch2 b/external/Catch2 index 506276c..0321d2f 160000 --- a/external/Catch2 +++ b/external/Catch2 @@ -1 +1 @@ -Subproject commit 506276c59217429c93abd2fe9507c7f45eb81072 +Subproject commit 0321d2fce328b5e2ad106a8230ff20e0d5bf5501 From 0f29f557f0432f62e1fc73b4a13dc7a60bcfa508 Mon Sep 17 00:00:00 2001 From: Alan Jowett Date: Tue, 17 Dec 2024 15:25:53 -0800 Subject: [PATCH 6/6] PR feedback Signed-off-by: Alan Jowett --- src/ke.cpp | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/src/ke.cpp b/src/ke.cpp index 2ee9855..0a6591e 100644 --- a/src/ke.cpp +++ b/src/ke.cpp @@ -363,7 +363,7 @@ KeSetSystemAffinityThreadEx(KAFFINITY affinity) } else { bool result = usersim_set_current_thread_affinity(&new_affinity, &old_affinity); if (result) { - return (KAFFINITY)old_affinity.Mask; + return old_affinity.Mask; } else { return 0; } @@ -392,7 +392,12 @@ KeSetSystemGroupAffinityThread(_In_ const PGROUP_AFFINITY Affinity, _Out_opt_ PG void KeRevertToUserGroupAffinityThread(PGROUP_AFFINITY PreviousAffinity) { - (void)usersim_set_current_thread_affinity(PreviousAffinity, nullptr); + bool result = usersim_set_current_thread_affinity(PreviousAffinity, nullptr); +#if defined(NDEBUG) + UNREFERENCED_PARAMETER(result); +#else + assert(result); +#endif } PKTHREAD @@ -654,10 +659,10 @@ class _usersim_emulated_dpc l.unlock(); _usersim_dispatch_locks[i].lock(); - //_usersim_current_irql = DISPATCH_LEVEL; + _usersim_current_irql = DISPATCH_LEVEL; work_item->work_item_routine(work_item, context, parameter_1, parameter_2); _usersim_dispatch_locks[i].unlock(); - //_usersim_current_irql = PASSIVE_LEVEL; + _usersim_current_irql = PASSIVE_LEVEL; l.lock(); } }