Kernel: Make spinlocks more safe

Kernel now makes sure thread is not holding any spinlocks when it tries
to lock a mutex or yield. Spinlocks are supposed to be only used for
short times without the possibility of yielding

kernel/CMakeLists.txt

@@ -43,6 +43,7 @@ set(KERNEL_SOURCES
 	kernel/Interruptable.cpp
 	kernel/InterruptController.cpp
 	kernel/kernel.cpp
+	kernel/Lock/SpinLock.cpp
 	kernel/Memory/DMARegion.cpp
 	kernel/Memory/FileBackedRegion.cpp
 	kernel/Memory/Heap.cpp

kernel/include/kernel/Lock/Mutex.h

@@ -88,6 +88,8 @@ namespace Kernel
 		void lock()
 		{
 			const auto tid = Thread::current_tid();
+			ASSERT(!tid || !Thread::current().has_spinlock());
+
 			if (tid == m_locker)
 				ASSERT(m_lock_depth > 0);
 			else
@@ -111,6 +113,8 @@ namespace Kernel
 		bool try_lock()
 		{
 			const auto tid = Thread::current_tid();
+			ASSERT(!tid || !Thread::current().has_spinlock());
+
 			if (tid == m_locker)
 				ASSERT(m_lock_depth > 0);
 			else

kernel/include/kernel/Lock/SpinLock.h

@@ -18,42 +18,11 @@ namespace Kernel
 	public:
 		SpinLock() = default;
 
-		InterruptState lock()
-		{
-			auto state = Processor::get_interrupt_state();
-			Processor::set_interrupt_state(InterruptState::Disabled);
+		InterruptState lock();
 
-			auto id = Processor::current_id().as_u32();
-			ASSERT(m_locker.load(BAN::MemoryOrder::memory_order_relaxed) != id);
+		bool try_lock_interrupts_disabled();
 
-			auto expected = PROCESSOR_NONE.as_u32();
-			while (!m_locker.compare_exchange(expected, id, BAN::MemoryOrder::memory_order_acquire))
-			{
-				Processor::pause();
-				expected = PROCESSOR_NONE.as_u32();
-			}
-
-			return state;
-		}
-
-		bool try_lock_interrupts_disabled()
-		{
-			ASSERT(Processor::get_interrupt_state() == InterruptState::Disabled);
-
-			auto id = Processor::current_id().as_u32();
-			ASSERT(m_locker.load(BAN::MemoryOrder::memory_order_relaxed) != id);
-
-			auto expected = PROCESSOR_NONE.as_u32();
-			return m_locker.compare_exchange(expected, id, BAN::MemoryOrder::memory_order_acquire);
-		}
-
-		void unlock(InterruptState state)
-		{
-			ASSERT(Processor::get_interrupt_state() == InterruptState::Disabled);
-			ASSERT(current_processor_has_lock());
-			m_locker.store(PROCESSOR_NONE.as_u32(), BAN::MemoryOrder::memory_order_release);
-			Processor::set_interrupt_state(state);
-		}
+		void unlock(InterruptState state);
 
 		bool current_processor_has_lock() const
 		{

kernel/include/kernel/Thread.h

@@ -101,6 +101,10 @@ namespace Kernel
 		void save_sse();
 		void load_sse();
 
+		void add_spinlock() { m_spinlock_count++; }
+		void remove_spinlock() { m_spinlock_count--; }
+		bool has_spinlock() const { return !!m_spinlock_count; }
+
 		void add_mutex() { m_mutex_count++; }
 		void remove_mutex() { m_mutex_count--; }
 
@@ -137,6 +141,7 @@ namespace Kernel
 		SpinLock                   m_signal_lock;
 		static_assert(_SIGMAX < 64);
 
+		BAN::Atomic<uint32_t>      m_spinlock_count       { 0 };
 		BAN::Atomic<uint32_t>      m_mutex_count          { 0 };
 
 		alignas(16) uint8_t m_sse_storage[512] {};

kernel/kernel/Lock/SpinLock.cpp

@@ -0,0 +1,55 @@
+#include <kernel/Lock/SpinLock.h>
+#include <kernel/Thread.h>
+
+namespace Kernel
+{
+
+	InterruptState SpinLock::lock()
+	{
+		auto state = Processor::get_interrupt_state();
+		Processor::set_interrupt_state(InterruptState::Disabled);
+
+		auto id = Processor::current_id().as_u32();
+		ASSERT(m_locker.load(BAN::MemoryOrder::memory_order_relaxed) != id);
+
+		auto expected = PROCESSOR_NONE.as_u32();
+		while (!m_locker.compare_exchange(expected, id, BAN::MemoryOrder::memory_order_acquire))
+		{
+			Processor::pause();
+			expected = PROCESSOR_NONE.as_u32();
+		}
+
+		if (Thread::current_tid())
+			Thread::current().add_spinlock();
+
+		return state;
+	}
+
+	bool SpinLock::try_lock_interrupts_disabled()
+	{
+		ASSERT(Processor::get_interrupt_state() == InterruptState::Disabled);
+
+		auto id = Processor::current_id().as_u32();
+		ASSERT(m_locker.load(BAN::MemoryOrder::memory_order_relaxed) != id);
+
+		auto expected = PROCESSOR_NONE.as_u32();
+		if (!m_locker.compare_exchange(expected, id, BAN::MemoryOrder::memory_order_acquire))
+			return false;
+
+		if (Thread::current_tid())
+			Thread::current().add_spinlock();
+
+		return true;
+	}
+
+	void SpinLock::unlock(InterruptState state)
+	{
+		ASSERT(Processor::get_interrupt_state() == InterruptState::Disabled);
+		ASSERT(current_processor_has_lock());
+		m_locker.store(PROCESSOR_NONE.as_u32(), BAN::MemoryOrder::memory_order_release);
+		if (Thread::current_tid())
+			Thread::current().remove_spinlock();
+		Processor::set_interrupt_state(state);
+	}
+
+}

kernel/kernel/Processor.cpp

@@ -337,6 +337,8 @@ namespace Kernel
 		auto state = get_interrupt_state();
 		set_interrupt_state(InterruptState::Disabled);
 
+		ASSERT(!Thread::current().has_spinlock());
+
 		auto& processor_info = s_processors[current_id().as_u32()];
 
 		{

kernel/kernel/Thread.cpp

@@ -140,6 +140,8 @@ namespace Kernel
 
 	pid_t Thread::current_tid()
 	{
+		if (Processor::count() == 0)
+			return 0;
 		return Processor::scheduler().current_tid();
 	}