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

@@ -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 };

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);

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))
+					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,12 +387,11 @@ namespace Kernel
 		{
 			if (!node->blocked)
 				return;
-			if (node != m_current)
+			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);
 			update_most_loaded_node_queue(node, &m_run_queue);
 		}
@@ -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)
 		{