[Zion] Add a framework for better process exit.

zion/memory/kernel_vmm.h

@@ -16,6 +16,7 @@ class KernelVmm {
 
   static uint64_t AcquireKernelStack();
 
+  // Takes the base address to the stack. I.e. the highest value in it.
   static void FreeKernelStack(uint64_t);
 
  private:

zion/object/process.cpp

@@ -5,6 +5,7 @@
 #include "memory/paging_util.h"
 #include "memory/physical_memory.h"
 #include "object/thread.h"
+#include "scheduler/process_manager.h"
 #include "scheduler/scheduler.h"
 
 namespace {
@@ -49,11 +50,11 @@ glcr::RefPtr<Thread> Process::GetThread(uint64_t tid) {
 void Process::CheckState() {
   MutexHolder lock(mutex_);
   for (uint64_t i = 0; i < threads_.size(); i++) {
-    if (threads_[i]->GetState() != Thread::FINISHED) {
+    if (!threads_[i]->IsDying()) {
       return;
     }
   }
-  state_ = FINISHED;
+  Exit();
 }
 
 glcr::RefPtr<Capability> Process::ReleaseCapability(uint64_t cid) {
@@ -67,3 +68,28 @@ glcr::RefPtr<Capability> Process::GetCapability(uint64_t cid) {
 uint64_t Process::AddExistingCapability(const glcr::RefPtr<Capability>& cap) {
   return caps_.AddExistingCapability(cap);
 }
+
+void Process::Exit() {
+  // TODO: Check this state elsewhere to ensure that we don't for instance
+  // create a running thread on a finished process.
+  state_ = FINISHED;
+
+  for (uint64_t i = 0; i < threads_.size(); i++) {
+    if (!threads_[i]->IsDying()) {
+      threads_[i]->Cleanup();
+    }
+  }
+
+  // From this point onward no threads should be able to reach userspace.
+
+  // TODO: Unmap all userspace mappings.
+  // TODO: Clear capabilities.
+
+  // TODO: In the future consider removing this from the process manager.
+  // I need to think through the implications because the process object
+  // will be kept alive by the process that created it most likely.
+
+  if (gScheduler->CurrentProcess().id_ == id_) {
+    gScheduler->Yield();
+  }
+}

zion/object/process.h

@@ -61,6 +61,8 @@ class Process : public KernelObject {
 
   State GetState() { return state_; }
 
+  void Exit();
+
  private:
   friend class glcr::MakeRefCountedFriend<Process>;
   Process();

zion/object/thread.cpp

@@ -2,6 +2,7 @@
 
 #include "common/gdt.h"
 #include "debug/debug.h"
+#include "memory/kernel_vmm.h"
 #include "memory/paging_util.h"
 #include "object/process.h"
 #include "scheduler/scheduler.h"
@@ -68,14 +69,27 @@ void Thread::Exit() {
 #if K_THREAD_DEBUG
   dbgln("Exiting");
 #endif
-  state_ = FINISHED;
+  auto curr_thread = gScheduler->CurrentThread();
+  if (curr_thread->tid() != id_) {
+    panic("Thread::Exit called from [{}.{}] on [{}.{}]", curr_thread->pid(),
+          curr_thread->tid(), pid(), tid());
+  }
+  Cleanup();
+  gScheduler->Yield();
+}
+
+void Thread::Cleanup() {
+  state_ = CLEANUP;
   process_.CheckState();
   while (blocked_threads_.size() != 0) {
     auto thread = blocked_threads_.PopFront();
     thread->SetState(Thread::RUNNABLE);
     gScheduler->Enqueue(thread);
   }
-  gScheduler->Yield();
+  state_ = FINISHED;
+  // TODO: Race condition when called from exit, once kernel stack manager
+  // actually reuses stacks this will cause an issue
+  KernelVmm::FreeKernelStack(rsp0_start_);
 }
 
 void Thread::Wait() {
@@ -86,7 +100,7 @@ void Thread::Wait() {
   // 3. B finishes.
   // 4. Context Switch B -> A
   // 5. A forever blocks on B.
-  if (state_ == Thread::FINISHED) {
+  if (IsDying()) {
     return;
   }
   auto thread = gScheduler->CurrentThread();

zion/object/thread.h

@@ -29,6 +29,7 @@ class Thread : public KernelObject, public glcr::IntrusiveListNode<Thread> {
     RUNNING,
     RUNNABLE,
     BLOCKED,
+    CLEANUP,
     FINISHED,
   };
   static glcr::RefPtr<Thread> RootThread(Process& root_proc);
@@ -51,8 +52,17 @@ class Thread : public KernelObject, public glcr::IntrusiveListNode<Thread> {
   // State Management.
   State GetState() { return state_; };
   void SetState(State state) { state_ = state; }
+  bool IsDying() { return state_ == CLEANUP || state_ == FINISHED; }
+
+  // Exits this thread.
+  // Allows all blocked threads to run and releases the kernel stack.
+  // This function should only be called by the running thread on itself
+  // as it will yield.
   void Exit();
 
+  // Like Exit except it does not yield.
+  void Cleanup();
+
   void Wait();
 
  private:

zion/scheduler/process_manager.h

@@ -12,6 +12,8 @@ class ProcessManager {
   static void Init();
 
   void InsertProcess(const glcr::RefPtr<Process>& proc);
+  void RemoveProcess(uint64_t id);
+
   Process& FromId(uint64_t id);
 
  private:

zion/scheduler/scheduler.cpp

@@ -43,6 +43,8 @@ void Scheduler::Preempt() {
     return;
   }
 
+  ClearDeadThreadsFromFront();
+
   asm volatile("cli");
   if (current_thread_ == sleep_thread_) {
     // Sleep should never be preempted. (We should yield it if another thread
@@ -66,9 +68,14 @@ void Scheduler::Preempt() {
 
 void Scheduler::Yield() {
   if (!enabled_) {
+    // This is expected to fire once at the start when we enqueue the first
+    // thread before the scheduler is enabled. Maybe we should get rid of it?
     dbgln("WARN Scheduler skipped yield.");
     return;
   }
+
+  ClearDeadThreadsFromFront();
+
   asm volatile("cli");
 
   glcr::RefPtr<Thread> prev = current_thread_;
@@ -78,7 +85,9 @@ void Scheduler::Yield() {
       return;
     } else {
       current_thread_ = runnable_threads_.PopFront();
-      prev->SetState(Thread::RUNNABLE);
+      if (!prev->IsDying()) {
+        prev->SetState(Thread::RUNNABLE);
+      }
     }
   } else {
     if (runnable_threads_.size() == 0) {
@@ -90,3 +99,10 @@ void Scheduler::Yield() {
 
   SwapToCurrent(*prev);
 }
+
+void Scheduler::ClearDeadThreadsFromFront() {
+  while (runnable_threads_.size() > 0 &&
+         runnable_threads_.PeekFront()->IsDying()) {
+    runnable_threads_.PopFront();
+  }
+}

zion/scheduler/scheduler.h

@@ -33,6 +33,8 @@ class Scheduler {
 
   Scheduler();
   void SwapToCurrent(Thread& prev);
+
+  void ClearDeadThreadsFromFront();
 };
 
 extern Scheduler* gScheduler;

zion/syscall/process.cpp

@@ -8,8 +8,7 @@
 z_err_t ProcessExit(ZProcessExitReq* req) {
   auto curr_thread = gScheduler->CurrentThread();
   dbgln("Exit code: {x}", req->code);
-  // FIXME: kill process here.
-  curr_thread->Exit();
+  curr_thread->process().Exit();
   panic("Returned from thread exit");
   return glcr::UNIMPLEMENTED;
 }