From b58186265e0264fe28ad8758552625095b2a69f2 Mon Sep 17 00:00:00 2001
From: Drew Galbraith <dgalbraith33@gmail.com>
Date: Mon, 29 May 2023 23:09:39 -0700
Subject: [PATCH] Split Yield and Preempt into separate scheduling functions.

This switch makes the logic for each much easier to parse.
---
 zion/interrupt/interrupt.cpp |  2 +-
 zion/scheduler/scheduler.cpp | 76 ++++++++++++++++++++++--------------
 zion/scheduler/scheduler.h   |  7 ++++
 3 files changed, 54 insertions(+), 31 deletions(-)

diff --git a/zion/interrupt/interrupt.cpp b/zion/interrupt/interrupt.cpp
index 88f38cb..509fc13 100644
--- a/zion/interrupt/interrupt.cpp
+++ b/zion/interrupt/interrupt.cpp
@@ -123,7 +123,7 @@ extern "C" void interrupt_timer(InterruptFrame*) {
     dbgln("timer: %us", cnt * 50 / 1000);
   }
   outb(PIC1_COMMAND, PIC_EOI);
-  sched::Yield();
+  sched::Preempt();
 }
 
 void EnablePic() {
diff --git a/zion/scheduler/scheduler.cpp b/zion/scheduler/scheduler.cpp
index c9c10eb..78bd686 100644
--- a/zion/scheduler/scheduler.cpp
+++ b/zion/scheduler/scheduler.cpp
@@ -34,55 +34,70 @@ class Scheduler {
   void InsertProcess(Process* process) { proc_list_.PushBack(process); }
   void Enqueue(Thread* thread) { runnable_threads_.PushBack(thread); }
 
+  void SwapToCurrent(Thread& prev) {
+    if (current_thread_->GetState() != Thread::RUNNABLE) {
+      panic("Swapping to non-runnable thread.");
+    }
+    current_thread_->SetState(Thread::RUNNING);
+
+    context_switch(prev.Rsp0Ptr(), current_thread_->Rsp0Ptr());
+
+    asm volatile("sti");
+  }
+
+  void Preempt() {
+    if (!enabled_) {
+      return;
+    }
+
+    asm volatile("cli");
+    if (current_thread_ == sleep_thread_) {
+      // Sleep should never be preempted. (We should yield it if another thread
+      // becomes scheduleable).
+      return;
+    }
+
+    if (runnable_threads_.size() == 0) {
+      // Continue.
+      return;
+    }
+
+    SharedPtr<Thread> prev = current_thread_;
+    prev->SetState(Thread::RUNNABLE);
+    current_thread_ = runnable_threads_.CycleFront(prev);
+
+    SwapToCurrent(*prev);
+  }
+
   void Yield() {
     if (!enabled_) {
+      dbgln("WARN Scheduler skipped yield.");
       return;
     }
     asm volatile("cli");
 
     SharedPtr<Thread> prev = current_thread_;
-    SharedPtr<Thread> next;
     if (prev == sleep_thread_) {
       if (runnable_threads_.size() == 0) {
-        // Continue sleeping.
+        panic("Sleep thread yielded without next.");
         return;
       } else {
         // FIXME: Memory operation.
-        next = runnable_threads_.PopFront();
+        current_thread_ = runnable_threads_.PopFront();
         prev->SetState(Thread::RUNNABLE);
       }
     } else {
-      // Normal thread running.
-      if (prev->GetState() == Thread::RUNNING) {
-        if (runnable_threads_.size() == 0) {
-          // This thread can continue.
-          return;
-        }
-        prev->SetState(Thread::RUNNABLE);
-        next = runnable_threads_.CycleFront(prev);
+      if (runnable_threads_.size() == 0) {
+        current_thread_ = sleep_thread_;
+        dbgln("Sleeping");
+        DumpProcessStates(proc_list_);
       } else {
-        // Thread blocked/exited.
-        if (runnable_threads_.size() == 0) {
-          next = sleep_thread_;
-          dbgln("Sleeping");
-          DumpProcessStates(proc_list_);
-        } else {
-          // FIXME: Memory operation.
-          next = runnable_threads_.PopFront();
-        }
+        // FIXME: Memory operation.
+        current_thread_ = runnable_threads_.PopFront();
       }
     }
 
-    if (next->GetState() != Thread::RUNNABLE) {
-      panic("Non-runnable thread in the queue");
-    }
-
-    next->SetState(Thread::RUNNING);
-    current_thread_ = next;
-
-    context_switch(prev->Rsp0Ptr(), next->Rsp0Ptr());
-
-    asm volatile("sti");
+    SwapToCurrent(*prev);
   }
 
  private:
@@ -110,6 +125,7 @@ Scheduler& GetScheduler() {
 void InitScheduler() { gScheduler = new Scheduler(); }
 void EnableScheduler() { GetScheduler().Enable(); }
 
+void Preempt() { GetScheduler().Preempt(); }
 void Yield() { GetScheduler().Yield(); }
 
 void InsertProcess(Process* process) { GetScheduler().InsertProcess(process); }
diff --git a/zion/scheduler/scheduler.h b/zion/scheduler/scheduler.h
index e08b2ed..550af20 100644
--- a/zion/scheduler/scheduler.h
+++ b/zion/scheduler/scheduler.h
@@ -12,6 +12,13 @@ void InitScheduler();
 // Enables the scheduler such that processes will yield on ticks.
 void EnableScheduler();
 
+// Preempts the current thread and flags it as runnable in the queue.
+// Generally used by the timer to move to the next timeslice.
+void Preempt();
+
+// Current thread yields and is not rescheduled until some external process
+// adds it.
+// Used when a thread blocks or exits.
 void Yield();
 
 // Scheduler will take ownership