From edc95e1c09d84e2ed2815493bdba0b340911b969 Mon Sep 17 00:00:00 2001
From: Bananymous <oskari.alaranta@bananymous.com>
Date: Sat, 11 Jul 2026 05:30:57 +0300
Subject: [PATCH] Kernel: Disable SMP message handling while blocking a thread

There were some race conditions on a thread getting unblocked before
being moved to the block queue
---
 kernel/include/kernel/Processor.h | 11 ++++++++++-
 kernel/kernel/Processor.cpp       |  2 ++
 kernel/kernel/Scheduler.cpp       | 22 ++++++++++++----------
 3 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/kernel/include/kernel/Processor.h b/kernel/include/kernel/Processor.h
index b0c85506..3359133e 100644
--- a/kernel/include/kernel/Processor.h
+++ b/kernel/include/kernel/Processor.h
@@ -104,10 +104,17 @@ namespace Kernel
 		static void pause()
 		{
 			__builtin_ia32_pause();
-			if (is_smp_enabled())
+			if (is_smp_enabled() && !smp_messages_disabled())
 				handle_smp_messages();
 		}
 
+		static bool smp_messages_disabled() { return read_gs_sized<bool>(offsetof(Processor, m_smp_messages_disabled)); }
+		static void set_disable_smp_messages(bool disabled)
+		{
+			ASSERT(smp_messages_disabled() != disabled);
+			write_gs_sized<bool>(offsetof(Processor, m_smp_messages_disabled), disabled);
+		}
+
 		vaddr_t stack_top_vaddr() const { return m_stack_vaddr + s_stack_size; }
 		paddr_t stack_top_paddr() const { return m_stack_paddr + s_stack_size; }
 
@@ -198,6 +205,8 @@ namespace Kernel
 		ProcessorID m_id { 0 };
 		uint8_t m_index { 0 };
 
+		bool m_smp_messages_disabled { false };
+
 		vaddr_t m_thread_syscall_stack;
 
 		Thread* m_sse_thread { nullptr };
diff --git a/kernel/kernel/Processor.cpp b/kernel/kernel/Processor.cpp
index 98ce7663..e216ab9a 100644
--- a/kernel/kernel/Processor.cpp
+++ b/kernel/kernel/Processor.cpp
@@ -447,6 +447,8 @@ namespace Kernel
 		auto processor_id = current_id();
 		auto& processor = s_processors[processor_id.m_id];
 
+		ASSERT(!processor.m_smp_messages_disabled);
+
 		auto* pending = processor.m_smp_pending.exchange(nullptr);
 		if (pending == nullptr)
 			return set_interrupt_state(state);
diff --git a/kernel/kernel/Scheduler.cpp b/kernel/kernel/Scheduler.cpp
index c3d50f70..5ea5cab0 100644
--- a/kernel/kernel/Scheduler.cpp
+++ b/kernel/kernel/Scheduler.cpp
@@ -235,8 +235,12 @@ namespace Kernel
 					add_current_to_most_loaded(m_current->blocked ? &m_block_queue : &m_run_queue);
 					if (!m_current->blocked)
 						m_run_queue.add_thread_to_back(m_current);
-					else if (m_block_queue.add_thread_with_wake_time(m_current))
-						update_wake_up_deadline();
+					else
+					{
+						if (m_block_queue.add_thread_with_wake_time(m_current))
+							update_wake_up_deadline();
+						Processor::set_disable_smp_messages(false);
+					}
 					break;
 				case Thread::State::NotStarted:
 					ASSERT(!m_current->blocked);
@@ -383,13 +387,12 @@ namespace Kernel
 		{
 			if (!node->blocked)
 				return;
-			if (node != m_current)
-				m_block_queue.remove_node(node);
+			ASSERT(node != m_current);
+			m_block_queue.remove_node(node);
 			if (auto* blocker = node->blocker.load())
 				blocker->remove_thread_from_block_queue(node);
 			node->blocked = false;
-			if (node != m_current)
-				m_run_queue.add_thread_to_back(node);
+			m_run_queue.add_thread_to_back(node);
 			update_most_loaded_node_queue(node, &m_run_queue);
 		}
 		else
@@ -602,7 +605,6 @@ namespace Kernel
 			}
 
 			thread_info.node->processor_id = least_loaded_id;
-
 			Processor::send_smp_message(least_loaded_id, {
 				.type = Processor::SMPMessage::Type::NewThread,
 				.new_thread = thread_info.node
@@ -676,14 +678,14 @@ namespace Kernel
 		ASSERT(m_current->processor_id == Processor::current_id());
 		ASSERT(!m_current->blocked);
 
+		Processor::set_disable_smp_messages(true);
+
 		m_current->blocked = true;
 		m_current->wake_time_ns = wake_time_ns;
 
-		if (blocker)
+		if (blocker != nullptr)
 			blocker->add_thread_to_block_queue(m_current);
 
-		update_most_loaded_node_queue(m_current, &m_block_queue);
-
 		uint32_t lock_depth = 0;
 		if (mutex != nullptr)
 		{