Skip to content

Commit

Permalink
Add group affinity (#240)
Browse files Browse the repository at this point in the history 
* Switch to group affinity

Signed-off-by: Alan Jowett <[email protected]>

* PR feedback

Signed-off-by: Alan Jowett <[email protected]>

* Bug fix

Signed-off-by: Alan Jowett <[email protected]>

* Fix bug in test

Signed-off-by: Alan Jowett <[email protected]>

* Revert catch2 change

Signed-off-by: Alan Jowett <[email protected]>

* PR feedback

Signed-off-by: Alan Jowett <[email protected]>

---------

Signed-off-by: Alan Jowett <[email protected]>
Signed-off-by: Alan Jowett <[email protected]>
Co-authored-by: Alan Jowett <[email protected]>
  • Loading branch information
@Alan-Jowett
Alan-Jowett and Alan Jowett authored Dec 18, 2024
1 parent 5017e62 commit ee55fba
Show file tree
Hide file tree
Showing 4 changed files with 115 additions and 112 deletions.
8 changes: 4 additions & 4 deletions inc/usersim/ke.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
204 changes: 108 additions & 96 deletions src/ke.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<std::mutex> _usersim_dispatch_locks;
Expand Down Expand Up @@ -110,6 +113,71 @@ 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)
{
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;
}
_usersim_thread_priority_cache = priority;
}
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.
*/
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) {
result = true;
} else {
result = SetThreadGroupAffinity(GetCurrentThread(), new_affinity, old_affinity);
}

if (result) {
if (old_affinity != nullptr) {
*old_affinity = _usersim_group_affinity_cache;
}
_usersim_group_affinity_cache = *new_affinity;
}
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,
Expand All @@ -125,55 +193,15 @@ _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 (_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);
} 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
Expand All @@ -195,13 +223,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();
}

Expand All @@ -223,12 +249,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);
Expand Down Expand Up @@ -330,21 +352,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 old_affinity.Mask;
} else {
return 0;
}
}
return (KAFFINITY)old_affinity.Mask;
}

_IRQL_requires_min_(PASSIVE_LEVEL) _IRQL_requires_max_(APC_LEVEL) NTKERNELAPI VOID
Expand All @@ -356,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 (!SetThreadGroupAffinity(GetCurrentThread(), Affinity, PreviousAffinity)) {
if (!usersim_set_current_thread_affinity(Affinity, PreviousAffinity)) {
DWORD error = GetLastError();
#if defined(NDEBUG)
UNREFERENCED_PARAMETER(error);
Expand All @@ -369,29 +392,18 @@ KeSetSystemGroupAffinityThread(_In_ const PGROUP_AFFINITY Affinity, _Out_opt_ PG
void
KeRevertToUserGroupAffinityThread(PGROUP_AFFINITY PreviousAffinity)
{
SetThreadGroupAffinity(GetCurrentThread(), PreviousAffinity, NULL);
bool result = usersim_set_current_thread_affinity(PreviousAffinity, nullptr);
#if defined(NDEBUG)
UNREFERENCED_PARAMETER(result);
#else
assert(result);
#endif
}

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
Expand Down
6 changes: 0 additions & 6 deletions src/platform.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -426,12 +426,6 @@ 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);

void
usersim_restore_current_thread_affinity(uintptr_t old_thread_affinity_mask);

typedef _Return_type_success_(return >= 0) LONG NTSTATUS;

/**
Expand Down
9 changes: 3 additions & 6 deletions tests/ke_test.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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]")
Expand Down Expand Up @@ -181,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 = {};
Expand Down

0 comments on commit ee55fba

Please sign in to comment.