refactor: EventManager and Publisher classes with subscription handling capabilities

Author: gilmar-sales

include/Freyr/Core/EventManager.hpp

@@ -4,78 +4,207 @@
 
 namespace FREYR_NAMESPACE
 {
+    template <typename TEvent>
+    class Publisher;
+
+    struct ListenerHandle
+    {
+        size_t id { 0 };
+        bool   IsValid() const { return id != 0; }
+    };
+
     class IPublisher
     {
       public:
-        virtual ~IPublisher() = default;
+        virtual ~IPublisher()                 = default;
+        virtual void ClearInactiveListeners() = 0;
     };
 
     template <typename TEvent>
     class Publisher final : public IPublisher
     {
+      private:
+        struct Listener
+        {
+            fr::function<void(const TEvent&)> callback;
+            std::weak_ptr<ListenerHandle>     handle;
+
+            Listener(fr::function<void(const TEvent&)>&& cb, Ref<ListenerHandle> handle) :
+                callback(std::move(cb)), handle(handle)
+            {
+            }
+        };
+
+        mutable std::shared_mutex mMutex;
+        std::vector<Listener>     mListeners;
+        std::atomic<size_t>       mNextId { 1 };
+        std::atomic<bool>         mNeedsCleanup { false };
+
+        std::mutex            mPendingMutex;
+        std::vector<Listener> mPendingListeners;
+
       public:
         ~Publisher() override = default;
 
-        void Subscribe(auto&& listener)
+        [[nodiscard]] Ref<ListenerHandle> Subscribe(auto&& listener)
         {
-            mListeners.emplace_back(std::move(listener));
+            const size_t id = mNextId.fetch_add(1, std::memory_order_relaxed);
+
+            auto handle = skr::MakeRef<ListenerHandle>(id);
+            {
+                std::lock_guard lock(mPendingMutex);
+                mPendingListeners.emplace_back(
+                    fr::function<void(const TEvent&)>(std::forward<decltype(listener)>(listener)),
+                    handle);
+            }
+
+            return handle;
         }
 
-        void Publish(TEvent event)
+        void Publish(const TEvent& event)
         {
-            for (auto i = 0; i < mListeners.size(); i++)
+            MergePendingListeners();
+
+            std::shared_lock lock(mMutex);
+
+            for (Listener& listener : mListeners)
+            {
+                if (!listener.handle.expired())
+                {
+                    listener.callback(event);
+                }
+            }
+
+            // Clean up after publishing if needed
+            if (mNeedsCleanup.load(std::memory_order_acquire))
             {
-                mListeners[i](event);
+                lock.unlock();
+                ClearInactiveListeners();
             }
         }
 
-    private:
-        std::vector<fr::function<void(TEvent)>> mListeners;
+        // Overload for rvalue events
+        void Publish(TEvent&& event) { Publish(static_cast<const TEvent&>(event)); }
+
+        void ClearInactiveListeners() override
+        {
+            std::unique_lock lock(mMutex);
+
+            mListeners.erase(std::remove_if(mListeners.begin(),
+                                            mListeners.end(),
+                                            [](const Listener& l) { return l.handle.expired(); }),
+                             mListeners.end());
+
+            mNeedsCleanup.store(false, std::memory_order_release);
+        }
+
+        size_t ListenerCount() const
+        {
+            std::shared_lock lock(mMutex);
+            return std::count_if(mListeners.begin(), mListeners.end(), [](const Listener& l) {
+                return !l.handle.expired();
+            });
+        }
+
+      private:
+        void MergePendingListeners()
+        {
+            std::vector<Listener> toMerge;
+
+            {
+                std::lock_guard lock(mPendingMutex);
+                if (mPendingListeners.empty())
+                    return;
+                toMerge = std::move(mPendingListeners);
+                mPendingListeners.clear();
+            }
+
+            if (!toMerge.empty())
+            {
+                std::unique_lock lock(mMutex);
+                mListeners.reserve(mListeners.size() + toMerge.size());
+                mListeners.insert(mListeners.end(),
+                                  std::make_move_iterator(toMerge.begin()),
+                                  std::make_move_iterator(toMerge.end()));
+            }
+        }
     };
 
     class EventManager
     {
+      private:
+        mutable std::shared_mutex                mMutex;
+        std::vector<std::unique_ptr<IPublisher>> mPublishers;
+        std::unordered_map<size_t, size_t>       mEventIdToIndex;
+
       public:
-        EventManager() : mPublishers(256) {}
-        ~EventManager()
+        EventManager() { mPublishers.reserve(256); }
+
+        ~EventManager() = default;
+
+        template <typename T>
+            requires IsEvent<T>
+        [[nodiscard]] Ref<ListenerHandle> Subscribe(auto&& listener)
         {
-            for (const auto publisher : mPublishers)
-            {
-                delete publisher;
-            }
+            return GetPublisher<T>()->Subscribe(std::forward<decltype(listener)>(listener));
         }
 
         template <typename T>
             requires IsEvent<T>
-        void Subscribe(auto&& listener)
+        void Send(const T& event)
         {
-            GetPublisher<T>()->Subscribe(listener);
+            GetPublisher<T>()->Publish(event);
         }
 
         template <typename T>
             requires IsEvent<T>
-        void Send(T event)
+        void Send(T&& event)
         {
-            GetPublisher<T>()->Publish(event);
+            GetPublisher<T>()->Publish(std::forward<T>(event));
+        }
+
+        void Cleanup()
+        {
+            std::shared_lock lock(mMutex);
+            for (auto& publisher : mPublishers)
+            {
+                if (publisher)
+                {
+                    publisher->ClearInactiveListeners();
+                }
+            }
         }
 
       private:
         template <typename T>
             requires IsEvent<T>
         Publisher<T>* GetPublisher()
         {
-            if (GetEventId<T>() + 1 > mPublishers.size())
-                mPublishers.resize(GetEventId<T>() + 16);
+            const size_t eventId = GetEventId<T>();
 
-            if (!mPublishers[GetEventId<T>()])
             {
-                mPublishers[GetEventId<T>()] = new Publisher<T>();
+                std::shared_lock lock(mMutex);
+                auto             it = mEventIdToIndex.find(eventId);
+                if (it != mEventIdToIndex.end())
+                {
+                    return static_cast<Publisher<T>*>(mPublishers[it->second].get());
+                }
             }
 
-            return static_cast<Publisher<T>*>(mPublishers[GetEventId<T>()]);
-        }
+            std::unique_lock lock(mMutex);
 
-        std::vector<IPublisher*> mPublishers;
+            auto it = mEventIdToIndex.find(eventId);
+            if (it != mEventIdToIndex.end())
+            {
+                return static_cast<Publisher<T>*>(mPublishers[it->second].get());
+            }
+
+            size_t index = mPublishers.size();
+            mPublishers.push_back(std::make_unique<Publisher<T>>());
+            mEventIdToIndex[eventId] = index;
+
+            return static_cast<Publisher<T>*>(mPublishers[index].get());
+        }
     };
 
-} // namespace FREYR_NAMESPACE
+} // namespace FREYR_NAMESPACE