Use native TLS for fiber context across thread migration

This commit is contained in:
nmzik
2026-09-16 11:22:05 +02:00
parent 9b2111dc01
commit ded6964953
2 changed files with 102 additions and 39 deletions
+72 -28
View File
@@ -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);
+30 -11
View File
@@ -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");