1 files changed, 55 insertions, 41 deletions

diff --git a/kernel/kapi/devices/device_registry.cpp b/kernel/kapi/devices/device_registry.cpp
index f5c20a6d..80ddde65 100644
--- a/kernel/kapi/devices/device_registry.cpp
+++ b/kernel/kapi/devices/device_registry.cpp
@@ -4,6 +4,7 @@
 #include <kapi/system.hpp>
 
 #include <kstd/memory.hpp>
+#include <kstd/mutex.hpp>
 #include <kstd/print.hpp>
 #include <kstd/string.hpp>
 #include <kstd/vector.hpp>
@@ -27,7 +28,7 @@ namespace kapi::devices
       system::panic("[OS:DEV] Device registry has already been initialized");
     }
 
-    instance = device_registry{};
+    instance.emplace();
   }
 
   auto device_registry::get() -> device_registry &
@@ -50,33 +51,38 @@ namespace kapi::devices
     kstd::println("[OS:DEV] Registering device {}", device->name());
 
     auto added = false;
-
-    // TODO: lock the registry
-    auto found = m_devices.find(device->name());
-    if (found != m_devices.end())
     {
-      if (!found->second.expired())
+      auto guard = kstd::lock_guard{m_lock};
+
+      auto found = m_devices.find(device->name());
+      if (found != m_devices.end())
       {
-        added = false;
-      }
+        if (!found->second.expired())
+        {
+          added = false;
+        }
 
-      found->second = device;
-      added = true;
-    }
-    else
-    {
-      added = m_devices.emplace(device->name(), device).second;
+        found->second = device;
+        added = true;
+      }
+      else
+      {
+        added = m_devices.emplace(device->name(), device).second;
+      }
     }
-    // TODO: unlock registry
 
     if (!added)
     {
       return false;
     }
 
-    // TODO: consider if we need to lock the registry observers here, since somebody may add an observer in the meantime
-    // once multiple codepaths can run at the same time.
-    for (auto const & observer : m_observers)
+    auto observers = kstd::vector<kstd::weak_ptr<device_registry_observer>>{};
+    {
+      auto guard = kstd::lock_guard{m_lock};
+      observers = m_observers;
+    }
+
+    for (auto const & observer : observers)
     {
       if (auto locked_observer = observer.lock())
       {
@@ -84,36 +90,43 @@ namespace kapi::devices
       }
     }
 
-    // TODO: lock registry
-    erase_if(m_observers, [](auto const & observer) { return observer.expired(); });
-    // TODO: unlock registry
+    {
+      auto guard = kstd::lock_guard{m_lock};
+      erase_if(m_observers, [](auto const & observer) { return observer.expired(); });
+    }
 
     return true;
   }
 
   auto device_registry::remove(device & device) -> bool
   {
-    // TODO: lock registry
-
-    auto found = m_devices.find(device.name());
-    if (found == m_devices.end())
     {
-      return false;
-    }
+      auto guard = kstd::lock_guard{m_lock};
 
-    auto locked = found->second.lock();
-    if (!locked || locked.get() != &device)
-    {
-      return false;
-    }
+      auto found = m_devices.find(device.name());
+      if (found == m_devices.end())
+      {
+        return false;
+      }
 
-    found->second = kstd::weak_ptr<struct device>{};
+      auto locked = found->second.lock();
+      if (!locked || locked.get() != &device)
+      {
+        return false;
+      }
 
-    // TODO: unlock registry
+      found->second = kstd::weak_ptr<struct device>{};
+    }
 
     kstd::println("[OS:DEV] Unregistering device {}", device.name());
 
-    for (auto const & observer : m_observers)
+    auto observers = kstd::vector<kstd::weak_ptr<device_registry_observer>>{};
+    {
+      auto guard = kstd::lock_guard{m_lock};
+      observers = m_observers;
+    }
+
+    for (auto const & observer : observers)
     {
       if (auto locked_observer = observer.lock())
       {
@@ -121,16 +134,17 @@ namespace kapi::devices
       }
     }
 
-    // TODO: lock registry
-    erase_if(m_observers, [](auto const & observer) { return observer.expired(); });
-    // TODO: unlock registry
+    {
+      auto guard = kstd::lock_guard{m_lock};
+      erase_if(m_observers, [](auto const & observer) { return observer.expired(); });
+    }
 
     return true;
   }
 
   auto device_registry::find(std::string_view name) const -> kstd::shared_ptr<device>
   {
-    // TODO: lock registry
+    auto guard = kstd::lock_guard{m_lock};
 
     auto found = m_devices.find(kstd::string{name});
     if (found == m_devices.end())
@@ -143,7 +157,7 @@ namespace kapi::devices
 
   auto device_registry::all() const -> kstd::vector<kstd::shared_ptr<device>>
   {
-    // TODO: lock registry
+    auto guard = kstd::lock_guard{m_lock};
 
     auto result = kstd::vector<kstd::shared_ptr<device>>{};
 
@@ -160,7 +174,7 @@ namespace kapi::devices
 
   auto device_registry::subscribe(kstd::weak_ptr<device_registry_observer> observer) -> void
   {
-    // TODO: lock registry
+    auto guard = kstd::lock_guard{m_lock};
 
     erase_if(m_observers, [](auto const & observer) { return observer.expired(); });
     m_observers.push_back(observer);