From 2f52001c6da9103a20b4580596f26628e5f69afa Mon Sep 17 00:00:00 2001
From: Bananymous <bananymousosq@gmail.com>
Date: Thu, 27 Jul 2023 18:36:44 +0300
Subject: [PATCH] Kernel: Locks allow locking after locker is invalid

SpinLock and RecursiveSpinLock will now allow locking after the initial
locker is invalid. This allows us to kill threads even if they are holding
internal locks
---
 kernel/include/kernel/SpinLock.h              |  2 +-
 kernel/include/kernel/Storage/ATAController.h |  1 -
 kernel/kernel/SpinLock.cpp                    | 27 ++++++++++++++-----
 3 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/kernel/include/kernel/SpinLock.h b/kernel/include/kernel/SpinLock.h
index cb0a82a25d..2eaef9ecaf 100644
--- a/kernel/include/kernel/SpinLock.h
+++ b/kernel/include/kernel/SpinLock.h
@@ -19,7 +19,7 @@ namespace Kernel
 		bool is_locked() const;
 
 	private:
-		volatile int m_lock = 0;
+		volatile pid_t m_locker = -1;
 	};
 
 	class RecursiveSpinLock
diff --git a/kernel/include/kernel/Storage/ATAController.h b/kernel/include/kernel/Storage/ATAController.h
index a79796db18..ab8f5dd684 100644
--- a/kernel/include/kernel/Storage/ATAController.h
+++ b/kernel/include/kernel/Storage/ATAController.h
@@ -1,7 +1,6 @@
 #pragma once
 
 #include <kernel/PCI.h>
-#include <kernel/SpinLock.h>
 #include <kernel/Storage/StorageController.h>
 
 namespace Kernel
diff --git a/kernel/kernel/SpinLock.cpp b/kernel/kernel/SpinLock.cpp
index 0d664d73b7..f65b4fde67 100644
--- a/kernel/kernel/SpinLock.cpp
+++ b/kernel/kernel/SpinLock.cpp
@@ -1,24 +1,39 @@
 #include <kernel/Scheduler.h>
 #include <kernel/SpinLock.h>
+#include <kernel/CriticalScope.h>
 
 namespace Kernel
 {
 
 	void SpinLock::lock()
 	{
-		while (__sync_lock_test_and_set(&m_lock, 1))
-			while (m_lock)
-				Scheduler::get().reschedule();
+		pid_t tid = Scheduler::current_tid();
+		while (true)
+		{
+			{
+				CriticalScope _;
+				ASSERT(m_locker != tid);
+				if (m_locker == -1 || !Scheduler::is_valid_tid(m_locker))
+				{
+					m_locker = tid;
+					break;
+				}
+			}
+			Scheduler::get().reschedule();
+		}
 	}
 
 	void SpinLock::unlock()
 	{
-		__sync_lock_release(&m_lock);
+		CriticalScope _;
+		ASSERT(m_locker == Scheduler::current_tid());
+		m_locker = -1;
 	}
 
 	bool SpinLock::is_locked() const
 	{
-		return m_lock;
+		CriticalScope _;
+		return m_locker != -1;
 	}
 
 	void RecursiveSpinLock::lock()
@@ -33,7 +48,7 @@ namespace Kernel
 				m_lock_depth++;
 				break;
 			}
-			if (m_locker == -1)
+			if (m_locker == -1 || !Scheduler::is_valid_tid(m_locker))
 			{
 				m_locker = tid;
 				m_lock_depth = 1;