Simplify fiber ownership and preserve small guest stacks

This commit is contained in:
nmzik 2026-09-16 07:15:07 +02:00
parent 86e5e01403
commit 71f52f9d90
2 changed files with 198 additions and 146 deletions

View File

@ -29,11 +29,9 @@
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <mutex>
#include <random>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>
#if KYTY_PLATFORM == KYTY_PLATFORM_WINDOWS
@ -2356,70 +2354,24 @@ 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_starting_fiber = nullptr;
static thread_local FiberObject* g_pending_idle_fiber = nullptr;
static thread_local FiberCpuContext g_thread_fiber_context {};
static std::mutex g_fiber_owner_mutex;
static std::unordered_map<FiberObject*, uint64_t> g_fiber_owner_thread;
static std::unordered_map<uint64_t, FiberObject*> g_fiber_current_by_thread;
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 {};
static uint32_t FiberLoadState(const FiberObject* fiber) {
auto& state = const_cast<uint32_t&>(fiber->state);
return std::atomic_ref<uint32_t>(state).load(std::memory_order_acquire);
}
static uint64_t FiberCurrentHostThreadId() {
#if KYTY_PLATFORM == KYTY_PLATFORM_WINDOWS
return static_cast<uint64_t>(::GetCurrentThreadId());
#else
return std::hash<std::thread::id> {}(std::this_thread::get_id());
#endif
}
static void FiberSetOwner(FiberObject* fiber) {
std::lock_guard lock(g_fiber_owner_mutex);
g_fiber_owner_thread[fiber] = FiberCurrentHostThreadId();
}
static void FiberClearOwner(FiberObject* fiber) {
std::lock_guard lock(g_fiber_owner_mutex);
g_fiber_owner_thread.erase(fiber);
}
static uint64_t FiberGetOwner(FiberObject* fiber) {
std::lock_guard lock(g_fiber_owner_mutex);
auto it = g_fiber_owner_thread.find(fiber);
return it != g_fiber_owner_thread.end() ? it->second : 0;
}
static void FiberSetCurrentFiber(FiberObject* fiber) {
g_current_fiber = fiber;
std::lock_guard lock(g_fiber_owner_mutex);
const auto thread_id = FiberCurrentHostThreadId();
if (fiber != nullptr) {
g_fiber_current_by_thread[thread_id] = fiber;
} else {
g_fiber_current_by_thread.erase(thread_id);
}
}
static void FiberStoreState(FiberObject* fiber, uint32_t state) {
std::atomic_ref<uint32_t>(fiber->state).store(state, std::memory_order_release);
if (state == FIBER_STATE_RUNNING) {
FiberSetOwner(fiber);
} else {
FiberClearOwner(fiber);
}
}
static void FiberDeferIdle(FiberObject* fiber) {
g_pending_idle_fiber = fiber;
}
static void FiberCommitDeferredIdle() {
// 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;
// Only publish IDLE after switching away from the departing fiber's stack.
auto* fiber = g_pending_idle_fiber;
g_pending_idle_fiber = nullptr;
if (fiber != nullptr) {
@ -2428,60 +2380,9 @@ static void FiberCommitDeferredIdle() {
}
static bool FiberCompareExchangeState(FiberObject* fiber, uint32_t expected, uint32_t desired) {
const bool ok = std::atomic_ref<uint32_t>(fiber->state)
.compare_exchange_strong(expected, desired, std::memory_order_acq_rel,
std::memory_order_acquire);
if (ok) {
if (desired == FIBER_STATE_RUNNING) {
FiberSetOwner(fiber);
} else {
FiberClearOwner(fiber);
}
}
return ok;
}
static bool FiberWaitAndEnterRunning(FiberObject* fiber, uint32_t* observed_state) {
const auto start = std::chrono::steady_clock::now();
uint32_t spin = 0;
auto& state = fiber->state;
auto ref = std::atomic_ref<uint32_t>(state);
for (;;) {
uint32_t expected = FIBER_STATE_IDLE;
if (ref.compare_exchange_strong(expected, FIBER_STATE_RUNNING, std::memory_order_acq_rel,
std::memory_order_acquire)) {
FiberSetOwner(fiber);
return true;
}
if (observed_state != nullptr) {
*observed_state = expected;
}
if (expected != FIBER_STATE_RUNNING) {
return false;
}
if (std::chrono::steady_clock::now() - start >= std::chrono::milliseconds(5)) {
return false;
}
if (spin++ < 64) {
std::this_thread::yield();
} else {
std::this_thread::sleep_for(std::chrono::microseconds(100));
}
}
}
static bool FiberRepairStaleRunningOnThisThread(FiberObject* fiber, uint32_t observed_state) {
if (observed_state != FIBER_STATE_RUNNING || fiber == g_current_fiber) {
return false;
}
const auto owner = FiberGetOwner(fiber);
if (owner != 0) {
return false;
}
return FiberCompareExchangeState(fiber, FIBER_STATE_RUNNING, FIBER_STATE_IDLE);
return std::atomic_ref<uint32_t>(fiber->state)
.compare_exchange_strong(expected, desired, std::memory_order_acq_rel,
std::memory_order_acquire);
}
static void FiberSetContextValid(FiberObject* fiber, bool valid) {
@ -2538,35 +2439,25 @@ static void FiberRestoreContext(FiberCpuContext* ctx, uint64_t ret) {
}
#endif
[[noreturn]] static void FiberStartTrampoline();
[[noreturn]] static KYTY_SYSV_ABI void FiberStartTrampoline(FiberObject* fiber);
[[noreturn]] static void FiberStartOnGuestStack(FiberObject* fiber) {
[[noreturn]] static KYTY_SYSV_ABI void FiberStartOnGuestStack(FiberObject* fiber) {
FiberCpuContext ctx {};
const auto stack_top = reinterpret_cast<uintptr_t>(fiber->addr_context) +
static_cast<uintptr_t>(fiber->size_context);
auto rsp = (stack_top & ~static_cast<uintptr_t>(0x0f));
#if KYTY_PLATFORM == KYTY_PLATFORM_WINDOWS
rsp -= 4u * sizeof(uint64_t);
#endif
rsp -= sizeof(uint64_t);
*reinterpret_cast<uint64_t*>(rsp) = 0;
ctx.rsp = rsp;
ctx.rip = reinterpret_cast<uint64_t>(&FiberStartTrampoline);
ctx.rdi = reinterpret_cast<uint64_t>(fiber);
g_starting_fiber = fiber;
FiberRestoreContext(&ctx, 1);
}
[[noreturn]] static void FiberStartTrampoline() {
auto* fiber = g_starting_fiber;
g_starting_fiber = nullptr;
if (fiber == nullptr) {
EXIT("Fiber start without current fiber\n");
}
FiberCommitDeferredIdle();
FiberSetCurrentFiber(fiber);
[[noreturn]] static KYTY_SYSV_ABI void FiberStartTrampoline(FiberObject* fiber) {
FiberCompleteSwitch(fiber);
fiber->entry(fiber->arg_on_initialize, fiber->arg_on_run);
@ -2574,7 +2465,7 @@ static void FiberRestoreContext(FiberCpuContext* ctx, uint64_t ret) {
FiberSetContextValid(fiber, false);
fiber->arg_on_return = 0;
g_thread_return_fiber = fiber;
FiberSetCurrentFiber(nullptr);
g_current_fiber = nullptr;
FiberRestoreContext(&g_thread_fiber_context, 1);
}
@ -2704,16 +2595,19 @@ int32_t KYTY_SYSV_ABI FiberRun(FiberObject* fiber, uint64_t arg_on_run, uint64_t
FiberStartOnGuestStack(fiber);
}
FiberCommitDeferredIdle();
FiberSetCurrentFiber(nullptr);
auto* returned_fiber = (g_thread_return_fiber != nullptr ? g_thread_return_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);
if (arg_on_return != nullptr) {
*arg_on_return = returned_fiber->arg_on_return;
*arg_on_return = return_value;
}
return (FiberLoadState(returned_fiber) == FIBER_STATE_TERMINATED ? FIBER_ERROR_STATE : OK);
return return_code;
}
int32_t KYTY_SYSV_ABI FiberSwitch(FiberObject* fiber, uint64_t arg_on_run,
@ -2727,14 +2621,7 @@ int32_t KYTY_SYSV_ABI FiberSwitch(FiberObject* fiber, uint64_t arg_on_run,
return FIBER_ERROR_PERMISSION;
}
for (;;) {
uint32_t observed_state = 0;
if (FiberWaitAndEnterRunning(fiber, &observed_state)) {
break;
}
if (FiberRepairStaleRunningOnThisThread(fiber, observed_state)) {
continue;
}
if (!FiberCompareExchangeState(fiber, FIBER_STATE_IDLE, FIBER_STATE_RUNNING)) {
return FIBER_ERROR_STATE;
}
@ -2745,16 +2632,14 @@ int32_t KYTY_SYSV_ABI FiberSwitch(FiberObject* fiber, uint64_t arg_on_run,
if (FiberSaveContext(&caller->saved_context) == 0) {
FiberSetContextValid(caller, true);
FiberDeferIdle(caller);
g_pending_idle_fiber = caller;
if (fiber->context_valid) {
FiberRestoreContext(&fiber->saved_context, 1);
}
FiberStartOnGuestStack(fiber);
}
FiberCommitDeferredIdle();
FiberSetCurrentFiber(caller);
FiberStoreState(caller, FIBER_STATE_RUNNING);
FiberCompleteSwitch(caller);
if (arg_on_return != nullptr) {
*arg_on_return = caller->arg_on_run;
@ -2788,13 +2673,11 @@ int32_t KYTY_SYSV_ABI FiberReturnToThread(uint64_t arg_on_return, uint64_t* arg_
if (FiberSaveContext(&fiber->saved_context) == 0) {
FiberSetContextValid(fiber, true);
g_thread_return_fiber = fiber;
FiberDeferIdle(fiber);
g_pending_idle_fiber = fiber;
FiberRestoreContext(&g_thread_fiber_context, 1);
}
FiberCommitDeferredIdle();
FiberSetCurrentFiber(fiber);
FiberStoreState(fiber, FIBER_STATE_RUNNING);
FiberCompleteSwitch(fiber);
if (arg_on_run != nullptr) {
*arg_on_run = fiber->arg_on_run;
}

View File

@ -14,7 +14,9 @@
#include <algorithm>
#include <array>
#include <atomic>
#include <bit>
#include <chrono>
#include <cmath>
#include <cinttypes>
#include <cstdint>
@ -22,6 +24,7 @@
#include <cstring>
#include <filesystem>
#include <string>
#include <thread>
#include <vector>
#if defined(__linux__) || KYTY_PLATFORM == KYTY_PLATFORM_WINDOWS
@ -46,6 +49,19 @@
#endif
#endif
namespace Libs::Fiber {
struct FiberObject;
struct FiberOptParam;
using FiberEntry = KYTY_SYSV_ABI void (*)(uint64_t, uint64_t);
int32_t KYTY_SYSV_ABI FiberInitialize(FiberObject*, const char*, FiberEntry, uint64_t, void*,
uint64_t, const FiberOptParam*, uint32_t);
int32_t KYTY_SYSV_ABI FiberFinalize(FiberObject*);
int32_t KYTY_SYSV_ABI FiberRun(FiberObject*, uint64_t, uint64_t*);
int32_t KYTY_SYSV_ABI FiberSwitch(FiberObject*, uint64_t, uint64_t*);
int32_t KYTY_SYSV_ABI FiberGetSelf(FiberObject**);
int32_t KYTY_SYSV_ABI FiberReturnToThread(uint64_t, uint64_t*);
} // namespace Libs::Fiber
namespace {
using Libs::LibKernel::Memory::VirtualQueryInfo;
@ -2800,10 +2816,160 @@ void TestPackedReciprocalSquareRoot() {
}
#endif
#if defined(__x86_64__) || defined(_M_X64)
constexpr int32_t FiberErrorState = -2141650938; // SCE_FIBER_ERROR_STATE
struct FiberRoundTrip {
Libs::Fiber::FiberObject* first;
Libs::Fiber::FiberObject* second;
std::atomic<bool> running {false};
std::atomic<bool> release {false};
int first_errors = 0;
int second_errors = 0;
uint32_t first_visits = 0;
uint32_t second_visits = 0;
uint64_t first_arg = 0;
uint64_t second_arg = 0;
};
[[noreturn]] void KYTY_SYSV_ABI FirstFiberEntry(uint64_t initial, uint64_t arg) {
auto& data = *reinterpret_cast<FiberRoundTrip*>(initial);
for (;;) {
Libs::Fiber::FiberObject* self = nullptr;
data.first_errors |= Libs::Fiber::FiberGetSelf(&self);
data.first_errors |= self != data.first;
data.first_errors |= Libs::Fiber::FiberSwitch(self, 0, nullptr) != FiberErrorState;
if (data.first_errors != 0) {
__builtin_trap();
}
++data.first_visits;
data.first_arg = arg;
if (arg == 20) {
data.running.store(true, std::memory_order_release);
while (!data.release.load(std::memory_order_acquire)) {
asm volatile("pause");
}
}
data.first_errors |= Libs::Fiber::FiberSwitch(data.second, arg + 1, &arg);
if (data.first_errors != 0) {
__builtin_trap();
}
data.first_errors |= Libs::Fiber::FiberReturnToThread(arg + 1, &arg);
}
}
[[noreturn]] void KYTY_SYSV_ABI SecondFiberEntry(uint64_t initial, uint64_t arg) {
auto& data = *reinterpret_cast<FiberRoundTrip*>(initial);
for (;;) {
Libs::Fiber::FiberObject* self = nullptr;
data.second_errors |= Libs::Fiber::FiberGetSelf(&self);
data.second_errors |= self != data.second;
if (data.second_errors != 0) {
__builtin_trap();
}
++data.second_visits;
data.second_arg = arg;
data.second_errors |= Libs::Fiber::FiberSwitch(data.first, arg + 1, &arg);
if (data.second_errors != 0) {
__builtin_trap();
}
}
}
void TestSmallFiberStacksAndMigration() {
const char* test = "SmallFiberStacksAndMigration";
// 256-byte objects, 8-byte object alignment, 16-byte context
// alignment, and a 512-byte minimum context. The game supplies 2048 bytes.
for (const size_t stack_size: {512u, 2048u}) {
alignas(8) std::array<uint8_t, 256> first_object {};
alignas(8) std::array<uint8_t, 256> second_object {};
constexpr size_t GuardSize = 4096;
alignas(16) std::array<uint8_t, GuardSize + 2048 + 64> first_stack;
alignas(16) std::array<uint8_t, GuardSize + 2048 + 64> second_stack;
first_stack.fill(0xa5);
second_stack.fill(0xa5);
FiberRoundTrip data {reinterpret_cast<Libs::Fiber::FiberObject*>(first_object.data()),
reinterpret_cast<Libs::Fiber::FiberObject*>(second_object.data())};
CheckOk(test, Libs::Fiber::FiberInitialize(data.first, "first", FirstFiberEntry,
reinterpret_cast<uint64_t>(&data), first_stack.data() + GuardSize, stack_size,
nullptr, 0x0a000000), "initialize first fiber");
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");
uint64_t first_base = 0;
uint64_t second_base = 0;
std::memcpy(&first_base, first_stack.data() + GuardSize, sizeof(first_base));
std::memcpy(&second_base, second_stack.data() + GuardSize, sizeof(second_base));
const auto check_stacks = [&] {
for (const auto* stack: {&first_stack, &second_stack}) {
const auto is_guard = [](uint8_t byte) { return byte == 0xa5; };
Check(test, std::all_of(stack->begin(), stack->begin() + GuardSize, is_guard) &&
std::all_of(stack->begin() + GuardSize + stack_size, stack->end(), is_guard),
"fiber wrote outside its supplied stack");
}
Check(test, std::memcmp(&first_base, first_stack.data() + GuardSize, sizeof(first_base)) == 0 &&
std::memcmp(&second_base, second_stack.data() + GuardSize, sizeof(second_base)) == 0,
"fiber overwrote the bottom of its stack");
};
uint64_t returned = 0;
CheckOk(test, Libs::Fiber::FiberRun(data.first, 10, &returned), "run first fiber");
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_stacks();
std::atomic<bool> done {false};
int worker_result = -1;
int worker_self_result = -1;
Libs::Fiber::FiberObject* worker_self = data.first;
std::thread worker([&] {
worker_result = Libs::Fiber::FiberRun(data.first, 20, &returned);
worker_self_result = Libs::Fiber::FiberGetSelf(&worker_self);
done.store(true, std::memory_order_release);
});
const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5);
while (!data.running.load(std::memory_order_acquire) &&
!done.load(std::memory_order_acquire) && std::chrono::steady_clock::now() < deadline) {
std::this_thread::yield();
}
const bool running = data.running.load(std::memory_order_acquire);
Libs::Fiber::FiberObject* main_self = data.first;
const int main_self_result = Libs::Fiber::FiberGetSelf(&main_self);
const int busy_run = running ? Libs::Fiber::FiberRun(data.first, 30, nullptr) : 0;
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,
"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 &&
data.first_errors == 0 && data.second_errors == 0,
"migration lost fiber identity, arguments, or resumable contexts");
check_stacks();
CheckOk(test, Libs::Fiber::FiberFinalize(data.first), "finalize first suspended fiber");
CheckOk(test, Libs::Fiber::FiberFinalize(data.second), "finalize second suspended fiber");
std::printf("[host] %s stack=%zu ok\n", test, stack_size);
}
}
#endif
} // namespace
int main(int argc, char** argv) {
InitSubsystems();
#if defined(__x86_64__) || defined(_M_X64)
if (argc == 2 && std::strcmp(argv[1], "--fiber-only") == 0) {
RunTest(TestSmallFiberStacksAndMigration);
return g_failed_tests == 0 ? 0 : 1;
}
#endif
#if defined(__linux__) || KYTY_PLATFORM == KYTY_PLATFORM_WINDOWS
if (argc == 2 && std::strcmp(argv[1], "--rsqrt-only") == 0) {
RunTest(TestPackedReciprocalSquareRoot);
@ -2815,6 +2981,9 @@ int main(int argc, char** argv) {
return g_failed_tests == 0 ? 0 : 1;
}
#if defined(__x86_64__) || defined(_M_X64)
RunTest(TestSmallFiberStacksAndMigration);
#endif
#if defined(__linux__) || KYTY_PLATFORM == KYTY_PLATFORM_WINDOWS
RunTest(TestPackedReciprocalSquareRoot);
#endif