From 13f5d09cfd5885f986cac25ba03774abed0c83b5 Mon Sep 17 00:00:00 2001
From: apio <blobs.trading@gmail.com>
Date: Wed, 7 Dec 2022 15:03:34 +0100
Subject: [PATCH] Add a Scheduler!!

---
 kernel/CMakeLists.txt           |   3 +
 kernel/src/thread/Scheduler.cpp | 144 ++++++++++++++++++++++++++++++++
 kernel/src/thread/Scheduler.h   |  19 +++++
 3 files changed, 166 insertions(+)
 create mode 100644 kernel/src/thread/Scheduler.cpp
 create mode 100644 kernel/src/thread/Scheduler.h

diff --git a/kernel/CMakeLists.txt b/kernel/CMakeLists.txt
index e20cca1c..5e8e03c6 100644
--- a/kernel/CMakeLists.txt
+++ b/kernel/CMakeLists.txt
@@ -12,6 +12,8 @@ set(SOURCES
     src/boot/Init.cpp
     src/arch/Serial.cpp
     src/arch/Timer.cpp
+    src/thread/Thread.cpp
+    src/thread/Scheduler.cpp
 )
 
 if("${ARCH}" MATCHES "x86_64")
@@ -22,6 +24,7 @@ if("${ARCH}" MATCHES "x86_64")
         src/arch/x86_64/MMU.cpp
         src/arch/x86_64/CPU.cpp
         src/arch/x86_64/Timer.cpp
+        src/arch/x86_64/Thread.cpp
     )
 endif()
 
diff --git a/kernel/src/thread/Scheduler.cpp b/kernel/src/thread/Scheduler.cpp
new file mode 100644
index 00000000..c4058e24
--- /dev/null
+++ b/kernel/src/thread/Scheduler.cpp
@@ -0,0 +1,144 @@
+#include "thread/Scheduler.h"
+#include "Log.h"
+#include "arch/CPU.h"
+#include "arch/MMU.h"
+#include "memory/KernelVM.h"
+#include "memory/MemoryManager.h"
+#include <luna/Stack.h>
+
+static Thread g_idle;
+static Thread* g_current = nullptr;
+
+static const usize TICKS_PER_TIMESLICE = 20;
+
+namespace Scheduler
+{
+    void init()
+    {
+        g_idle.id = 0;
+        g_idle.init_regs_kernel();
+        g_idle.set_ip((u64)CPU::idle_loop);
+        g_idle.is_idle = true;
+
+        g_idle.ticks_left = 1;
+
+        // Map some stack for the idle task
+        u64 idle_stack_vm = KernelVM::alloc_one_page().release_value();
+        MemoryManager::alloc_at(idle_stack_vm, 1, MMU::NoExecute | MMU::ReadWrite).release_value();
+
+        Stack idle_stack{idle_stack_vm, ARCH_PAGE_SIZE};
+        g_idle.set_sp(idle_stack.top());
+
+        kinfoln("CREATED IDLE THREAD: id %lu with ip %lx and sp %lx", g_idle.id, g_idle.ip(), g_idle.sp());
+
+        g_current = &g_idle;
+    }
+
+    Thread* current()
+    {
+        return g_current;
+    }
+
+    Result<void> kernel_thread_alloc_stack_and_append_impl(Thread* thread)
+    {
+        // FIXME: We will leak the thread if VM allocation or alloc_at fail.
+        u64 thread_stack_vm = TRY(KernelVM::alloc_several_pages(4));
+        TRY(MemoryManager::alloc_at(thread_stack_vm, 4, MMU::NoExecute | MMU::ReadWrite));
+
+        Stack thread_stack{thread_stack_vm, ARCH_PAGE_SIZE * 4};
+        thread->set_sp(thread_stack.top());
+
+        g_threads.append(thread);
+
+        kinfoln("CREATED THREAD: id %lu with ip %lx and sp %lx", thread->id, thread->ip(), thread->sp());
+
+        return {};
+    }
+
+    Result<void> new_kernel_thread(u64 address)
+    {
+        Thread* thread = TRY(new_thread());
+        thread->init_regs_kernel();
+        thread->set_ip(address);
+
+        return kernel_thread_alloc_stack_and_append_impl(thread);
+    }
+
+    Result<void> new_kernel_thread(void (*func)(void))
+    {
+        Thread* thread = TRY(new_thread());
+        thread->init_regs_kernel();
+        thread->set_ip((u64)func);
+
+        return kernel_thread_alloc_stack_and_append_impl(thread);
+    }
+
+    Result<void> new_kernel_thread(void (*func)(void*), void* arg)
+    {
+        Thread* thread = TRY(new_thread());
+        thread->init_regs_kernel();
+        thread->set_ip((u64)func);
+        thread->set_arguments((u64)arg, 0, 0, 0);
+
+        return kernel_thread_alloc_stack_and_append_impl(thread);
+    }
+
+    Thread* pick_task()
+    {
+        Thread* old = g_current;
+        if (old->is_idle)
+        {
+            auto maybe_first = g_threads.last();
+            if (maybe_first.has_error()) // No threads!!
+                return &g_idle;
+            g_current = old = maybe_first.value();
+        }
+
+        do {
+            auto maybe_next = g_threads.next(g_current);
+            if (maybe_next.has_error()) g_current = g_threads.first().value();
+            else
+                g_current = maybe_next.value();
+
+            if (true) // FIXME: Check if the current task is runnable.
+                break;
+        } while (g_current != old);
+
+        return g_current;
+    }
+
+    void generic_switch_context(Thread* old_thread, Thread* new_thread, Registers* regs)
+    {
+        if (old_thread != new_thread) switch_context(old_thread, new_thread, regs);
+
+        if (new_thread->is_idle)
+        {
+            new_thread->ticks_left = 1; // The idle task only runs for 1 tick so we can check for new runnable tasks
+                                        // as fast as possible.
+        }
+        else
+            new_thread->ticks_left = TICKS_PER_TIMESLICE;
+    }
+
+    void switch_task(Registers* regs)
+    {
+        Thread* old_thread = g_current;
+        Thread* new_thread = pick_task();
+        generic_switch_context(old_thread, new_thread, regs);
+    }
+
+    void invoke(Registers* regs)
+    {
+        CPU::disable_interrupts();
+
+        g_current->ticks++;
+
+        if (is_in_kernel(regs)) g_current->ticks_in_kernel++;
+        else
+            g_current->ticks_in_user++;
+
+        g_current->ticks_left--;
+
+        if (!g_current->ticks_left) switch_task(regs);
+    }
+}
\ No newline at end of file
diff --git a/kernel/src/thread/Scheduler.h b/kernel/src/thread/Scheduler.h
new file mode 100644
index 00000000..95b12451
--- /dev/null
+++ b/kernel/src/thread/Scheduler.h
@@ -0,0 +1,19 @@
+#pragma once
+#include "thread/Thread.h"
+
+namespace Scheduler
+{
+    void init();
+
+    Thread* current();
+
+    Result<void> new_kernel_thread(u64 address);
+    Result<void> new_kernel_thread(void (*func)(void));
+    Result<void> new_kernel_thread(void (*func)(void*), void* arg);
+
+    Thread* pick_task();
+
+    void switch_task(Registers* regs);
+
+    void invoke(Registers* regs);
+}
\ No newline at end of file