Merge pull request #14148 from jordan-woyak/hookable-event-add-remove-inside-trigger

HookableEvent: Allow hooks to be added and removed from within a Trigger callback.

1 files changed, 50 insertions, 15 deletions

diff --git a/Source/Core/Common/HookableEvent.h b/Source/Core/Common/HookableEvent.h
index 58a07c1fe9..906f3e225b 100644
--- a/Source/Core/Common/HookableEvent.h
+++ b/Source/Core/Common/HookableEvent.h
@@ -3,6 +3,7 @@
 
 #pragma once
 
+#include <algorithm>
 #include <functional>
 #include <memory>
 #include <mutex>
@@ -49,37 +50,59 @@ public:
   using CallbackType = std::function<void(CallbackArgs...)>;
 
   // Returns a handle that will unregister the listener when destroyed.
-  // Note: Attempting to add/remove hooks of the event within the callback itself will NOT work.
   [[nodiscard]] EventHook Register(CallbackType callback)
   {
     DEBUG_LOG_FMT(COMMON, "Registering event hook handler");
-    auto handle = std::make_unique<HookImpl>(m_storage, std::move(callback));
-
     std::lock_guard lg(m_storage->listeners_mutex);
-    m_storage->listeners.push_back(handle.get());
-    return handle;
+
+    auto& new_listener =
+        m_storage->listeners.emplace_back(std::make_unique<Listener>(std::move(callback)));
+
+    return std::make_unique<HookImpl>(m_storage, new_listener.get());
   }
 
+  // Invokes all registered callbacks.
+  // Hooks added from within a callback will be invoked.
+  // Hooks removed from within a callback will be skipped,
+  //  but destruction of the hook's callback will be delayed until Trigger() completes.
   void Trigger(const CallbackArgs&... args)
   {
     std::lock_guard lg(m_storage->listeners_mutex);
-    for (auto* const handle : m_storage->listeners)
-      std::invoke(handle->callback, args...);
+    m_storage->is_triggering = true;
+
+    // Avoiding an actual iterator because the container may be modified.
+    for (std::size_t i = 0; i != m_storage->listeners.size(); ++i)
+    {
+      auto& listener = m_storage->listeners[i];
+
+      if (listener->is_pending_removal)
+        continue;
+
+      std::invoke(listener->callback, args...);
+    }
+
+    m_storage->is_triggering = false;
+    std::erase_if(m_storage->listeners, std::mem_fn(&Listener::is_pending_removal));
   }
 
 private:
-  struct HookImpl;
+  struct Listener
+  {
+    const CallbackType callback;
+    bool is_pending_removal{};
+  };
 
   struct Storage
   {
-    std::mutex listeners_mutex;
-    std::vector<HookImpl*> listeners;
+    std::recursive_mutex listeners_mutex;
+    std::vector<std::unique_ptr<Listener>> listeners;
+    bool is_triggering{};
   };
 
   struct HookImpl final : HookBase
   {
-    HookImpl(const std::shared_ptr<Storage> storage, CallbackType func)
-        : weak_storage{storage}, callback{std::move(func)}
+    HookImpl(std::weak_ptr<Storage> storage, Listener* listener)
+        : weak_storage{std::move(storage)}, listener_ptr{listener}
     {
     }
 
@@ -95,11 +118,23 @@ private:
       DEBUG_LOG_FMT(COMMON, "Removing event hook handler");
 
       std::lock_guard lg(storage->listeners_mutex);
-      std::erase(storage->listeners, this);
+
+      if (storage->is_triggering)
+      {
+        // Just mark our listener for removal.
+        // Trigger() will erase it for us.
+        listener_ptr->is_pending_removal = true;
+      }
+      else
+      {
+        // Remove our listener.
+        storage->listeners.erase(std::ranges::find_if(
+            storage->listeners, [&](auto& ptr) { return ptr.get() == listener_ptr; }));
+      }
     }
 
-    std::weak_ptr<Storage> weak_storage;
-    const CallbackType callback;
+    const std::weak_ptr<Storage> weak_storage;
+    Listener* const listener_ptr;  // "owned" by the above Storage.
   };
 
   // shared_ptr storage allows hooks to forget their connection if they outlive the event itself.