kapi: implement a simple, tracked mutex

6 files changed, 143 insertions, 0 deletions

diff --git a/arch/x86_64/kapi/cpu.cpp b/arch/x86_64/kapi/cpu.cpp
index c45703e5..100f53df 100644
--- a/arch/x86_64/kapi/cpu.cpp
+++ b/arch/x86_64/kapi/cpu.cpp
@@ -82,4 +82,10 @@ namespace kapi::cpu
     return core_index > 0;
   }
 
+  auto current_id() -> id
+  {
+    // TODO: implement actual ID reading once APs are being set up.
+    return id{.value = 0};
+  }
+
 }  // namespace kapi::cpu
diff --git a/kapi/kapi/cpu.hpp b/kapi/kapi/cpu.hpp
index 41e00e40..9a50f67e 100644
--- a/kapi/kapi/cpu.hpp
+++ b/kapi/kapi/cpu.hpp
@@ -7,6 +7,7 @@
 #include <kstd/memory.hpp>
 
 #include <cstdint>
+#include <limits>
 
 namespace kapi::cpu
 {
@@ -62,6 +63,17 @@ namespace kapi::cpu
     bool is_user_mode{};
   };
 
+  //! A stable id for a single logical CPU core.
+  struct id
+  {
+    constexpr auto friend operator==(id const &, id const &) noexcept -> bool = default;
+
+    std::uint64_t value{};
+  };
+
+  //! The invalid CPU id.
+  constexpr auto static invalid_id = id{std::numeric_limits<std::uint64_t>::max()};
+
   //! @}
 
   //! @addtogroup kapi-cpu-kernel-defined
@@ -94,6 +106,11 @@ namespace kapi::cpu
   //! @return true iff. the CPU topology was discovered successfully, false otherwise.
   auto discover_topology() -> bool;
 
+  //! Get the current CPU id.
+  //!
+  //! @return The id of the CPU core currently executing the call to this function.
+  [[nodiscard]] auto current_id() -> id;
+
   //! @}
 
 }  // namespace kapi::cpu
@@ -132,4 +149,13 @@ struct kstd::formatter<enum kapi::cpu::exception::type>
   }
 };
 
+template<>
+struct kstd::formatter<kapi::cpu::id> : kstd::formatter<std::uint64_t>
+{
+  constexpr auto format(kapi::cpu::id const & id, kstd::format_context & ctx) const -> void
+  {
+    kstd::formatter<std::uint64_t>::format(id.value, ctx);
+  }
+};
+
 #endif
diff --git a/kapi/kapi/tracked_mutex.hpp b/kapi/kapi/tracked_mutex.hpp
new file mode 100644
index 00000000..11b99e12
--- /dev/null
+++ b/kapi/kapi/tracked_mutex.hpp
@@ -0,0 +1,48 @@
+#ifndef TEACHOS_KAPI_TRACKED_MUTEX_HPP
+#define TEACHOS_KAPI_TRACKED_MUTEX_HPP
+
+#include <kapi/cpu.hpp>
+
+#include <atomic>
+
+namespace kapi
+{
+
+  //! @addtogroup kapi-kernel-defined
+  //! @{
+
+  struct tracked_mutex
+  {
+    constexpr tracked_mutex() = default;
+    tracked_mutex(tracked_mutex const &) = delete;
+    tracked_mutex(tracked_mutex &&) = delete;
+    auto operator=(tracked_mutex const &) -> tracked_mutex & = delete;
+    auto operator=(tracked_mutex &&) -> tracked_mutex & = delete;
+
+    //! Lock this mutex, potentially blocking if it already locked.
+    //!
+    //! @warning This function will panic if a recursive lock on a single core is detected!
+    auto lock() -> void;
+
+    //! Try to lock this mutex.
+    //!
+    //! @note This function will not panic if a recursive lock on a single core is detected, but rather it will return
+    //! false.
+    //!
+    //! @return true iff. the mutex was locked, false otherwise.
+    [[nodiscard]] auto try_lock() -> bool;
+
+    //! Unlock this mutex.
+    //!
+    //! @warning This function will panic if the core executing the call does not own the lock on this mutex!
+    auto unlock() -> void;
+
+  private:
+    std::atomic<kapi::cpu::id> m_owner{kapi::cpu::invalid_id};
+  };
+
+  //! @}
+
+}  // namespace kapi
+
+#endif
\ No newline at end of file
diff --git a/kernel/CMakeLists.txt b/kernel/CMakeLists.txt
index d22ad107..a97c1e4a 100644
--- a/kernel/CMakeLists.txt
+++ b/kernel/CMakeLists.txt
@@ -24,6 +24,7 @@ target_sources("kernel_lib" PRIVATE
     "kapi/interrupts.cpp"
     "kapi/memory.cpp"
     "kapi/system.cpp"
+    "kapi/tracked_mutex.cpp"
 
     # Kernel-defined KSTD Implementation
     "kstd/os.cpp"
diff --git a/kernel/kapi/tracked_mutex.cpp b/kernel/kapi/tracked_mutex.cpp
new file mode 100644
index 00000000..6608169c
--- /dev/null
+++ b/kernel/kapi/tracked_mutex.cpp
@@ -0,0 +1,55 @@
+#include <kapi/tracked_mutex.hpp>
+
+#include <kapi/cpu.hpp>
+#include <kapi/system.hpp>
+
+#include <atomic>
+
+namespace kapi
+{
+
+  auto tracked_mutex::lock() -> void
+  {
+    auto self = kapi::cpu::current_id();
+    auto expected = kapi::cpu::invalid_id;
+
+    while (!m_owner.compare_exchange_weak(expected, self, std::memory_order::acquire, std::memory_order::relaxed))
+    {
+      if (expected == self)
+      {
+        system::panic("[OS] CPU {} tried to reacquire a mutex it already holds!", self);
+      }
+      expected = kapi::cpu::invalid_id;
+    }
+  }
+
+  auto tracked_mutex::try_lock() -> bool
+  {
+    auto self = kapi::cpu::current_id();
+    auto expected = kapi::cpu::invalid_id;
+
+    if (m_owner.compare_exchange_strong(expected, self, std::memory_order::acquire, std::memory_order::relaxed))
+    {
+      return true;
+    }
+
+    if (expected == self)
+    {
+      system::panic("[OS] CPU {} tried to reacquire a mutex it already holds!", self);
+    }
+
+    return false;
+  }
+
+  auto tracked_mutex::unlock() -> void
+  {
+    auto self = kapi::cpu::current_id();
+    auto expected = self;
+
+    if (!m_owner.compare_exchange_strong(expected, kapi::cpu::invalid_id, std::memory_order::release))
+    {
+      system::panic("[OS] CPU {} released a mutex it did not hold!", self);
+    }
+  }
+
+}  // namespace kapi
\ No newline at end of file
diff --git a/kernel/kernel/test_support/kapi/cpu.cpp b/kernel/kernel/test_support/kapi/cpu.cpp
index 5d956331..a82dc348 100644
--- a/kernel/kernel/test_support/kapi/cpu.cpp
+++ b/kernel/kernel/test_support/kapi/cpu.cpp
@@ -3,7 +3,9 @@
 #include <kapi/cpu.hpp>
 
 #include <atomic>
+#include <functional>
 #include <stdexcept>
+#include <thread>
 
 namespace
 {
@@ -34,6 +36,11 @@ namespace kapi::cpu
     return true;
   }
 
+  auto current_id() -> id
+  {
+    return id{.value = std::hash<std::thread::id>{}(std::this_thread::get_id())};
+  }
+
 }  // namespace kapi::cpu
 
 namespace kernel::tests::cpu