From ff83f967d868d30bdc47cac62cb2903d2f9610d9 Mon Sep 17 00:00:00 2001
From: Bananymous <bananymousosq@gmail.com>
Date: Mon, 29 May 2023 19:38:09 +0300
Subject: [PATCH] Kernel: Make RecursiveSpinLock thread safe

also SpinLock is now implemented with gcc builtins
---
 kernel/CMakeLists.txt            |  1 -
 kernel/arch/x86_64/SpinLock.S    | 17 -----------
 kernel/include/kernel/SpinLock.h |  4 +--
 kernel/kernel/SpinLock.cpp       | 52 ++++++++++++++++++++------------
 4 files changed, 35 insertions(+), 39 deletions(-)
 delete mode 100644 kernel/arch/x86_64/SpinLock.S

diff --git a/kernel/CMakeLists.txt b/kernel/CMakeLists.txt
index ec88a26b25..fdc0d94955 100644
--- a/kernel/CMakeLists.txt
+++ b/kernel/CMakeLists.txt
@@ -69,7 +69,6 @@ if("${BANAN_ARCH}" STREQUAL "x86_64")
 		arch/x86_64/IDT.cpp
 		arch/x86_64/interrupts.S
 		arch/x86_64/MMU.cpp
-		arch/x86_64/SpinLock.S
 		arch/x86_64/Thread.S
 	)
 elseif("${BANAN_ARCH}" STREQUAL "i386")
diff --git a/kernel/arch/x86_64/SpinLock.S b/kernel/arch/x86_64/SpinLock.S
deleted file mode 100644
index ad27179bc8..0000000000
--- a/kernel/arch/x86_64/SpinLock.S
+++ /dev/null
@@ -1,17 +0,0 @@
-.global spinlock_lock_asm
-spinlock_lock_asm:
-	lock; btsq $0, (%rdi)
-	jnc .done
-.retry:
-	pause
-	testq $1, (%rdi)
-	jne .retry
-	lock; btsq $0, (%rdi)
-	jc .retry
-.done:
-	ret
-
-.global spinlock_unlock_asm
-spinlock_unlock_asm:
-	movl $0, (%rdi)
-	ret
\ No newline at end of file
diff --git a/kernel/include/kernel/SpinLock.h b/kernel/include/kernel/SpinLock.h
index d3d2471a14..cb0a82a25d 100644
--- a/kernel/include/kernel/SpinLock.h
+++ b/kernel/include/kernel/SpinLock.h
@@ -19,7 +19,7 @@ namespace Kernel
 		bool is_locked() const;
 
 	private:
-		int m_lock = 0;
+		volatile int m_lock = 0;
 	};
 
 	class RecursiveSpinLock
@@ -34,7 +34,7 @@ namespace Kernel
 		bool is_locked() const;
 
 	private:
-		pid_t m_locker = 0;
+		pid_t m_locker = -1;
 		uint32_t m_lock_depth = 0;
 		SpinLock m_lock;
 	};
diff --git a/kernel/kernel/SpinLock.cpp b/kernel/kernel/SpinLock.cpp
index fa6d55cf39..bb65372440 100644
--- a/kernel/kernel/SpinLock.cpp
+++ b/kernel/kernel/SpinLock.cpp
@@ -4,17 +4,16 @@
 namespace Kernel
 {
 
-	extern "C" void spinlock_lock_asm(int*);
-	extern "C" void spinlock_unlock_asm(int*);
-
 	void SpinLock::lock()
 	{
-		spinlock_lock_asm(&m_lock);
+		while (__sync_lock_test_and_set(&m_lock, 1))
+			while (m_lock)
+				__builtin_ia32_pause();
 	}
 
 	void SpinLock::unlock()
 	{
-		spinlock_unlock_asm(&m_lock);
+		__sync_lock_release(&m_lock);
 	}
 
 	bool SpinLock::is_locked() const
@@ -24,35 +23,50 @@ namespace Kernel
 
 	void RecursiveSpinLock::lock()
 	{
-		// FIXME: is this thread safe?
-		if (m_locker == Scheduler::current_tid())
-		{
-			m_lock_depth++;
-		}
-		else
+		pid_t tid = Scheduler::current_tid();
+
+		while (true)
 		{
+			// Wait for us to be the locker or the lock being free
+			while (m_locker != -1 && m_locker != tid)
+				__builtin_ia32_pause();
+
 			m_lock.lock();
-			ASSERT(m_locker == 0);
-			m_locker = Scheduler::current_tid();
-			m_lock_depth = 1;
+			if (m_locker == tid)
+			{
+				m_lock_depth++;
+				break;
+			}
+			if (m_locker == -1)
+			{
+				m_locker = tid;
+				m_lock_depth = 1;
+				break;
+			}
+			m_lock.unlock();
 		}
+
+		m_lock.unlock();
 	}
 
 	void RecursiveSpinLock::unlock()
 	{
+		m_lock.lock();
+		
 		ASSERT(m_lock_depth > 0);
+		ASSERT(m_locker == Scheduler::current_tid());
 
 		m_lock_depth--;
+
 		if (m_lock_depth == 0)
-		{
-			m_locker = 0;
-			m_lock.unlock();
-		}
+			m_locker = -1;
+		
+		m_lock.unlock();
 	}
 
 	bool RecursiveSpinLock::is_locked() const
 	{
-		return m_lock.is_locked();
+		return m_locker != -1;
 	}
 
 }
\ No newline at end of file