Merge pull request #6041 from bunnei/fiber-leaks

common: fiber: Use weak_ptr when yielding.

5 files changed, 35 insertions, 37 deletions

diff --git a/src/common/fiber.cpp b/src/common/fiber.cpp
index 3c1eefcb7..39532ff58 100644
--- a/src/common/fiber.cpp
+++ b/src/common/fiber.cpp
@@ -116,16 +116,19 @@ void Fiber::Rewind() {
     boost::context::detail::jump_fcontext(impl->rewind_context, this);
 }
 
-void Fiber::YieldTo(std::shared_ptr<Fiber> from, std::shared_ptr<Fiber> to) {
-    ASSERT_MSG(from != nullptr, "Yielding fiber is null!");
-    ASSERT_MSG(to != nullptr, "Next fiber is null!");
-    to->impl->guard.lock();
-    to->impl->previous_fiber = from;
-    auto transfer = boost::context::detail::jump_fcontext(to->impl->context, to.get());
-    ASSERT(from->impl->previous_fiber != nullptr);
-    from->impl->previous_fiber->impl->context = transfer.fctx;
-    from->impl->previous_fiber->impl->guard.unlock();
-    from->impl->previous_fiber.reset();
+void Fiber::YieldTo(std::weak_ptr<Fiber> weak_from, Fiber& to) {
+    to.impl->guard.lock();
+    to.impl->previous_fiber = weak_from.lock();
+
+    auto transfer = boost::context::detail::jump_fcontext(to.impl->context, &to);
+
+    // "from" might no longer be valid if the thread was killed
+    if (auto from = weak_from.lock()) {
+        ASSERT(from->impl->previous_fiber != nullptr);
+        from->impl->previous_fiber->impl->context = transfer.fctx;
+        from->impl->previous_fiber->impl->guard.unlock();
+        from->impl->previous_fiber.reset();
+    }
 }
 
 std::shared_ptr<Fiber> Fiber::ThreadToFiber() {
diff --git a/src/common/fiber.h b/src/common/fiber.h
index f7f587f8c..f2a8ff29a 100644
--- a/src/common/fiber.h
+++ b/src/common/fiber.h
@@ -41,7 +41,7 @@ public:
 
     /// Yields control from Fiber 'from' to Fiber 'to'
     /// Fiber 'from' must be the currently running fiber.
-    static void YieldTo(std::shared_ptr<Fiber> from, std::shared_ptr<Fiber> to);
+    static void YieldTo(std::weak_ptr<Fiber> weak_from, Fiber& to);
     [[nodiscard]] static std::shared_ptr<Fiber> ThreadToFiber();
 
     void SetRewindPoint(std::function<void(void*)>&& rewind_func, void* rewind_param);
diff --git a/src/core/cpu_manager.cpp b/src/core/cpu_manager.cpp
index 8f04fb8f5..bdb374792 100644
--- a/src/core/cpu_manager.cpp
+++ b/src/core/cpu_manager.cpp
@@ -148,7 +148,7 @@ void CpuManager::MultiCoreRunSuspendThread() {
         auto core = kernel.GetCurrentHostThreadID();
         auto& scheduler = *kernel.CurrentScheduler();
         Kernel::KThread* current_thread = scheduler.GetCurrentThread();
-        Common::Fiber::YieldTo(current_thread->GetHostContext(), core_data[core].host_context);
+        Common::Fiber::YieldTo(current_thread->GetHostContext(), *core_data[core].host_context);
         ASSERT(scheduler.ContextSwitchPending());
         ASSERT(core == kernel.GetCurrentHostThreadID());
         scheduler.RescheduleCurrentCore();
@@ -245,7 +245,7 @@ void CpuManager::SingleCoreRunSuspendThread() {
         auto core = kernel.GetCurrentHostThreadID();
         auto& scheduler = *kernel.CurrentScheduler();
         Kernel::KThread* current_thread = scheduler.GetCurrentThread();
-        Common::Fiber::YieldTo(current_thread->GetHostContext(), core_data[0].host_context);
+        Common::Fiber::YieldTo(current_thread->GetHostContext(), *core_data[0].host_context);
         ASSERT(scheduler.ContextSwitchPending());
         ASSERT(core == kernel.GetCurrentHostThreadID());
         scheduler.RescheduleCurrentCore();
@@ -271,7 +271,7 @@ void CpuManager::PreemptSingleCore(bool from_running_enviroment) {
         scheduler.Unload(scheduler.GetCurrentThread());
 
         auto& next_scheduler = kernel.Scheduler(current_core);
-        Common::Fiber::YieldTo(current_thread->GetHostContext(), next_scheduler.ControlContext());
+        Common::Fiber::YieldTo(current_thread->GetHostContext(), *next_scheduler.ControlContext());
     }
 
     // May have changed scheduler
@@ -363,7 +363,7 @@ void CpuManager::RunThread(std::size_t core) {
 
         auto current_thread = system.Kernel().CurrentScheduler()->GetCurrentThread();
         data.is_running = true;
-        Common::Fiber::YieldTo(data.host_context, current_thread->GetHostContext());
+        Common::Fiber::YieldTo(data.host_context, *current_thread->GetHostContext());
         data.is_running = false;
         data.is_paused = true;
         data.exit_barrier->Wait();
diff --git a/src/core/hle/kernel/k_scheduler.cpp b/src/core/hle/kernel/k_scheduler.cpp
index 6e89c3042..e7de48476 100644
--- a/src/core/hle/kernel/k_scheduler.cpp
+++ b/src/core/hle/kernel/k_scheduler.cpp
@@ -734,7 +734,7 @@ void KScheduler::ScheduleImpl() {
     }
     guard.unlock();
 
-    Common::Fiber::YieldTo(*old_context, switch_fiber);
+    Common::Fiber::YieldTo(*old_context, *switch_fiber);
     /// When a thread wakes up, the scheduler may have changed to other in another core.
     auto& next_scheduler = *system.Kernel().CurrentScheduler();
     next_scheduler.SwitchContextStep2();
@@ -769,13 +769,8 @@ void KScheduler::SwitchToCurrent() {
                     break;
                 }
             }
-            std::shared_ptr<Common::Fiber>* next_context;
-            if (next_thread != nullptr) {
-                next_context = &next_thread->GetHostContext();
-            } else {
-                next_context = &idle_thread->GetHostContext();
-            }
-            Common::Fiber::YieldTo(switch_fiber, *next_context);
+            auto thread = next_thread ? next_thread : idle_thread;
+            Common::Fiber::YieldTo(switch_fiber, *thread->GetHostContext());
         } while (!is_switch_pending());
     }
 }
diff --git a/src/tests/common/fibers.cpp b/src/tests/common/fibers.cpp
index d94492fc6..751cbe196 100644
--- a/src/tests/common/fibers.cpp
+++ b/src/tests/common/fibers.cpp
@@ -67,7 +67,7 @@ void TestControl1::DoWork() {
         value++;
     }
     results[id] = value;
-    Fiber::YieldTo(work_fibers[id], thread_fibers[id]);
+    Fiber::YieldTo(work_fibers[id], *thread_fibers[id]);
 }
 
 void TestControl1::ExecuteThread(u32 id) {
@@ -76,7 +76,7 @@ void TestControl1::ExecuteThread(u32 id) {
     thread_fibers[id] = thread_fiber;
     work_fibers[id] = std::make_shared<Fiber>(std::function<void(void*)>{WorkControl1}, this);
     items[id] = rand() % 256;
-    Fiber::YieldTo(thread_fibers[id], work_fibers[id]);
+    Fiber::YieldTo(thread_fibers[id], *work_fibers[id]);
     thread_fibers[id]->Exit();
 }
 
@@ -117,11 +117,11 @@ public:
         for (u32 i = 0; i < 12000; i++) {
             value1 += i;
         }
-        Fiber::YieldTo(fiber1, fiber3);
+        Fiber::YieldTo(fiber1, *fiber3);
         const u32 id = thread_ids.Get();
         assert1 = id == 1;
         value2 += 5000;
-        Fiber::YieldTo(fiber1, thread_fibers[id]);
+        Fiber::YieldTo(fiber1, *thread_fibers[id]);
     }
 
     void DoWork2() {
@@ -129,7 +129,7 @@ public:
             ;
         value2 = 2000;
         trap = false;
-        Fiber::YieldTo(fiber2, fiber1);
+        Fiber::YieldTo(fiber2, *fiber1);
         assert3 = false;
     }
 
@@ -137,19 +137,19 @@ public:
         const u32 id = thread_ids.Get();
         assert2 = id == 0;
         value1 += 1000;
-        Fiber::YieldTo(fiber3, thread_fibers[id]);
+        Fiber::YieldTo(fiber3, *thread_fibers[id]);
     }
 
     void ExecuteThread(u32 id);
 
     void CallFiber1() {
         const u32 id = thread_ids.Get();
-        Fiber::YieldTo(thread_fibers[id], fiber1);
+        Fiber::YieldTo(thread_fibers[id], *fiber1);
     }
 
     void CallFiber2() {
         const u32 id = thread_ids.Get();
-        Fiber::YieldTo(thread_fibers[id], fiber2);
+        Fiber::YieldTo(thread_fibers[id], *fiber2);
     }
 
     void Exit();
@@ -241,23 +241,23 @@ public:
 
     void DoWork1() {
         value1 += 1;
-        Fiber::YieldTo(fiber1, fiber2);
+        Fiber::YieldTo(fiber1, *fiber2);
         const u32 id = thread_ids.Get();
         value3 += 1;
-        Fiber::YieldTo(fiber1, thread_fibers[id]);
+        Fiber::YieldTo(fiber1, *thread_fibers[id]);
     }
 
     void DoWork2() {
         value2 += 1;
         const u32 id = thread_ids.Get();
-        Fiber::YieldTo(fiber2, thread_fibers[id]);
+        Fiber::YieldTo(fiber2, *thread_fibers[id]);
     }
 
     void ExecuteThread(u32 id);
 
     void CallFiber1() {
         const u32 id = thread_ids.Get();
-        Fiber::YieldTo(thread_fibers[id], fiber1);
+        Fiber::YieldTo(thread_fibers[id], *fiber1);
     }
 
     void Exit();
@@ -332,7 +332,7 @@ public:
 
     void Execute() {
         thread_fiber = Fiber::ThreadToFiber();
-        Fiber::YieldTo(thread_fiber, fiber1);
+        Fiber::YieldTo(thread_fiber, *fiber1);
         thread_fiber->Exit();
     }
 
@@ -340,7 +340,7 @@ public:
         fiber1->SetRewindPoint(std::function<void(void*)>{WorkControl4}, this);
         if (rewinded) {
             goal_reached = true;
-            Fiber::YieldTo(fiber1, thread_fiber);
+            Fiber::YieldTo(fiber1, *thread_fiber);
         }
         rewinded = true;
         fiber1->Rewind();