Scheduler with working threads.

Currently only one process but it is a start.

zion/CMakeLists.txt

@@ -8,6 +8,7 @@ add_executable(zion
   memory/kernel_heap.cpp
   memory/paging_util.cpp
   memory/physical_memory.cpp
+  scheduler/context_switch.s
   scheduler/process.cpp
   scheduler/scheduler.cpp
   scheduler/thread.cpp

zion/memory/kernel_heap.cpp

@@ -31,5 +31,6 @@ void* KernelHeap::Allocate(uint64_t size) {
 }
 
 void* operator new(uint64_t size) { return GetKernelHeap().Allocate(size); }
+void* operator new[](uint64_t size) { return GetKernelHeap().Allocate(size); }
 
 void operator delete(void*, uint64_t) {}

zion/scheduler/context_switch.s

@@ -0,0 +1,41 @@
+.global context_switch
+context_switch:
+  push %rax
+  push %rcx
+  push %rdx
+  push %rbx
+  push %rbp
+  push %rsi
+  push %rdi
+  push %r8
+  push %r9
+  push %r10
+  push %r11
+  push %r12
+  push %r13
+  push %r14
+  push %r15
+  mov %cr3, %rax
+  push %rax
+
+  mov %rsp, (%rdi) # Save rsp to the prev task.
+  mov (%rsi), %rsp # Load the next task's rsp.
+
+  pop %rax
+  mov %rax, %cr3
+  pop %r15
+  pop %r14
+  pop %r13
+  pop %r12
+  pop %r11
+  pop %r10
+  pop %r9
+  pop %r8
+  pop %rdi
+  pop %rsi
+  pop %rbp
+  pop %rbx
+  pop %rdx
+  pop %rcx
+  pop %rax
+  retq

zion/scheduler/process.cpp

@@ -14,7 +14,7 @@ Process* Process::RootProcess() {
   asm volatile("mov %%cr3, %0;" : "=r"(pml4_addr));
   Process* proc = new Process(0, pml4_addr);
   proc->thread_list_front_ = new ThreadEntry{
-      .thread = new Thread(proc, 0),
+      .thread = Thread::RootThread(proc),
       .next = nullptr,
   };
   proc->next_thread_id_ = 1;
@@ -22,6 +22,20 @@ Process* Process::RootProcess() {
   return proc;
 }
 
+Thread* Process::CreateThread() {
+  Thread* thread = new Thread(this, next_thread_id_++);
+
+  ThreadEntry* entry = thread_list_front_;
+  while (entry->next != nullptr) {
+    entry = entry->next;
+  }
+  entry->next = new ThreadEntry{
+      .thread = thread,
+      .next = nullptr,
+  };
+  return thread;
+}
+
 Thread* Process::GetThread(uint64_t tid) {
   ThreadEntry* entry = thread_list_front_;
   while (entry != nullptr) {

zion/scheduler/process.h

@@ -12,6 +12,7 @@ class Process {
   Process();
 
   uint64_t id() { return id_; }
+  uint64_t cr3() { return cr3_; }
 
   Thread* CreateThread();
   Thread* GetThread(uint64_t tid);

zion/scheduler/scheduler.cpp

@@ -5,6 +5,8 @@
 namespace sched {
 namespace {
 
+extern "C" void context_switch(uint64_t* current_esp, uint64_t* next_esp);
+
 // Simple linked list class with the intent of eventually replacing this with a
 // map.
 class ProcList {
@@ -47,13 +49,41 @@ class Scheduler {
     Process* root = Process::RootProcess();
     current_thread_ = root->GetThread(0);
     proc_list_.InsertProcess(Process::RootProcess());
+    // FIXME: Don't enqueue threads here.
+    Enqueue(root->CreateThread());
   }
   void Enable() { enabled_ = true; }
 
   Process& CurrentProcess() { return current_thread_->process(); }
   Thread& CurrentThread() { return *current_thread_; }
 
-  void Yield() {}
+  void Enqueue(Thread* thread) {
+    Thread* back = current_thread_;
+    while (back->next_thread_ != nullptr) {
+      back = back->next_thread_;
+    }
+    back->next_thread_ = thread;
+  }
+
+  void Yield() {
+    if (!enabled_) {
+      return;
+    }
+    asm volatile("cli");
+
+    if (current_thread_->next_thread_ == nullptr) {
+      dbgln("No next thread, continue");
+      return;
+    }
+
+    Thread* prev = current_thread_;
+    current_thread_ = current_thread_->next_thread_;
+    prev->next_thread_ = nullptr;
+    Enqueue(prev);
+    context_switch(prev->Rsp0Ptr(), current_thread_->Rsp0Ptr());
+
+    asm volatile("sti");
+  }
 
  private:
   bool enabled_ = false;

zion/scheduler/scheduler.h

@@ -18,6 +18,8 @@ void Yield();
 // of the created process.
 void InsertProcess(Process* proc);
 
+void EnqueueThread(Thread* thread);
+
 Process& CurrentProcess();
 Thread& CurrentThread();
 

zion/scheduler/thread.cpp

@@ -1,5 +1,34 @@
 #include "scheduler/thread.h"
 
+#include "debug/debug.h"
 #include "scheduler/process.h"
+#include "scheduler/scheduler.h"
+
+namespace {
+
+extern "C" void thread_init() {
+  asm("sti");
+  dbgln("New Thread!");
+  sched::Yield();
+  panic("End of thread.");
+}
+
+}  // namespace
+
+Thread* Thread::RootThread(Process* root_proc) { return new Thread(root_proc); }
+
+Thread::Thread(Process* proc, uint64_t tid) : process_(proc), id_(tid) {
+  uint64_t* stack = new uint64_t[512];
+  uint64_t* stack_ptr = stack + 511;
+  // 0: rip
+  *(stack_ptr) = reinterpret_cast<uint64_t>(thread_init);
+  // 1-4: rax, rcx, rdx, rbx
+  // 5: rbp
+  *(stack_ptr - 5) = reinterpret_cast<uint64_t>(stack_ptr + 1);
+  // 6-15: rsi, rdi, r8, r9, r10, r11, r12, r13, r14, r15
+  // 16: cr3
+  *(stack_ptr - 16) = proc->cr3();
+  rsp0_ = reinterpret_cast<uint64_t>(stack_ptr - 16);
+}
 
 uint64_t Thread::pid() { return process_->id(); }

zion/scheduler/thread.h

@@ -7,17 +7,27 @@ class Process;
 
 class Thread {
  public:
-  Thread(Process* proc, uint64_t thread_id) : process_(proc), id_(thread_id) {}
+  static Thread* RootThread(Process* root_proc);
+
+  explicit Thread(Process* proc, uint64_t tid);
 
   uint64_t tid() { return id_; };
   uint64_t pid();
 
   Process& process() { return *process_; }
 
+  uint64_t* Rsp0Ptr() { return &rsp0_; }
+
+  // FIXME: Probably make this private.
+  Thread* next_thread_;
+
  private:
+  // Special constructor for the root thread only.
+  Thread(Process* proc) : process_(proc), id_(0) {}
   Process* process_;
   uint64_t id_;
 
-  // Next task in queue for quick scheduling.
-  Thread* next_task_;
+  // Stack pointer to take on resume.
+  // Stack will contain the full thread context.
+  uint64_t rsp0_;
 };

zion/zion.cpp

@@ -18,6 +18,8 @@ extern "C" void zion() {
   phys_mem::InitPhysicalMemoryManager();
 
   sched::InitScheduler();
+  sched::EnableScheduler();
+  sched::Yield();
 
   dbgln("Sleeping!");
   while (1)