Split Yield and Preempt into separate scheduling functions.

This switch makes the logic for each much easier to parse.

zion/CMakeLists.txt

@@ -5,6 +5,7 @@ add_executable(zion
   debug/debug.cpp
   interrupt/interrupt.cpp
   interrupt/interrupt_enter.s
+  interrupt/timer.cpp
   loader/elf_loader.cpp
   loader/init_loader.cpp
   memory/kernel_heap.cpp

zion/interrupt/interrupt.cpp

@@ -2,7 +2,9 @@
 
 #include <stdint.h>
 
+#include "common/port.h"
 #include "debug/debug.h"
+#include "scheduler/scheduler.h"
 
 #define IDT_INTERRUPT_GATE 0x8E
 
@@ -66,7 +68,21 @@ extern "C" void isr_divide_by_zero();
 extern "C" void interrupt_divide_by_zero(void* frame) { panic("DIV0"); }
 
 extern "C" void isr_protection_fault();
-extern "C" void interrupt_protection_fault(void* frame) { panic("GP"); }
+extern "C" void interrupt_protection_fault(InterruptFrame* frame) {
+  dbgln("General Protection Fault");
+  uint64_t err = frame->error_code;
+  if (err & 0x1) {
+    dbgln("External Source");
+  }
+  if (err & 0x2) {
+    dbgln("IDT");
+  } else {
+    dbgln("GDT");
+  }
+  dbgln("Index: %u", err >> 3);
+
+  panic("GP");
+}
 
 extern "C" void isr_page_fault();
 extern "C" void interrupt_page_fault(InterruptFrame* frame) {
@@ -95,13 +111,41 @@ extern "C" void interrupt_page_fault(InterruptFrame* frame) {
   panic("PF");
 }
 
+#define PIC1_COMMAND 0x20
+#define PIC1_DATA 0x21
+#define PIC_EOI 0x20
+
+uint64_t cnt = 0;
+extern "C" void isr_timer();
+extern "C" void interrupt_timer(InterruptFrame*) {
+  cnt++;
+  if (cnt % 20 == 0) {
+    dbgln("timer: %us", cnt * 50 / 1000);
+  }
+  outb(PIC1_COMMAND, PIC_EOI);
+  sched::Preempt();
+}
+
+void EnablePic() {
+  outb(PIC1_COMMAND, 0x11);
+  outb(PIC1_DATA, 0x20);  // PIC1 offset.
+  outb(PIC1_DATA, 0x4);
+  outb(PIC1_DATA, 0x1);
+
+  // Mask all except the timer.
+  outb(PIC1_DATA, 0xE);
+}
+
 void InitIdt() {
   gIdt[0] = CreateDescriptor(isr_divide_by_zero);
   gIdt[13] = CreateDescriptor(isr_protection_fault);
   gIdt[14] = CreateDescriptor(isr_page_fault);
+  gIdt[32] = CreateDescriptor(isr_timer);
   InterruptDescriptorTablePointer idtp{
       .size = sizeof(gIdt),
       .base = reinterpret_cast<uint64_t>(gIdt),
   };
   asm volatile("lidt %0" ::"m"(idtp));
+
+  EnablePic();
 }

zion/interrupt/interrupt_enter.s

@@ -55,3 +55,5 @@ isr_\name:
 isr_handler divide_by_zero
 isr_handler protection_fault,1
 isr_handler page_fault,1
+
+isr_handler timer

zion/interrupt/timer.cpp

@@ -0,0 +1,41 @@
+#include "interrupt/timer.h"
+
+#include "common/port.h"
+
+namespace {
+
+// IO Ports
+constexpr uint8_t kPit0Ctl = 0x40;
+constexpr uint8_t kPit1Ctl = 0x41;
+constexpr uint8_t kPit2Ctl = 0x42;
+constexpr uint8_t kPitCmd = 0x43;
+
+// Command Register
+// Mode
+constexpr uint8_t kPitMode0 = 0x00;
+constexpr uint8_t kPitMode1 = 0x02;
+constexpr uint8_t kPitMode2 = 0x04;
+constexpr uint8_t kPitMode3 = 0x06;
+constexpr uint8_t kPitMode4 = 0x08;
+constexpr uint8_t kPitMode5 = 0x0A;
+// R/W
+constexpr uint8_t kCmdLatch = 0x00;
+constexpr uint8_t kCmdRwLow = 0x10;
+constexpr uint8_t kCmdRwHi = 0x20;
+constexpr uint8_t kCmdRwBoth = 0x30;
+// PIT Select
+constexpr uint8_t kSelect0 = 0x00;
+constexpr uint8_t kSelect1 = 0x40;
+constexpr uint8_t kSelect2 = 0x80;
+constexpr uint8_t kReadback = 0xC0;
+
+constexpr uint32_t kPitFrequency = 1193182;
+
+}  // namespace
+
+void SetFrequency(uint64_t hertz) {
+  uint16_t reload = kPitFrequency / hertz;
+  outb(kPitCmd, kPitMode3 | kCmdRwBoth | kSelect0);
+  outb(kPit0Ctl, reload & 0xFF);
+  outb(kPit0Ctl, reload >> 8);
+}

zion/interrupt/timer.h

@@ -0,0 +1,5 @@
+#pragma once
+
+#include <stdint.h>
+
+void SetFrequency(uint64_t hertz);

zion/lib/linked_list.h

@@ -44,16 +44,22 @@ class LinkedList {
     return ret;
   }
 
-  T CycleFront() {
+  /*
+   * Returns the front item in the list and pushes the passed item to the back.
+   *
+   * Done in one function to avoid a memory alloc/dealloc during scheduling.
+   **/
+  T CycleFront(const T& new_item) {
     if (size_ == 0 || front_ == nullptr) {
       panic("Cycling empty list");
     }
 
+    T ret = front_->item;
+    front_->item = new_item;
     if (size_ == 1) {
-      return front_->item;
+      return ret;
     }
 
-    T ret = front_->item;
     ListItem* old_front = front_;
     ListItem* iter = front_;
     front_ = front_->next;

zion/scheduler/scheduler.cpp

@@ -22,62 +22,90 @@ class Scheduler {
   Scheduler() {
     SharedPtr<Process> root = Process::RootProcess();
     sleep_thread_ = root->GetThread(0);
-    runnable_threads_.PushBack(sleep_thread_);
-    proc_list_.PushBack(Process::RootProcess());
+    // TODO: Implement a separate sleep thread?
+    current_thread_ = sleep_thread_;
+    proc_list_.PushBack(root);
   }
   void Enable() { enabled_ = true; }
 
   Process& CurrentProcess() { return CurrentThread().process(); }
-  Thread& CurrentThread() { return *runnable_threads_.PeekFront(); }
+  Thread& CurrentThread() { return *current_thread_; }
 
   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;
-    if (CurrentThread().GetState() == Thread::RUNNING) {
-      prev = runnable_threads_.CycleFront();
-      prev->SetState(Thread::RUNNABLE);
+    SharedPtr<Thread> prev = current_thread_;
+    if (prev == sleep_thread_) {
+      if (runnable_threads_.size() == 0) {
+        panic("Sleep thread yielded without next.");
+        return;
+      } else {
+        // FIXME: Memory operation.
+        current_thread_ = runnable_threads_.PopFront();
+        prev->SetState(Thread::RUNNABLE);
+      }
     } else {
-      // This technically is a memory operation but should only occur when a
-      // thread is blocking so may be ok?
-      prev = runnable_threads_.PopFront();
+      if (runnable_threads_.size() == 0) {
+        current_thread_ = sleep_thread_;
+        dbgln("Sleeping");
+        DumpProcessStates(proc_list_);
+      } else {
+        // FIXME: Memory operation.
+        current_thread_ = runnable_threads_.PopFront();
+      }
     }
 
-    SharedPtr<Thread> next;
-    if (runnable_threads_.size() == 0) {
-      next = sleep_thread_;
-      DumpProcessStates(proc_list_);
-    } else {
-      next = runnable_threads_.PeekFront();
-    }
-
-    if (next->GetState() != Thread::RUNNABLE) {
-      panic("Non-runnable thread in the queue");
-    }
-    // Needs to be before the next == prev check
-    // otherwise the active thread will be RUNNABLE instead of RUNNING.
-    next->SetState(Thread::RUNNING);
-
-    if (next == prev) {
-      dbgln("No next thread, continue");
-      return;
-    }
-
-    context_switch(prev->Rsp0Ptr(), next->Rsp0Ptr());
-
-    asm volatile("sti");
+    SwapToCurrent(*prev);
   }
 
  private:
   bool enabled_ = false;
   // TODO: move this to a separate process manager class.
   LinkedList<SharedPtr<Process>> proc_list_;
+
+  SharedPtr<Thread> current_thread_;
   LinkedList<SharedPtr<Thread>> runnable_threads_;
 
   SharedPtr<Thread> sleep_thread_;
@@ -97,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); }

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

zion/zion.cpp

@@ -3,6 +3,7 @@
 #include "common/gdt.h"
 #include "debug/debug.h"
 #include "interrupt/interrupt.h"
+#include "interrupt/timer.h"
 #include "loader/init_loader.h"
 #include "memory/kernel_heap.h"
 #include "memory/paging_util.h"
@@ -11,20 +12,31 @@
 #include "syscall/syscall.h"
 
 extern "C" void zion() {
+  dbgln("[boot] Init GDT & IDT.");
   InitGdt();
   InitIdt();
+
+  dbgln("[boot] Init Paging.");
   InitPaging();
 
+  dbgln("[boot] Init Physical Memory Manager.");
   phys_mem::InitBootstrapPageAllocation();
   KernelHeap heap(0xFFFFFFFF'40000000, 0xFFFFFFFF'80000000);
   phys_mem::InitPhysicalMemoryManager();
 
+  dbgln("[boot] Init syscalls.");
   InitSyscall();
 
+  dbgln("[boot] Init scheduler.");
+  // Schedule every 50ms.
+  SetFrequency(/* hertz= */ 20);
   sched::InitScheduler();
-  sched::EnableScheduler();
 
+  dbgln("[boot] Loading sys init program.");
   LoadInitProgram();
+
+  dbgln("[boot] Init finished, yielding.");
+  sched::EnableScheduler();
   sched::Yield();
 
   dbgln("Sleeping!");