From 1dabe7a22237cc615cb3f1239ef84d77be5f8b7b Mon Sep 17 00:00:00 2001
From: Bananymous <bananymousosq@gmail.com>
Date: Wed, 8 Mar 2023 03:23:46 +0200
Subject: [PATCH] Kernel: More rework on Scheduler

---
 kernel/include/kernel/Scheduler.h |  5 ++++-
 kernel/kernel/Scheduler.cpp       | 26 +++++++++++++-------------
 2 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/kernel/include/kernel/Scheduler.h b/kernel/include/kernel/Scheduler.h
index 64eb40550f..3a4bf70149 100644
--- a/kernel/include/kernel/Scheduler.h
+++ b/kernel/include/kernel/Scheduler.h
@@ -28,18 +28,21 @@ namespace Kernel
 
 		void wake_threads();
 		[[nodiscard]] bool save_current_thread();
-		void get_next_thread();
+		void remove_and_advance_current_thread();
+		void advance_current_thread();
 		[[noreturn]] void execute_current_thread();
 
 	private:
 		struct ActiveThread
 		{
+			ActiveThread(const BAN::RefPtr<Thread>& thread) : thread(thread) {}
 			BAN::RefPtr<Thread> thread;
 			uint64_t padding = 0;
 		};
 
 		struct SleepingThread
 		{
+			SleepingThread(const BAN::RefPtr<Thread>& thread, uint64_t wake_time) : thread(thread), wake_time(wake_time) {}
 			BAN::RefPtr<Thread> thread;
 			uint64_t wake_time;
 		};
diff --git a/kernel/kernel/Scheduler.cpp b/kernel/kernel/Scheduler.cpp
index 1624181b38..fd232dedea 100644
--- a/kernel/kernel/Scheduler.cpp
+++ b/kernel/kernel/Scheduler.cpp
@@ -53,6 +53,7 @@ namespace Kernel
 
 	void Scheduler::reschedule()
 	{
+		VERIFY_CLI();
 		ASSERT(InterruptController::get().is_in_service(PIT_IRQ));
 		InterruptController::get().eoi(PIT_IRQ);
 
@@ -81,23 +82,17 @@ namespace Kernel
 
 			// This should work as we released enough memory from sleeping thread
 			static_assert(sizeof(ActiveThread) == sizeof(SleepingThread));
-			MUST(m_active_threads.push_back({ thread, 0 }));
+			MUST(m_active_threads.emplace_back(thread));
+			thread.clear();
 		}
 	}
 
-	BAN::ErrorOr<void> Scheduler::add_thread(BAN::RefCounted<Thread> thread)
+	BAN::ErrorOr<void> Scheduler::add_thread(BAN::RefPtr<Thread> thread)
 	{
-		if (interrupts_disabled())
-		{
-			TRY(m_active_threads.push_back({ thread, 0 }));
-		}
-		else
-		{
-			DISABLE_INTERRUPTS();
-			TRY(m_active_threads.push_back({ thread, 0 }));
-			ENABLE_INTERRUPTS();
-		}
-		return {};
+		auto flags = disable_interrupts_and_get_flags();
+		BAN::ErrorOr<void> result = m_active_threads.emplace_back(thread);
+		restore_flags(flags);
+		return result;
 	}
 
 	void Scheduler::advance_current_thread()
@@ -176,6 +171,7 @@ namespace Kernel
 #pragma GCC optimize("O0")
 	void Scheduler::set_current_thread_sleeping(uint64_t wake_time)
 	{
+		VERIFY_STI();
 		DISABLE_INTERRUPTS();
 
 		ASSERT(m_current_thread);
@@ -183,7 +179,10 @@ namespace Kernel
 		auto sleeping = m_current_thread->thread;
 
 		if (save_current_thread())
+		{
+			ENABLE_INTERRUPTS();
 			return;
+		}
 		remove_and_advance_current_thread();
 
 		auto it = m_sleeping_threads.begin();
@@ -203,6 +202,7 @@ namespace Kernel
 
 	void Scheduler::set_current_thread_done()
 	{
+		VERIFY_STI();
 		DISABLE_INTERRUPTS();
 
 		ASSERT(m_current_thread);