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);