From d9a1e8a980a6a6d68db8eb66bcb0459bbea6ffd8 Mon Sep 17 00:00:00 2001
From: apio <blobs.trading@gmail.com>
Date: Wed, 10 May 2023 21:58:30 +0200
Subject: [PATCH] kernel: Add a KMutex class and use that for ATA::Controller
 locking

---
 kernel/src/arch/x86_64/disk/ATA.h |  3 +-
 kernel/src/lib/KMutex.h           | 78 +++++++++++++++++++++++++++++++
 2 files changed, 80 insertions(+), 1 deletion(-)
 create mode 100644 kernel/src/lib/KMutex.h

diff --git a/kernel/src/arch/x86_64/disk/ATA.h b/kernel/src/arch/x86_64/disk/ATA.h
index ad46feef..03d7cbd7 100644
--- a/kernel/src/arch/x86_64/disk/ATA.h
+++ b/kernel/src/arch/x86_64/disk/ATA.h
@@ -1,6 +1,7 @@
 #pragma once
 #include "arch/PCI.h"
 #include "fs/devices/DeviceRegistry.h"
+#include "lib/KMutex.h"
 #include <luna/Atomic.h>
 #include <luna/SharedPtr.h>
 
@@ -41,7 +42,7 @@ namespace ATA
         Channel m_primary_channel;
         Channel m_secondary_channel;
 
-        Atomic<int> m_hwlock = 0;
+        KMutex<100> m_lock;
     };
 
 }
diff --git a/kernel/src/lib/KMutex.h b/kernel/src/lib/KMutex.h
new file mode 100644
index 00000000..d55a947e
--- /dev/null
+++ b/kernel/src/lib/KMutex.h
@@ -0,0 +1,78 @@
+#pragma once
+#include "Log.h"
+#include "arch/CPU.h"
+#include "thread/Scheduler.h"
+#include "thread/Thread.h"
+#include <luna/CircularQueue.h>
+
+template <usize ConcurrentThreads> class KMutex
+{
+  public:
+    void lock()
+    {
+        int expected = 0;
+        while (!m_lock.compare_exchange_strong(expected, 1))
+        {
+            expected = 0;
+            auto* current = Scheduler::current();
+
+            // We cannot be interrupted between these functions, otherwise we might never exit the loop
+            CPU::disable_interrupts();
+            bool ok = m_blocked_threads.try_push(current);
+            if (!ok) kernel_sleep(10);
+            else
+                kernel_wait_for_event();
+            CPU::enable_interrupts();
+        }
+    };
+
+    void unlock()
+    {
+        int expected = 1;
+        if (!m_lock.compare_exchange_strong(expected, 0))
+        {
+            kwarnln("KMutex::unlock() called on an unlocked lock with value %d", expected);
+        }
+
+        Thread* blocked;
+        if (m_blocked_threads.try_pop(blocked)) blocked->wake_up();
+    }
+
+    bool try_lock()
+    {
+        int expected = 0;
+        return m_lock.compare_exchange_strong(expected, 1);
+    }
+
+  private:
+    CircularQueue<Thread*, ConcurrentThreads> m_blocked_threads;
+    Atomic<int> m_lock;
+};
+
+template <usize ConcurrentThreads> class ScopedKMutexLock
+{
+  public:
+    ScopedKMutexLock(KMutex<ConcurrentThreads>& lock)
+    {
+        m_lock = lock;
+        m_lock.lock();
+    }
+
+    ~ScopedKMutexLock()
+    {
+        if (!m_taken_over) m_lock.unlock();
+    }
+
+    ScopedKMutexLock(const ScopedKMutexLock&) = delete;
+    ScopedKMutexLock(ScopedKMutexLock&&) = delete;
+
+    KMutex<ConcurrentThreads>& take_over()
+    {
+        m_taken_over = true;
+        return m_lock;
+    }
+
+  private:
+    KMutex<ConcurrentThreads>& m_lock;
+    bool m_taken_over { false };
+};