Kernel: Replace last CriticalScopes in kernel with SpinLocks

kernel/include/kernel/Scheduler.h

@@ -30,9 +30,12 @@ namespace Kernel
 		static pid_t current_tid();
 
 		[[noreturn]] void execute_current_thread();
-		[[noreturn]] void _execute_current_thread();
+		[[noreturn]] void execute_current_thread_locked();
 		[[noreturn]] void delete_current_process_and_thread();
 
+		// This is no return if called on current thread
+		void terminate_thread(Thread*);
+
 	private:
 		Scheduler() = default;
 
@@ -43,6 +46,8 @@ namespace Kernel
 		void remove_and_advance_current_thread();
 		void advance_current_thread();
 
+		[[noreturn]] void execute_current_thread_stack_loaded();
+
 		BAN::ErrorOr<void> add_thread(Thread*);
 
 	private:
@@ -57,6 +62,8 @@ namespace Kernel
 			Semaphore*	semaphore;
 		};
 
+		SpinLockUnsafe m_lock;
+
 		Thread* m_idle_thread { nullptr };
 		BAN::LinkedList<SchedulerThread> m_active_threads;
 		BAN::LinkedList<SchedulerThread> m_sleeping_threads;

kernel/include/kernel/Thread.h

@@ -65,8 +65,6 @@ namespace Kernel
 		uintptr_t rip() const { return m_rip; }
 
 		void set_started() { ASSERT(m_state == State::NotStarted); m_state = State::Executing; }
-		// Thread will no longer execute. If called on current thread, does not return
-		void terminate();
 		State state() const { return m_state; }
 
 		vaddr_t stack_base() const { return m_stack->vaddr(); }

kernel/kernel/Process.cpp

@@ -230,7 +230,8 @@ namespace Kernel
 			m_threads.clear();
 
 			thread.setup_process_cleanup();
-			Scheduler::get().execute_current_thread();
+			// NOTE: This function is only called from scheduler when it is already locked
+			Scheduler::get().execute_current_thread_locked();
 			ASSERT_NOT_REACHED();
 		}
 
@@ -252,9 +253,9 @@ namespace Kernel
 		m_exit_status.exit_code = __WGENEXITCODE(status, signal);
 		for (auto* thread : m_threads)
 			if (thread != &Thread::current())
-				thread->terminate();
+				Scheduler::get().terminate_thread(thread);
 		if (this == &Process::current())
-			Thread::current().terminate();
+			Scheduler::get().terminate_thread(&Thread::current());
 	}
 
 	size_t Process::proc_meminfo(off_t offset, BAN::ByteSpan buffer) const
@@ -330,8 +331,8 @@ namespace Kernel
 
 	BAN::ErrorOr<long> Process::sys_exit(int status)
 	{
+		ASSERT(this == &Process::current());
 		exit(status, 0);
-		Thread::current().terminate();
 		ASSERT_NOT_REACHED();
 	}
 

kernel/kernel/Scheduler.cpp

@@ -1,6 +1,5 @@
 #include <kernel/Arch.h>
 #include <kernel/Attributes.h>
-#include <kernel/CriticalScope.h>
 #include <kernel/GDT.h>
 #include <kernel/InterruptController.h>
 #include <kernel/Process.h>
@@ -8,15 +7,6 @@
 #include <kernel/Timer/Timer.h>
 
 #define SCHEDULER_VERIFY_STACK 1
-#define SCHEDULER_VERIFY_INTERRUPT_STATE 1
-
-#if SCHEDULER_VERIFY_INTERRUPT_STATE
-	#define VERIFY_STI() ASSERT(get_interrupt_state() == InterruptState::Enabled)
-	#define VERIFY_CLI() ASSERT(get_interrupt_state() == InterruptState::Disabled)
-#else
-	#define VERIFY_STI()
-	#define VERIFY_CLI()
-#endif
 
 namespace Kernel
 {
@@ -50,10 +40,11 @@ namespace Kernel
 
 	void Scheduler::start()
 	{
-		VERIFY_CLI();
+		ASSERT(get_interrupt_state() == InterruptState::Disabled);
+		m_lock.lock();
 		ASSERT(!m_active_threads.empty());
 		m_current_thread = m_active_threads.begin();
-		execute_current_thread();
+		execute_current_thread_locked();
 		ASSERT_NOT_REACHED();
 	}
 
@@ -71,45 +62,40 @@ namespace Kernel
 
 	void Scheduler::timer_reschedule()
 	{
-		VERIFY_CLI();
+		auto state = m_lock.lock();
-
 		wake_threads();
-
 		if (save_current_thread())
-			return;
+			return m_lock.unlock(state);
 		advance_current_thread();
-		execute_current_thread();
+		execute_current_thread_locked();
 		ASSERT_NOT_REACHED();
 	}
 
 	void Scheduler::reschedule()
 	{
-		set_interrupt_state(InterruptState::Disabled);
+		auto state = m_lock.lock();
-
 		if (save_current_thread())
-			return set_interrupt_state(InterruptState::Enabled);
+			return set_interrupt_state(state);
 		advance_current_thread();
-		execute_current_thread();
+		execute_current_thread_locked();
 		ASSERT_NOT_REACHED();
 	}
 
 	void Scheduler::reschedule_if_idling()
 	{
-		VERIFY_CLI();
+		auto state = m_lock.lock();
-
 		if (m_active_threads.empty() || &current_thread() != m_idle_thread)
-			return;
+			return m_lock.unlock(state);
-
 		if (save_current_thread())
-			return;
+			return m_lock.unlock(state);
 		m_current_thread = m_active_threads.begin();
-		execute_current_thread();
+		execute_current_thread_locked();
 		ASSERT_NOT_REACHED();
 	}
 
 	void Scheduler::wake_threads()
 	{
-		VERIFY_CLI();
+		ASSERT(m_lock.is_locked());
 
 		uint64_t current_time = SystemTimer::get().ms_since_boot();
 		while (!m_sleeping_threads.empty() && m_sleeping_threads.front().wake_time <= current_time)
@@ -124,14 +110,22 @@ namespace Kernel
 
 	BAN::ErrorOr<void> Scheduler::add_thread(Thread* thread)
 	{
-		CriticalScope _;
+		SpinLockGuard _(m_lock);
 		TRY(m_active_threads.emplace_back(thread));
 		return {};
 	}
 
+	void Scheduler::terminate_thread(Thread* thread)
+	{
+		SpinLockGuard _(m_lock);
+		thread->m_state = Thread::State::Terminated;
+		if (thread == &current_thread())
+			execute_current_thread_locked();
+	}
+
 	void Scheduler::advance_current_thread()
 	{
-		VERIFY_CLI();
+		ASSERT(m_lock.is_locked());
 
 		if (m_active_threads.empty())
 		{
@@ -144,7 +138,7 @@ namespace Kernel
 
 	void Scheduler::remove_and_advance_current_thread()
 	{
-		VERIFY_CLI();
+		ASSERT(m_lock.is_locked());
 
 		ASSERT(m_current_thread);
 
@@ -165,7 +159,7 @@ namespace Kernel
 	//       after getting the rsp
 	ALWAYS_INLINE bool Scheduler::save_current_thread()
 	{
-		VERIFY_CLI();
+		ASSERT(m_lock.is_locked());
 
 		uintptr_t rsp, rip;
 		push_callee_saved();
@@ -187,7 +181,7 @@ namespace Kernel
 
 	void Scheduler::delete_current_process_and_thread()
 	{
-		set_interrupt_state(InterruptState::Disabled);
+		m_lock.lock();
 
 		load_temp_stack();
 		PageTable::kernel().load();
@@ -201,23 +195,33 @@ namespace Kernel
 
 		delete thread;
 
-		execute_current_thread();
+		execute_current_thread_locked();
 		ASSERT_NOT_REACHED();
 	}
 
 	void Scheduler::execute_current_thread()
 	{
-		VERIFY_CLI();
+		m_lock.lock();
-
 		load_temp_stack();
 		PageTable::kernel().load();
-		_execute_current_thread();
+		execute_current_thread_stack_loaded();
 		ASSERT_NOT_REACHED();
 	}
 
-	NEVER_INLINE void Scheduler::_execute_current_thread()
+	void Scheduler::execute_current_thread_locked()
 	{
-		VERIFY_CLI();
+		load_temp_stack();
+		PageTable::kernel().load();
+		execute_current_thread_stack_loaded();
+		ASSERT_NOT_REACHED();
+	}
+
+	NEVER_INLINE void Scheduler::execute_current_thread_stack_loaded()
+	{
+		ASSERT(m_lock.is_locked());
+
+		load_temp_stack();
+		PageTable::kernel().load();
 
 #if SCHEDULER_VERIFY_STACK
 		vaddr_t rsp;
@@ -264,10 +268,12 @@ namespace Kernel
 		{
 			case Thread::State::NotStarted:
 				current->set_started();
+				m_lock.unlock(InterruptState::Disabled);
 				start_thread(current->rsp(), current->rip());
 			case Thread::State::Executing:
 				while (current->can_add_signal_to_execute())
 					current->handle_signal();
+				m_lock.unlock(InterruptState::Disabled);
 				continue_thread(current->rsp(), current->rip());
 			case Thread::State::Terminated:
 				ASSERT_NOT_REACHED();
@@ -278,7 +284,7 @@ namespace Kernel
 
 	void Scheduler::set_current_thread_sleeping_impl(uint64_t wake_time)
 	{
-		VERIFY_CLI();
+		ASSERT(m_lock.is_locked());
 
 		if (save_current_thread())
 		{
@@ -301,35 +307,27 @@ namespace Kernel
 		m_current_thread = {};
 		advance_current_thread();
 
-		execute_current_thread();
+		execute_current_thread_locked();
 		ASSERT_NOT_REACHED();
 	}
 
 	void Scheduler::set_current_thread_sleeping(uint64_t wake_time)
 	{
-		VERIFY_STI();
+		SpinLockGuard _(m_lock);
-		set_interrupt_state(InterruptState::Disabled);
-
-		ASSERT(m_current_thread);
-
 		m_current_thread->semaphore = nullptr;
 		set_current_thread_sleeping_impl(wake_time);
 	}
 
 	void Scheduler::block_current_thread(Semaphore* semaphore, uint64_t wake_time)
 	{
-		VERIFY_STI();
+		SpinLockGuard _(m_lock);
-		set_interrupt_state(InterruptState::Disabled);
-
-		ASSERT(m_current_thread);
-
 		m_current_thread->semaphore = semaphore;
 		set_current_thread_sleeping_impl(wake_time);
 	}
 
 	void Scheduler::unblock_threads(Semaphore* semaphore)
 	{
-		CriticalScope critical;
+		SpinLockGuard _(m_lock);
 
 		for (auto it = m_sleeping_threads.begin(); it != m_sleeping_threads.end();)
 		{
@@ -350,7 +348,7 @@ namespace Kernel
 
 	void Scheduler::unblock_thread(pid_t tid)
 	{
-		CriticalScope _;
+		SpinLockGuard _(m_lock);
 
 		for (auto it = m_sleeping_threads.begin(); it != m_sleeping_threads.end(); it++)
 		{

kernel/kernel/Thread.cpp

@@ -27,14 +27,6 @@ namespace Kernel
 			memcpy((void*)rsp, (void*)&value, sizeof(uintptr_t));
 	}
 
-	void Thread::terminate()
-	{
-		set_interrupt_state(InterruptState::Disabled);
-		m_state = Thread::State::Terminated;
-		if (this == &Thread::current())
-			Scheduler::get().execute_current_thread();
-	}
-
 	static pid_t s_next_tid = 1;
 
 	BAN::ErrorOr<Thread*> Thread::create_kernel(entry_t entry, void* data, Process* process)
@@ -193,9 +185,10 @@ namespace Kernel
 	{
 		m_state = State::NotStarted;
 		static entry_t entry(
-			[](void* process)
+			[](void* process_ptr)
 			{
-				((Process*)process)->cleanup_function();
+				auto& process = *reinterpret_cast<Process*>(process_ptr);
+				process.cleanup_function();
 				Scheduler::get().delete_current_process_and_thread();
 				ASSERT_NOT_REACHED();
 			}
@@ -395,7 +388,7 @@ namespace Kernel
 	void Thread::on_exit()
 	{
 		ASSERT(this == &Thread::current());
-		terminate();
+		Scheduler::get().terminate_thread(this);
 		ASSERT_NOT_REACHED();
 	}