Move Processes and Threads to be stored in SharedPtr

Reference counting lets us pass these around a bit more easily.

SharedPtr was lightly tested using uint64_t in the main zion funcion.

Also add a sleep functionality instead of panicking. Functionally the
same right now since we don't preempt.

zion/lib/linked_list.h

@@ -0,0 +1,100 @@
+#pragma once
+
+#include <stdint.h>
+
+#include "debug/debug.h"
+
+template <typename T>
+class LinkedList {
+ public:
+  LinkedList() {}
+
+  LinkedList(const LinkedList&) = delete;
+
+  uint64_t size() const { return size_; }
+
+  void PushBack(const T& item) {
+    size_++;
+    ListItem* new_item = new ListItem{
+        .item = item,
+        .next = nullptr,
+    };
+    if (front_ == nullptr) {
+      front_ = new_item;
+      return;
+    }
+    ListItem* litem = front_;
+    while (litem->next != nullptr) {
+      litem = litem->next;
+    }
+    litem->next = new_item;
+  }
+
+  T PopFront() {
+    if (size_ == 0 || front_ == nullptr) {
+      panic("Popping from empty list");
+    }
+
+    size_--;
+
+    ListItem* old_front = front_;
+    front_ = front_->next;
+    T ret = old_front->item;
+    delete old_front;
+    return ret;
+  }
+
+  T CycleFront() {
+    if (size_ == 0 || front_ == nullptr) {
+      panic("Cycling empty list");
+    }
+
+    if (size_ == 1) {
+      return front_->item;
+    }
+
+    T ret = front_->item;
+    ListItem* old_front = front_;
+    ListItem* iter = front_;
+    front_ = front_->next;
+    while (iter->next != nullptr) {
+      iter = iter->next;
+    }
+    iter->next = old_front;
+    old_front->next = nullptr;
+    return ret;
+  }
+
+  T PeekFront() const { return front_->item; }
+
+  struct ListItem {
+    T item;
+    ListItem* next;
+  };
+  class Iterator {
+   public:
+    Iterator(ListItem* item) : item_(item) {}
+
+    Iterator next() { return {item_->next}; }
+    Iterator& operator++() {
+      item_ = item_->next;
+      return *this;
+    }
+
+    T& operator*() { return item_->item; }
+    T& operator->() { return item_->item; }
+    bool operator==(const Iterator& other) { return item_ == other.item_; }
+    bool operator!=(const Iterator& other) { return item_ != other.item_; }
+
+   private:
+    ListItem* item_;
+  };
+
+  Iterator begin() { return {front_}; }
+  Iterator end() { return {nullptr}; }
+
+ private:
+  uint64_t size_ = 0;
+
+  ListItem* front_ = nullptr;
+};

zion/lib/shared_ptr.h

@@ -0,0 +1,63 @@
+#pragma once
+
+#include <stdint.h>
+
+#include "debug/debug.h"
+
+template <typename T>
+class SharedPtr {
+ public:
+  SharedPtr() : init_(false), ptr_(0), ref_cnt_(0) {}
+  // Takes ownership.
+  SharedPtr(T* ptr) {
+    ptr_ = ptr;
+    ref_cnt_ = new uint64_t(1);
+  }
+
+  SharedPtr(const SharedPtr<T>& other)
+      : ptr_(other.ptr_), ref_cnt_(other.ref_cnt_) {
+    (*ref_cnt_)++;
+  }
+
+  SharedPtr& operator=(const SharedPtr<T>& other) {
+    Cleanup();
+    ptr_ = other.ptr_;
+    ref_cnt_ = other.ref_cnt_;
+    (*ref_cnt_)++;
+    return *this;
+  }
+
+  ~SharedPtr() { Cleanup(); }
+
+  T& operator*() { return *ptr_; }
+  const T& operator*() const { return *ptr_; }
+  T* operator->() { return ptr_; }
+  const T* operator->() const { return ptr_; }
+
+  T* ptr() { return ptr_; }
+
+  bool operator==(const SharedPtr<T>& other) { return ptr_ == other.ptr_; }
+
+  bool empty() { return !init_; }
+
+ private:
+  bool init_ = true;
+  T* ptr_;
+  uint64_t* ref_cnt_;
+
+  void Cleanup() {
+    if (!init_) {
+      return;
+    }
+    if (--(*ref_cnt_) == 0) {
+      dbgln("Deleting shared ptr: %m", ptr_);
+      delete ptr_;
+      delete ref_cnt_;
+    }
+  }
+};
+
+template <typename T, class... A>
+SharedPtr<T> MakeShared(A... args) {
+  return {new T(args...)};
+}

zion/scheduler/process.cpp

@@ -12,14 +12,11 @@ static uint64_t gNextId = 1;
 
 }
 
-Process* Process::RootProcess() {
+SharedPtr<Process> Process::RootProcess() {
   uint64_t pml4_addr = 0;
   asm volatile("mov %%cr3, %0;" : "=r"(pml4_addr));
-  Process* proc = new Process(0, pml4_addr);
+  SharedPtr<Process> proc(new Process(0, pml4_addr));
-  proc->thread_list_front_ = new ThreadEntry{
+  proc->threads_.PushBack(Thread::RootThread(proc.ptr()));
-      .thread = Thread::RootThread(proc),
-      .next = nullptr,
-  };
   proc->next_thread_id_ = 1;
 
   return proc;
@@ -33,42 +30,29 @@ Process::Process(uint64_t elf_ptr) : id_(gNextId++), state_(RUNNING) {
 
 void Process::CreateThread(uint64_t elf_ptr) {
   Thread* thread = new Thread(this, next_thread_id_++, elf_ptr);
-  ThreadEntry* tentry = new ThreadEntry{
+  threads_.PushBack(thread);
-      .thread = thread,
-      .next = nullptr,
-  };
-
-  if (thread_list_front_ == nullptr) {
-    thread_list_front_ = tentry;
-  } else {
-    ThreadEntry* entry = thread_list_front_;
-    while (entry->next != nullptr) {
-      entry = entry->next;
-    }
-    entry->next = tentry;
-  }
   sched::EnqueueThread(thread);
 }
 
-Thread* Process::GetThread(uint64_t tid) {
+SharedPtr<Thread> Process::GetThread(uint64_t tid) {
-  ThreadEntry* entry = thread_list_front_;
+  auto iter = threads_.begin();
-  while (entry != nullptr) {
+  while (iter != threads_.end()) {
-    if (entry->thread->tid() == tid) {
+    if (iter->tid() == tid) {
-      return entry->thread;
+      return *iter;
     }
+    ++iter;
   }
   panic("Bad thread access.");
   return nullptr;
 }
 
 void Process::CheckState() {
-  ThreadEntry* entry = thread_list_front_;
+  auto iter = threads_.begin();
-
+  while (iter != threads_.end()) {
-  while (entry != nullptr) {
+    if (iter->GetState() != Thread::FINISHED) {
-    if (entry->thread->GetState() != Thread::FINISHED) {
       return;
     }
-    entry = entry->next;
+    ++iter;
   }
   state_ = FINISHED;
 }

zion/scheduler/process.h

@@ -2,6 +2,9 @@
 
 #include <stdint.h>
 
+#include "lib/linked_list.h"
+#include "lib/shared_ptr.h"
+
 // Forward decl due to cyclic dependency.
 class Thread;
 
@@ -13,15 +16,14 @@ class Process {
     RUNNING,
     FINISHED,
   };
-  // Caller takes ownership of returned process.
+  static SharedPtr<Process> RootProcess();
-  static Process* RootProcess();
   Process(uint64_t elf_ptr);
 
-  uint64_t id() { return id_; }
+  uint64_t id() const { return id_; }
-  uint64_t cr3() { return cr3_; }
+  uint64_t cr3() const { return cr3_; }
 
   void CreateThread(uint64_t elf_ptr);
-  Thread* GetThread(uint64_t tid);
+  SharedPtr<Thread> GetThread(uint64_t tid);
 
   // Checks the state of all child threads and transitions to
   // finished if all have finished.
@@ -37,10 +39,5 @@ class Process {
 
   uint64_t next_thread_id_ = 0;
 
-  // FIXME: Make a better data structure for this.
+  LinkedList<SharedPtr<Thread>> threads_;
-  struct ThreadEntry {
-    Thread* thread;
-    ThreadEntry* next;
-  };
-  ThreadEntry* thread_list_front_;
 };

zion/scheduler/scheduler.cpp

@@ -1,77 +1,37 @@
 #include "scheduler/scheduler.h"
 
 #include "debug/debug.h"
+#include "lib/linked_list.h"
 
 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
+void DumpProcessStates(LinkedList<SharedPtr<Process>>& proc_list) {
-// map.
+  dbgln("Process States: %u", proc_list.size());
-class ProcList {
+  auto iter = proc_list.begin();
- public:
+  while (iter != proc_list.end()) {
-  ProcList() {}
+    dbgln("%u: %u", iter->id(), iter->GetState());
-
+    ++iter;
-  // Takes ownership.
-  void InsertProcess(Process* proc) {
-    if (front_ == nullptr) {
-      front_ = new ProcEntry{
-          .proc = proc,
-          .next = nullptr,
-      };
-      return;
-    }
-
-    ProcEntry* back = front_;
-    while (back->next != nullptr) {
-      back = back->next;
-    }
-
-    back->next = new ProcEntry{
-        .proc = proc,
-        .next = nullptr,
-    };
   }
-
+}
-  void DumpProcessStates() {
-    ProcEntry* p = front_;
-    dbgln("Process States:");
-    while (p != nullptr) {
-      dbgln("%u: %u", p->proc->id(), p->proc->GetState());
-      p = p->next;
-    }
-  }
-
- private:
-  struct ProcEntry {
-    Process* proc;
-    ProcEntry* next;
-  };
-
-  ProcEntry* front_ = nullptr;
-};
 
 class Scheduler {
  public:
   Scheduler() {
-    Process* root = Process::RootProcess();
+    SharedPtr<Process> root = Process::RootProcess();
-    current_thread_ = root->GetThread(0);
+    sleep_thread_ = root->GetThread(0);
-    proc_list_.InsertProcess(Process::RootProcess());
+    runnable_threads_.PushBack(sleep_thread_);
+    proc_list_.PushBack(Process::RootProcess());
   }
   void Enable() { enabled_ = true; }
 
-  Process& CurrentProcess() { return current_thread_->process(); }
+  Process& CurrentProcess() { return CurrentThread().process(); }
-  Thread& CurrentThread() { return *current_thread_; }
+  Thread& CurrentThread() { return *runnable_threads_.PeekFront(); }
 
-  void InsertProcess(Process* process) { proc_list_.InsertProcess(process); }
+  void InsertProcess(Process* process) { proc_list_.PushBack(process); }
-  void Enqueue(Thread* thread) {
+  void Enqueue(Thread* thread) { runnable_threads_.PushBack(thread); }
-    Thread* back = current_thread_;
-    while (back->next_thread_ != nullptr) {
-      back = back->next_thread_;
-    }
-    back->next_thread_ = thread;
-  }
 
   void Yield() {
     if (!enabled_) {
@@ -79,37 +39,48 @@ class Scheduler {
     }
     asm volatile("cli");
 
-    if (current_thread_->next_thread_ == nullptr) {
+    SharedPtr<Thread> prev;
-      if (current_thread_->GetState() == Thread::RUNNING) {
+    if (CurrentThread().GetState() == Thread::RUNNING) {
-        dbgln("No next thread, continue");
+      prev = runnable_threads_.CycleFront();
-        return;
+      prev->SetState(Thread::RUNNABLE);
-      } else {
+    } else {
-        proc_list_.DumpProcessStates();
+      // This technically is a memory operation but should only occur when a
-        panic("FIXME: Implement Sleep");
+      // thread is blocking so may be ok?
-      }
+      prev = runnable_threads_.PopFront();
     }
 
-    Thread* prev = current_thread_;
+    SharedPtr<Thread> next;
-    current_thread_ = current_thread_->next_thread_;
+    if (runnable_threads_.size() == 0) {
-    prev->next_thread_ = nullptr;
+      next = sleep_thread_;
-    if (prev->pid() != 0 && prev->GetState() == Thread::RUNNING) {
+      DumpProcessStates(proc_list_);
-      prev->SetState(Thread::RUNNABLE);
+    } else {
-      Enqueue(prev);
+      next = runnable_threads_.PeekFront();
     }
-    if (current_thread_->GetState() != Thread::RUNNABLE) {
+
+    if (next->GetState() != Thread::RUNNABLE) {
       panic("Non-runnable thread in the queue");
     }
-    current_thread_->SetState(Thread::RUNNING);
+    // Needs to be before the next == prev check
-    context_switch(prev->Rsp0Ptr(), current_thread_->Rsp0Ptr());
+    // 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");
   }
 
  private:
   bool enabled_ = false;
-  ProcList proc_list_;
+  // TODO: move this to a separate process manager class.
+  LinkedList<SharedPtr<Process>> proc_list_;
+  LinkedList<SharedPtr<Thread>> runnable_threads_;
 
-  Thread* current_thread_;
+  SharedPtr<Thread> sleep_thread_;
 };
 
 static Scheduler* gScheduler = nullptr;

zion/scheduler/thread.cpp

@@ -19,9 +19,11 @@ extern "C" void thread_init() {
 
 }  // namespace
 
-Thread* Thread::RootThread(Process* root_proc) { return new Thread(root_proc); }
+SharedPtr<Thread> Thread::RootThread(Process* root_proc) {
+  return new Thread(root_proc);
+}
 
-Thread::Thread(Process* proc, uint64_t tid, uint64_t elf_ptr)
+Thread::Thread(const SharedPtr<Process>& proc, uint64_t tid, uint64_t elf_ptr)
     : process_(proc), id_(tid), elf_ptr_(elf_ptr) {
   uint64_t* stack = new uint64_t[512];
   uint64_t* stack_ptr = stack + 511;

zion/scheduler/thread.h

@@ -2,6 +2,8 @@
 
 #include <stdint.h>
 
+#include "lib/shared_ptr.h"
+
 // Forward decl due to cyclic dependency.
 class Process;
 
@@ -15,9 +17,10 @@ class Thread {
     BLOCKED,
     FINISHED,
   };
-  static Thread* RootThread(Process* root_proc);
+  static SharedPtr<Thread> RootThread(Process* root_proc);
 
-  explicit Thread(Process* proc, uint64_t tid, uint64_t elf_ptr);
+  explicit Thread(const SharedPtr<Process>& proc, uint64_t tid,
+                  uint64_t elf_ptr);
 
   uint64_t tid() { return id_; };
   uint64_t pid();
@@ -35,13 +38,10 @@ class Thread {
   void SetState(State state) { state_ = state; }
   void Exit();
 
-  // 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_;
+  SharedPtr<Process> process_;
   uint64_t id_;
   State state_ = RUNNABLE;