Kernel: Always return from Thread::wait_on

We need to always return from Thread::wait_on, even when a thread
is being killed. This is necessary so that the kernel call stack
can clean up and release references held by it. Then, right before
transitioning back to user mode, we check if the thread is
supposed to die, and at that point change the thread state to
Dying to prevent further scheduling of this thread.

This addresses some possible resource leaks similar to #3073

1 files changed, 15 insertions, 0 deletions

diff --git a/Kernel/Scheduler.cpp b/Kernel/Scheduler.cpp
index dcfdd4074c..7d0a477903 100644
--- a/Kernel/Scheduler.cpp
+++ b/Kernel/Scheduler.cpp
@@ -365,6 +365,21 @@ bool Scheduler::pick_next()
 
     ScopedSpinLock lock(g_scheduler_lock);
 
+    if (current_thread->should_die() && current_thread->state() == Thread::Running) {
+        // Rather than immediately killing threads, yanking the kernel stack
+        // away from them (which can lead to e.g. reference leaks), we always
+        // allow Thread::wait_on to return. This allows the kernel stack to
+        // clean up and eventually we'll get here shortly before transitioning
+        // back to user mode (from Processor::exit_trap). At this point we
+        // no longer want to schedule this thread. We can't wait until
+        // Scheduler::enter_current because we don't want to allow it to
+        // transition back to user mode.
+#ifdef SCHEDULER_DEBUG
+        dbg() << "Scheduler[" << Processor::current().id() << "]: Thread " << *current_thread << " is dying";
+#endif
+        current_thread->set_state(Thread::Dying);
+    }
+
     // Check and unblock threads whose wait conditions have been met.
     Scheduler::for_each_nonrunnable([&](Thread& thread) {
         thread.consider_unblock(now_sec, now_usec);