diff --git a/src/libs/libKernel.cpp b/src/libs/libKernel.cpp index 334e8e85b..0c5784cdf 100644 --- a/src/libs/libKernel.cpp +++ b/src/libs/libKernel.cpp @@ -41,9 +41,11 @@ #include #elif defined(__APPLE__) #include +#include #include #else #include +#include #include #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(TlsGetValue(g_fiber_context_key)); +#else + return static_cast(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(fiber->state); @@ -2368,12 +2406,12 @@ static void FiberStoreState(FiberObject* fiber, uint32_t state) { std::atomic_ref(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); diff --git a/tests/VirtualMemoryAllocationTests.cpp b/tests/VirtualMemoryAllocationTests.cpp index 1b5f98f58..663691317 100644 --- a/tests/VirtualMemoryAllocationTests.cpp +++ b/tests/VirtualMemoryAllocationTests.cpp @@ -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(&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 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");