Use native TLS for fiber context across thread migration

src/libs/libKernel.cpp

@@ -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;
+struct FiberThreadContext {
-static thread_local FiberObject*    g_thread_return_fiber = nullptr;
+	FiberObject*    current_fiber;
-static thread_local FiberObject*    g_pending_idle_fiber  = nullptr;
+	FiberObject*    returned_fiber;
-static thread_local FiberCpuContext g_thread_fiber_context {};
+	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.
+static void FiberCompleteSwitch(FiberObject* current) {
-[[gnu::noinline]] static void FiberCompleteSwitch(FiberObject* current) {
+	auto* context = FiberGetThreadContext();
-	g_current_fiber = current;
+	context->current_fiber = current;
 	// Only publish IDLE after switching away from the departing fiber's stack.
-	auto* fiber          = g_pending_idle_fiber;
+	auto* fiber = context->pending_idle_fiber;
-	g_pending_idle_fiber = nullptr;
+	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;
+	auto* context = FiberGetThreadContext();
-	g_current_fiber = nullptr;
+	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_run    = arg_on_run;
-	fiber->arg_on_return  = 0;
+	fiber->arg_on_return = 0;
-	g_thread_return_fiber = nullptr;
+	// 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;
+		context->returned_fiber = fiber;
-		g_pending_idle_fiber = fiber;
+		context->pending_idle_fiber = fiber;
-		FiberRestoreContext(&g_thread_fiber_context, 1);
+		FiberRestoreContext(&context->cpu_context, 1);
 	}
 
 	FiberCompleteSwitch(fiber);

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<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;
+		int worker_repeat_result = -1;
-		Libs::Fiber::FiberObject* worker_self = data.first;
+		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_context_clear = thread_context_is_clear();
-			worker_self_result = Libs::Fiber::FiberGetSelf(&worker_self);
+			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, thread_context_is_clear() && main_self_result == OK && main_self == nullptr,
-		Check(test, Libs::Fiber::FiberGetSelf(&main_self_after) == OK && main_self_after == nullptr &&
-		                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,
+		CheckOk(test, worker_repeat_result, "repeat run on another thread");
-		      "return to thread retained a current fiber");
+		Check(test, worker_context_clear, "worker retained a context outside a fiber run");
-		Check(test, returned == 23 && data.first_visits == 2 && data.second_visits == 2 &&
+		Check(test, worker_returned == 23 && worker_repeat_returned == 33 &&
-		                data.first_arg == 20 && data.second_arg == 21 &&
+		                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");