mirror of
https://github.com/KytyPS5/KytyPS5.git
synced 2026-10-19 18:36:53 +02:00
Use native TLS for fiber context across thread migration
This commit is contained in:
+72
-28
@@ -41,9 +41,11 @@
|
||||
#include <windows.h>
|
||||
#elif defined(__APPLE__)
|
||||
#include <csignal>
|
||||
#include <pthread.h>
|
||||
#include <sys/ucontext.h>
|
||||
#else
|
||||
#include <csignal>
|
||||
#include <pthread.h>
|
||||
#include <ucontext.h>
|
||||
#endif
|
||||
|
||||
@@ -2354,10 +2356,46 @@ static bool FiberIsValid(const FiberObject* fiber) {
|
||||
fiber->magic_end == FIBER_MAGIC_END;
|
||||
}
|
||||
|
||||
static thread_local FiberObject* g_current_fiber = nullptr;
|
||||
static thread_local FiberObject* g_thread_return_fiber = nullptr;
|
||||
static thread_local FiberObject* g_pending_idle_fiber = nullptr;
|
||||
static thread_local FiberCpuContext g_thread_fiber_context {};
|
||||
struct FiberThreadContext {
|
||||
FiberObject* current_fiber;
|
||||
FiberObject* returned_fiber;
|
||||
FiberObject* pending_idle_fiber;
|
||||
FiberCpuContext cpu_context;
|
||||
};
|
||||
|
||||
static const auto g_fiber_context_key = [] {
|
||||
#if KYTY_PLATFORM == KYTY_PLATFORM_WINDOWS
|
||||
const auto key = TlsAlloc();
|
||||
if (key == TLS_OUT_OF_INDEXES) {
|
||||
std::abort();
|
||||
}
|
||||
#else
|
||||
pthread_key_t key;
|
||||
if (pthread_key_create(&key, nullptr) != 0) {
|
||||
std::abort();
|
||||
}
|
||||
#endif
|
||||
return key;
|
||||
}();
|
||||
|
||||
// Native lookup must stay fresh when a fiber resumes on another host thread.
|
||||
static FiberThreadContext* FiberGetThreadContext() {
|
||||
#if KYTY_PLATFORM == KYTY_PLATFORM_WINDOWS
|
||||
return static_cast<FiberThreadContext*>(TlsGetValue(g_fiber_context_key));
|
||||
#else
|
||||
return static_cast<FiberThreadContext*>(pthread_getspecific(g_fiber_context_key));
|
||||
#endif
|
||||
}
|
||||
|
||||
static void FiberSetThreadContext(FiberThreadContext* context) {
|
||||
#if KYTY_PLATFORM == KYTY_PLATFORM_WINDOWS
|
||||
if (TlsSetValue(g_fiber_context_key, context) == 0) {
|
||||
#else
|
||||
if (pthread_setspecific(g_fiber_context_key, context) != 0) {
|
||||
#endif
|
||||
std::abort();
|
||||
}
|
||||
}
|
||||
|
||||
static uint32_t FiberLoadState(const FiberObject* fiber) {
|
||||
auto& state = const_cast<uint32_t&>(fiber->state);
|
||||
@@ -2368,12 +2406,12 @@ static void FiberStoreState(FiberObject* fiber, uint32_t state) {
|
||||
std::atomic_ref<uint32_t>(fiber->state).store(state, std::memory_order_release);
|
||||
}
|
||||
|
||||
// A fiber may resume on another host thread. Do not cache TLS across the switch.
|
||||
[[gnu::noinline]] static void FiberCompleteSwitch(FiberObject* current) {
|
||||
g_current_fiber = current;
|
||||
static void FiberCompleteSwitch(FiberObject* current) {
|
||||
auto* context = FiberGetThreadContext();
|
||||
context->current_fiber = current;
|
||||
// Only publish IDLE after switching away from the departing fiber's stack.
|
||||
auto* fiber = g_pending_idle_fiber;
|
||||
g_pending_idle_fiber = nullptr;
|
||||
auto* fiber = context->pending_idle_fiber;
|
||||
context->pending_idle_fiber = nullptr;
|
||||
if (fiber != nullptr) {
|
||||
FiberStoreState(fiber, FIBER_STATE_IDLE);
|
||||
}
|
||||
@@ -2464,10 +2502,11 @@ static void FiberRestoreContext(FiberCpuContext* ctx, uint64_t ret) {
|
||||
FiberStoreState(fiber, FIBER_STATE_TERMINATED);
|
||||
FiberSetContextValid(fiber, false);
|
||||
fiber->arg_on_return = 0;
|
||||
g_thread_return_fiber = fiber;
|
||||
g_current_fiber = nullptr;
|
||||
auto* context = FiberGetThreadContext();
|
||||
context->returned_fiber = fiber;
|
||||
context->current_fiber = nullptr;
|
||||
|
||||
FiberRestoreContext(&g_thread_fiber_context, 1);
|
||||
FiberRestoreContext(&context->cpu_context, 1);
|
||||
}
|
||||
|
||||
int32_t KYTY_SYSV_ABI FiberInitialize(FiberObject* fiber, const char* name, FiberEntry entry,
|
||||
@@ -2577,31 +2616,33 @@ int32_t KYTY_SYSV_ABI FiberRun(FiberObject* fiber, uint64_t arg_on_run, uint64_t
|
||||
if (!FiberIsValid(fiber)) {
|
||||
return FIBER_ERROR_INVALID;
|
||||
}
|
||||
if (g_current_fiber != nullptr) {
|
||||
if (FiberGetThreadContext() != nullptr) {
|
||||
return FIBER_ERROR_PERMISSION;
|
||||
}
|
||||
if (!FiberCompareExchangeState(fiber, FIBER_STATE_IDLE, FIBER_STATE_RUNNING)) {
|
||||
return FIBER_ERROR_STATE;
|
||||
}
|
||||
|
||||
fiber->arg_on_run = arg_on_run;
|
||||
fiber->arg_on_return = 0;
|
||||
g_thread_return_fiber = nullptr;
|
||||
fiber->arg_on_run = arg_on_run;
|
||||
fiber->arg_on_return = 0;
|
||||
// Register on the host stack; switching fibers only reads the native TLS slot.
|
||||
FiberThreadContext context {};
|
||||
FiberSetThreadContext(&context);
|
||||
|
||||
if (FiberSaveContext(&g_thread_fiber_context) == 0) {
|
||||
if (FiberSaveContext(&context.cpu_context) == 0) {
|
||||
if (fiber->context_valid) {
|
||||
FiberRestoreContext(&fiber->saved_context, 1);
|
||||
}
|
||||
FiberStartOnGuestStack(fiber);
|
||||
}
|
||||
|
||||
auto* returned_fiber = (g_thread_return_fiber != nullptr ? g_thread_return_fiber : fiber);
|
||||
auto* returned_fiber = (context.returned_fiber != nullptr ? context.returned_fiber : fiber);
|
||||
const auto return_value = returned_fiber->arg_on_return;
|
||||
const auto return_code =
|
||||
FiberLoadState(returned_fiber) == FIBER_STATE_TERMINATED ? FIBER_ERROR_STATE : OK;
|
||||
// Copy results before another thread can acquire and resume the fiber.
|
||||
g_thread_return_fiber = nullptr;
|
||||
FiberCompleteSwitch(nullptr);
|
||||
FiberSetThreadContext(nullptr);
|
||||
|
||||
if (arg_on_return != nullptr) {
|
||||
*arg_on_return = return_value;
|
||||
@@ -2617,7 +2658,8 @@ int32_t KYTY_SYSV_ABI FiberSwitch(FiberObject* fiber, uint64_t arg_on_run,
|
||||
if (!FiberIsValid(fiber)) {
|
||||
return FIBER_ERROR_INVALID;
|
||||
}
|
||||
if (g_current_fiber == nullptr) {
|
||||
auto* context = FiberGetThreadContext();
|
||||
if (context == nullptr) {
|
||||
return FIBER_ERROR_PERMISSION;
|
||||
}
|
||||
|
||||
@@ -2625,14 +2667,14 @@ int32_t KYTY_SYSV_ABI FiberSwitch(FiberObject* fiber, uint64_t arg_on_run,
|
||||
return FIBER_ERROR_STATE;
|
||||
}
|
||||
|
||||
auto* caller = g_current_fiber;
|
||||
auto* caller = context->current_fiber;
|
||||
|
||||
fiber->arg_on_run = arg_on_run;
|
||||
fiber->arg_on_return = 0;
|
||||
|
||||
if (FiberSaveContext(&caller->saved_context) == 0) {
|
||||
FiberSetContextValid(caller, true);
|
||||
g_pending_idle_fiber = caller;
|
||||
context->pending_idle_fiber = caller;
|
||||
if (fiber->context_valid) {
|
||||
FiberRestoreContext(&fiber->saved_context, 1);
|
||||
}
|
||||
@@ -2655,7 +2697,8 @@ int32_t KYTY_SYSV_ABI FiberGetSelf(FiberObject** fiber) {
|
||||
return FIBER_ERROR_NULL;
|
||||
}
|
||||
|
||||
*fiber = g_current_fiber;
|
||||
const auto* context = FiberGetThreadContext();
|
||||
*fiber = context != nullptr ? context->current_fiber : nullptr;
|
||||
|
||||
return OK;
|
||||
}
|
||||
@@ -2663,18 +2706,19 @@ int32_t KYTY_SYSV_ABI FiberGetSelf(FiberObject** fiber) {
|
||||
int32_t KYTY_SYSV_ABI FiberReturnToThread(uint64_t arg_on_return, uint64_t* arg_on_run) {
|
||||
PRINT_NAME();
|
||||
|
||||
if (g_current_fiber == nullptr) {
|
||||
auto* context = FiberGetThreadContext();
|
||||
if (context == nullptr) {
|
||||
return FIBER_ERROR_PERMISSION;
|
||||
}
|
||||
|
||||
auto* fiber = g_current_fiber;
|
||||
auto* fiber = context->current_fiber;
|
||||
fiber->arg_on_return = arg_on_return;
|
||||
|
||||
if (FiberSaveContext(&fiber->saved_context) == 0) {
|
||||
FiberSetContextValid(fiber, true);
|
||||
g_thread_return_fiber = fiber;
|
||||
g_pending_idle_fiber = fiber;
|
||||
FiberRestoreContext(&g_thread_fiber_context, 1);
|
||||
context->returned_fiber = fiber;
|
||||
context->pending_idle_fiber = fiber;
|
||||
FiberRestoreContext(&context->cpu_context, 1);
|
||||
}
|
||||
|
||||
FiberCompleteSwitch(fiber);
|
||||
|
||||
@@ -2818,6 +2818,7 @@ void TestPackedReciprocalSquareRoot() {
|
||||
|
||||
#if defined(__x86_64__) || defined(_M_X64)
|
||||
constexpr int32_t FiberErrorState = -2141650938; // SCE_FIBER_ERROR_STATE
|
||||
constexpr int32_t FiberErrorPermission = -2141650939;
|
||||
|
||||
struct FiberRoundTrip {
|
||||
Libs::Fiber::FiberObject* first;
|
||||
@@ -2896,6 +2897,13 @@ void TestSmallFiberStacksAndMigration() {
|
||||
CheckOk(test, Libs::Fiber::FiberInitialize(data.second, "second", SecondFiberEntry,
|
||||
reinterpret_cast<uint64_t>(&data), second_stack.data() + GuardSize, stack_size,
|
||||
nullptr, 0x0a000000), "initialize second fiber");
|
||||
const auto thread_context_is_clear = [&] {
|
||||
Libs::Fiber::FiberObject* self = data.first;
|
||||
return Libs::Fiber::FiberGetSelf(&self) == OK && self == nullptr &&
|
||||
Libs::Fiber::FiberSwitch(data.first, 0, nullptr) == FiberErrorPermission &&
|
||||
Libs::Fiber::FiberReturnToThread(0, nullptr) == FiberErrorPermission;
|
||||
};
|
||||
Check(test, thread_context_is_clear(), "thread retained a fiber context before its first run");
|
||||
uint64_t first_base = 0;
|
||||
uint64_t second_base = 0;
|
||||
std::memcpy(&first_base, first_stack.data() + GuardSize, sizeof(first_base));
|
||||
@@ -2916,15 +2924,21 @@ void TestSmallFiberStacksAndMigration() {
|
||||
Check(test, returned == 13 && data.first_visits == 1 && data.second_visits == 1 &&
|
||||
data.first_arg == 10 && data.second_arg == 11,
|
||||
"switch/return arguments were not preserved");
|
||||
Check(test, thread_context_is_clear(), "run retained its thread context after returning");
|
||||
check_stacks();
|
||||
|
||||
std::atomic<bool> done {false};
|
||||
int worker_result = -1;
|
||||
int worker_self_result = -1;
|
||||
Libs::Fiber::FiberObject* worker_self = data.first;
|
||||
int worker_repeat_result = -1;
|
||||
bool worker_context_clear = false;
|
||||
uint64_t worker_returned = 0;
|
||||
uint64_t worker_repeat_returned = 0;
|
||||
std::thread worker([&] {
|
||||
worker_result = Libs::Fiber::FiberRun(data.first, 20, &returned);
|
||||
worker_self_result = Libs::Fiber::FiberGetSelf(&worker_self);
|
||||
worker_context_clear = thread_context_is_clear();
|
||||
worker_result = Libs::Fiber::FiberRun(data.first, 20, &worker_returned);
|
||||
worker_context_clear = thread_context_is_clear() && worker_context_clear;
|
||||
worker_repeat_result = Libs::Fiber::FiberRun(data.first, 30, &worker_repeat_returned);
|
||||
worker_context_clear = thread_context_is_clear() && worker_context_clear;
|
||||
done.store(true, std::memory_order_release);
|
||||
});
|
||||
const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5);
|
||||
@@ -2939,19 +2953,24 @@ void TestSmallFiberStacksAndMigration() {
|
||||
const int busy_finalize = running ? Libs::Fiber::FiberFinalize(data.first) : 0;
|
||||
data.release.store(true, std::memory_order_release);
|
||||
worker.join();
|
||||
Libs::Fiber::FiberObject* main_self_after = data.first;
|
||||
Check(test, Libs::Fiber::FiberGetSelf(&main_self_after) == OK && main_self_after == nullptr &&
|
||||
main_self_result == OK && main_self == nullptr,
|
||||
Check(test, thread_context_is_clear() && main_self_result == OK && main_self == nullptr,
|
||||
"migrated fiber changed the original thread's current fiber");
|
||||
Check(test, running && busy_run == FiberErrorState && busy_finalize == FiberErrorState,
|
||||
"running fiber was not exclusively owned");
|
||||
CheckOk(test, worker_result, "resume fiber on another thread");
|
||||
Check(test, worker_self_result == OK && worker_self == nullptr,
|
||||
"return to thread retained a current fiber");
|
||||
Check(test, returned == 23 && data.first_visits == 2 && data.second_visits == 2 &&
|
||||
data.first_arg == 20 && data.second_arg == 21 &&
|
||||
CheckOk(test, worker_repeat_result, "repeat run on another thread");
|
||||
Check(test, worker_context_clear, "worker retained a context outside a fiber run");
|
||||
Check(test, worker_returned == 23 && worker_repeat_returned == 33 &&
|
||||
data.first_visits == 3 && data.second_visits == 3 &&
|
||||
data.first_arg == 30 && data.second_arg == 31 &&
|
||||
data.first_errors == 0 && data.second_errors == 0,
|
||||
"migration lost fiber identity, arguments, or resumable contexts");
|
||||
CheckOk(test, Libs::Fiber::FiberRun(data.first, 40, &returned), "migrate fiber back to original thread");
|
||||
Check(test, returned == 43 && data.first_visits == 4 && data.second_visits == 4 &&
|
||||
data.first_arg == 40 && data.second_arg == 41 &&
|
||||
data.first_errors == 0 && data.second_errors == 0,
|
||||
"return migration reused an expired thread context");
|
||||
Check(test, thread_context_is_clear(), "repeated run retained its thread context after returning");
|
||||
check_stacks();
|
||||
CheckOk(test, Libs::Fiber::FiberFinalize(data.first), "finalize first suspended fiber");
|
||||
CheckOk(test, Libs::Fiber::FiberFinalize(data.second), "finalize second suspended fiber");
|
||||
|
||||
Reference in New Issue
Block a user