Skip to content

Add processor-topology test hooks #311

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Alan-Jowett wants to merge 2 commits into
base: main
Choose a base branch
from
+417 −3
Open

Add processor-topology test hooks #311

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
411456f
Add processor-topology test hooks
Jul 1, 2026
6aab0fe
Address processor callback review comments
Alan-Jowett Jul 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 70 additions & 0 deletions inc/usersim/ke.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,76 @@ USERSIM_API
ULONG
KeQueryActiveProcessorCountEx(_In_ USHORT group_number);

typedef enum _KE_PROCESSOR_CHANGE_NOTIFY_STATE
{
KeProcessorAddStartNotify,
KeProcessorAddCompleteNotify,
KeProcessorAddFailureNotify,
} KE_PROCESSOR_CHANGE_NOTIFY_STATE;

typedef struct _KE_PROCESSOR_CHANGE_NOTIFY_CONTEXT
{
KE_PROCESSOR_CHANGE_NOTIFY_STATE State;
ULONG NtNumber;
} KE_PROCESSOR_CHANGE_NOTIFY_CONTEXT;
typedef KE_PROCESSOR_CHANGE_NOTIFY_CONTEXT* PKE_PROCESSOR_CHANGE_NOTIFY_CONTEXT;

#define KE_PROCESSOR_CHANGE_ADD_EXISTING 0x00000001

typedef VOID (*KE_PROCESSOR_CHANGE_CALLBACK)(
_In_ void* callback_context,
_In_ PKE_PROCESSOR_CHANGE_NOTIFY_CONTEXT change_context,
_Inout_ PNTSTATUS operation_status);

USERSIM_API
void*
KeRegisterProcessorChangeCallback(
_In_ KE_PROCESSOR_CHANGE_CALLBACK callback, _In_opt_ void* callback_context, _In_ ULONG flags);

USERSIM_API
void
KeDeregisterProcessorChangeCallback(_In_ void* callback_handle);

USERSIM_API
_Must_inspect_result_ NTSTATUS
usersim_set_active_processor_count(_In_ ULONG active_processor_count);

USERSIM_API
void
usersim_reset_active_processor_count();

/**
* @brief Start a simulated processor hot-add for the next inactive processor.
*
* The test hook models processor additions serially. Callers must start with
* the current active processor count and must finish the same processor with
* usersim_notify_processor_add_complete() or usersim_notify_processor_add_failure()
* before starting another add.
*/
USERSIM_API
_Must_inspect_result_ NTSTATUS
usersim_notify_processor_add_start(_In_ ULONG processor_index);

/**
* @brief Complete a simulated processor hot-add started by usersim_notify_processor_add_start().
*
* The processor index must match the currently pending add and becomes visible
* via KeQueryActiveProcessorCount*() once the completion notification succeeds.
*/
USERSIM_API
_Must_inspect_result_ NTSTATUS
usersim_notify_processor_add_complete(_In_ ULONG processor_index);

/**
* @brief Fail a simulated processor hot-add started by usersim_notify_processor_add_start().
*
* The processor index must match the currently pending add. Failing an add
* leaves the active processor count unchanged.
*/
USERSIM_API
_Must_inspect_result_ NTSTATUS
usersim_notify_processor_add_failure(_In_ ULONG processor_index);

USERSIM_API
KAFFINITY
KeSetSystemAffinityThreadEx(KAFFINITY affinity);
Expand Down
222 changes: 219 additions & 3 deletions src/ke.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,10 @@
#include "usersim/ke.h"
#include "utilities.h"

#include <algorithm>
#include <atomic>
#include <format>
#include <memory>
#include <mutex>
#include <optional>
#include <sstream>
Expand All @@ -32,13 +35,67 @@ thread_local int

static uint32_t _usersim_original_priority_class;
static std::vector<std::mutex> _usersim_dispatch_locks;
static std::mutex _usersim_processor_change_callback_lock;
static std::mutex _usersim_processor_change_state_lock;
static std::atomic<ULONG> _usersim_active_processor_count_override = 0;
static std::optional<ULONG> _usersim_pending_processor_add;

static TP_POOL* _usersim_threadpool = nullptr;
static std::optional<TP_CALLBACK_ENVIRON> _usersim_threadpool_callback_environment = std::nullopt;

typedef struct _usersim_processor_change_callback_registration
{
KE_PROCESSOR_CHANGE_CALLBACK callback;
void* callback_context;
} usersim_processor_change_callback_registration_t;

static std::vector<std::shared_ptr<usersim_processor_change_callback_registration_t>>
_usersim_processor_change_callbacks;

static NTSTATUS
_wait_for_kevent(_Inout_ KEVENT* event, _In_opt_ PLARGE_INTEGER timeout);

static _Must_inspect_result_ NTSTATUS
_usersim_invoke_processor_change_callback(
_In_ const std::shared_ptr<usersim_processor_change_callback_registration_t>& registration,
_In_ ULONG processor_index,
_In_ KE_PROCESSOR_CHANGE_NOTIFY_STATE state,
_Inout_ NTSTATUS& operation_status)
{
KE_PROCESSOR_CHANGE_NOTIFY_CONTEXT change_context = {
.State = state,
.NtNumber = processor_index,
};

registration->callback(registration->callback_context, &change_context, &operation_status);
return operation_status;
}

static _Must_inspect_result_ NTSTATUS
_usersim_notify_processor_change(_In_ ULONG processor_index, _In_ KE_PROCESSOR_CHANGE_NOTIFY_STATE state)
{
if (processor_index >= KeQueryMaximumProcessorCountEx(ALL_PROCESSOR_GROUPS)) {
return STATUS_INVALID_PARAMETER;
}

std::vector<std::shared_ptr<usersim_processor_change_callback_registration_t>> callbacks;
{
std::lock_guard<std::mutex> lock(_usersim_processor_change_callback_lock);
callbacks = _usersim_processor_change_callbacks;
Comment thread
Alan-Jowett marked this conversation as resolved.
}

NTSTATUS operation_status = STATUS_SUCCESS;

for (const auto& registration : callbacks) {
_usersim_invoke_processor_change_callback(registration, processor_index, state, operation_status);
if (!NT_SUCCESS(operation_status)) {
break;
}
}

return operation_status;
}

usersim_result_t
usersim_initialize_irql()
{
Expand Down Expand Up @@ -356,10 +413,169 @@ ULONG
KeQueryMaximumProcessorCountEx(_In_ USHORT group_number) { return GetMaximumProcessorCount(group_number); }

ULONG
KeQueryActiveProcessorCount() { return KeQueryMaximumProcessorCount(); }
KeQueryActiveProcessorCount() { return KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS); }

ULONG
KeQueryActiveProcessorCountEx(_In_ USHORT group_number) { return KeQueryMaximumProcessorCountEx(group_number); }
KeQueryActiveProcessorCountEx(_In_ USHORT group_number)
{
ULONG active_processor_count_override = _usersim_active_processor_count_override.load(std::memory_order_acquire);
if (active_processor_count_override == 0) {
return KeQueryMaximumProcessorCountEx(group_number);
}

return min(active_processor_count_override, KeQueryMaximumProcessorCountEx(group_number));
}

void*
KeRegisterProcessorChangeCallback(
_In_ KE_PROCESSOR_CHANGE_CALLBACK callback, _In_opt_ void* callback_context, _In_ ULONG flags)
{
std::shared_ptr<usersim_processor_change_callback_registration_t> registration(
new (std::nothrow) usersim_processor_change_callback_registration_t{
.callback = callback, .callback_context = callback_context});
if (registration == nullptr) {
return nullptr;
}

{
std::lock_guard<std::mutex> lock(_usersim_processor_change_callback_lock);
_usersim_processor_change_callbacks.push_back(registration);
}

if ((flags & KE_PROCESSOR_CHANGE_ADD_EXISTING) != 0) {
Comment thread
Alan-Jowett marked this conversation as resolved.
ULONG active_processor_count = KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS);
NTSTATUS status = STATUS_SUCCESS;
for (ULONG processor_index = 0; processor_index < active_processor_count; processor_index++) {
status = _usersim_invoke_processor_change_callback(
registration, processor_index, KeProcessorAddStartNotify, status);
if (!NT_SUCCESS(status)) {
KeDeregisterProcessorChangeCallback(registration.get());
return nullptr;
}
}

status = STATUS_SUCCESS;
for (ULONG processor_index = 0; processor_index < active_processor_count; processor_index++) {
status = _usersim_invoke_processor_change_callback(
registration, processor_index, KeProcessorAddCompleteNotify, status);
if (!NT_SUCCESS(status)) {
KeDeregisterProcessorChangeCallback(registration.get());
return nullptr;
}
}
}

return registration.get();
}

void
KeDeregisterProcessorChangeCallback(_In_ void* callback_handle)
{
std::lock_guard<std::mutex> lock(_usersim_processor_change_callback_lock);
auto iterator = std::find_if(
_usersim_processor_change_callbacks.begin(),
_usersim_processor_change_callbacks.end(),
[callback_handle](const std::shared_ptr<usersim_processor_change_callback_registration_t>& registration) {
return registration.get() == callback_handle;
});
if (iterator != _usersim_processor_change_callbacks.end()) {
_usersim_processor_change_callbacks.erase(iterator);
}
}

_Must_inspect_result_ NTSTATUS
usersim_set_active_processor_count(_In_ ULONG active_processor_count)
{
ULONG maximum_processor_count = KeQueryMaximumProcessorCountEx(ALL_PROCESSOR_GROUPS);
if (active_processor_count == 0 || active_processor_count > maximum_processor_count) {
return STATUS_INVALID_PARAMETER;
}

std::lock_guard<std::mutex> lock(_usersim_processor_change_state_lock);
_usersim_pending_processor_add.reset();
_usersim_active_processor_count_override.store(active_processor_count, std::memory_order_release);
return STATUS_SUCCESS;
}

void
usersim_reset_active_processor_count()
{
std::lock_guard<std::mutex> lock(_usersim_processor_change_state_lock);
_usersim_pending_processor_add.reset();
_usersim_active_processor_count_override.store(0, std::memory_order_release);
}

_Must_inspect_result_ NTSTATUS
usersim_notify_processor_add_start(_In_ ULONG processor_index)
Comment thread
Alan-Jowett marked this conversation as resolved.
{
{
std::lock_guard<std::mutex> lock(_usersim_processor_change_state_lock);
ULONG active_processor_count = KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS);
if (_usersim_pending_processor_add.has_value() || processor_index != active_processor_count) {
return STATUS_INVALID_PARAMETER;
}

_usersim_pending_processor_add = processor_index;
}

NTSTATUS status = _usersim_notify_processor_change(processor_index, KeProcessorAddStartNotify);
if (!NT_SUCCESS(status)) {
std::lock_guard<std::mutex> lock(_usersim_processor_change_state_lock);
if (_usersim_pending_processor_add == processor_index) {
_usersim_pending_processor_add.reset();
}
}

return status;
}

_Must_inspect_result_ NTSTATUS
usersim_notify_processor_add_complete(_In_ ULONG processor_index)
{
{
std::lock_guard<std::mutex> lock(_usersim_processor_change_state_lock);
ULONG active_processor_count = KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS);
if (!_usersim_pending_processor_add.has_value() || *_usersim_pending_processor_add != processor_index ||
processor_index != active_processor_count) {
return STATUS_INVALID_PARAMETER;
}
}

NTSTATUS status = _usersim_notify_processor_change(processor_index, KeProcessorAddCompleteNotify);

std::lock_guard<std::mutex> lock(_usersim_processor_change_state_lock);
if (_usersim_pending_processor_add == processor_index) {
_usersim_pending_processor_add.reset();
}
if (NT_SUCCESS(status)) {
ULONG active_processor_count_override =
_usersim_active_processor_count_override.load(std::memory_order_relaxed);
_usersim_active_processor_count_override.store(
max(active_processor_count_override, processor_index + 1), std::memory_order_release);
}

return status;
}

_Must_inspect_result_ NTSTATUS
usersim_notify_processor_add_failure(_In_ ULONG processor_index)
{
{
std::lock_guard<std::mutex> lock(_usersim_processor_change_state_lock);
if (!_usersim_pending_processor_add.has_value() || *_usersim_pending_processor_add != processor_index) {
return STATUS_INVALID_PARAMETER;
}
}

NTSTATUS status = _usersim_notify_processor_change(processor_index, KeProcessorAddFailureNotify);

std::lock_guard<std::mutex> lock(_usersim_processor_change_state_lock);
if (_usersim_pending_processor_add == processor_index) {
_usersim_pending_processor_add.reset();
}

return status;
}

KAFFINITY
KeSetSystemAffinityThreadEx(KAFFINITY affinity)
Expand Down Expand Up @@ -1093,4 +1309,4 @@ KeExpandKernelStackAndCalloutEx(
return STATUS_SUCCESS;
}

#pragma endregion events
#pragma endregion events
Loading