New ListenerSet design.

Author: Mysticial

Common/Cpp/CancellableScope.cpp

@@ -39,9 +39,6 @@ void Cancellable::add_cancel_listener(CancelListener& listener){
 void Cancellable::remove_cancel_listener(CancelListener& listener){
     m_impl->m_listeners.remove(listener);
 }
-bool Cancellable::try_add_cancel_listener(CancelListener& listener){
-    return m_impl->m_listeners.try_add(listener);
-}
 
 
 Cancellable::Cancellable()

Common/Cpp/CancellableScope.h

@@ -52,8 +52,6 @@ class Cancellable{
     void add_cancel_listener(CancelListener& listener);
     void remove_cancel_listener(CancelListener& listener);
 
-    bool try_add_cancel_listener(CancelListener& listener);
-
 
 public:
     virtual ~Cancellable();

Common/Cpp/ListenerSet.h

@@ -10,15 +10,20 @@
 #include <map>
 #include <atomic>
 #include "Common/Cpp/Concurrency/SpinLock.h"
-#include "Common/Cpp/LifetimeSanitizer.h"
 
-//#include <iostream>
+#include <iostream>
 //using std::cout;
 //using std::endl;
 
+
 //#define PA_DEBUG_ListenerSet
 
 
+#ifdef PA_DEBUG_ListenerSet
+#include "Common/Cpp/LifetimeSanitizer.h"
+#endif
+
+
 namespace PokemonAutomation{
 
 
@@ -27,29 +32,55 @@ template <typename ListenerType>
 class ListenerSet{
 public:
     bool empty() const{
-        return m_count.load(std::memory_order_relaxed) == 0;
+        return m_count.load(std::memory_order_acquire) == 0;
     }
     size_t count_unique() const{
-        return m_count.load(std::memory_order_relaxed);
+        return m_count.load(std::memory_order_acquire);
     }
 
+    //  Add a new listener. This will never fail unless it throws.
+    //  Deadlocking is not possible since there's only one local lock.
     void add(ListenerType& listener);
+
+    //  Remove a listener. This will deadlock if a listener tries to remove
+    //  from inside a callback.
     void remove(ListenerType& listener);
 
-    bool try_add(ListenerType& listener);
+    //  Same as above, but will return false if it needs to wait.
+    //  This can never deadlock.
     bool try_remove(ListenerType& listener);
 
     template <typename Function, class... Args>
     void run_method(Function function, Args&&... args);
 
+
 private:
     //  Optimization. Keep an atomic version of the count. This will let us
     //  skip the lock when there are no listeners.
     std::atomic<size_t> m_count;
 
     mutable SpinLock m_lock;
-//    mutable std::mutex m_lock;
-    std::map<ListenerType*, size_t> m_listeners;
+
+    struct Node{
+        SpinLock lock;
+        ListenerType& listener;
+        Node* next = nullptr;
+        Node** prevs_next = nullptr;
+
+#ifdef PA_DEBUG_ListenerSet
+        LifetimeSanitizer sanitizer;
+#endif
+
+        Node(ListenerSet& parent, ListenerType& p_listener)
+            : listener(p_listener)
+            , prevs_next(&parent.m_list)
+#ifdef PA_DEBUG_ListenerSet
+            , sanitizer("Node")
+#endif
+        {}
+    };
+    Node* m_list = nullptr;
+    std::map<ListenerType*, Node> m_listeners;
 
 #ifdef PA_DEBUG_ListenerSet
     LifetimeSanitizer m_sanitizer;
@@ -72,42 +103,75 @@ void ListenerSet<ListenerType>::add(ListenerType& listener){
     auto scope = m_sanitizer.check_scope();
 #endif
     WriteSpinLock lg(m_lock);
-    m_listeners[&listener]++;
+    auto ret = m_listeners.emplace(
+        std::piecewise_construct,
+        std::forward_as_tuple(&listener),
+        std::forward_as_tuple(*this, listener)
+        );
+    if (!ret.second){
+        return;
+    }
+    Node& node = ret.first->second;
+#ifdef PA_DEBUG_ListenerSet
+    node.sanitizer.check_usage();
+#endif
+    if (m_list != nullptr){
+        m_list->prevs_next = &node.next;
+    }
+    node.next = m_list;
+    m_list = &node;
     m_count.store(m_listeners.size(), std::memory_order_relaxed);
 }
 template <typename ListenerType>
 void ListenerSet<ListenerType>::remove(ListenerType& listener){
 #ifdef PA_DEBUG_ListenerSet
     auto scope = m_sanitizer.check_scope();
 #endif
-    WriteSpinLock lg(m_lock);
-    auto iter = m_listeners.find(&listener);
-    if (iter == m_listeners.end()){
-        return;
-    }
-    if (--iter->second == 0){
+    while (true){
+        WriteSpinLock lg(m_lock);
+        auto iter = m_listeners.find(&listener);
+        if (iter == m_listeners.end()){
+            return;
+        }
+
+        Node& node = iter->second;
+
+#ifdef PA_DEBUG_ListenerSet
+        node.sanitizer.check_usage();
+#endif
+
+        if (!node.lock.try_acquire_write()){
+            std::cout << "ListenerSet::remove(): Retry inner." << std::endl;
+            continue;
+        }
+
+//        std::cout << "node = " << &node.sanitizer << " : " << &node.prev->sanitizer << " : " << &node.next->sanitizer << std::endl;
+
+        *node.prevs_next = node.next;
+        if (node.next){
+#ifdef PA_DEBUG_ListenerSet
+            node.next->sanitizer.check_usage();
+#endif
+            node.next->prevs_next = node.prevs_next;
+        }
+
+#ifdef PA_DEBUG_ListenerSet
+        node.sanitizer.check_usage();
+#endif
+
         m_listeners.erase(iter);
+        m_count.store(m_listeners.size(), std::memory_order_relaxed);
+        return;
     }
-    m_count.store(m_listeners.size(), std::memory_order_relaxed);
 }
 
 
 
 template <typename ListenerType>
-bool ListenerSet<ListenerType>::try_add(ListenerType& listener){
+bool ListenerSet<ListenerType>::try_remove(ListenerType& listener){
 #ifdef PA_DEBUG_ListenerSet
     auto scope = m_sanitizer.check_scope();
 #endif
-    if (!m_lock.try_acquire_write()){
-        return false;
-    }
-    m_listeners[&listener]++;
-    m_count.store(m_listeners.size(), std::memory_order_relaxed);
-    m_lock.unlock_write();
-    return true;
-}
-template <typename ListenerType>
-bool ListenerSet<ListenerType>::try_remove(ListenerType& listener){
     if (!m_lock.try_acquire_write()){
         return false;
     }
@@ -116,9 +180,30 @@ bool ListenerSet<ListenerType>::try_remove(ListenerType& listener){
         m_lock.unlock_write();
         return true;
     }
-    if (--iter->second == 0){
-        m_listeners.erase(iter);
+
+    Node& node = iter->second;
+    if (!node.lock.try_acquire_write()){
+        std::cout << "ListenerSet::try_remove(): Fail inner." << std::endl;
+        return false;
     }
+
+#ifdef PA_DEBUG_ListenerSet
+    node.sanitizer.check_usage();
+#endif
+
+    *node.prevs_next = node.next;
+    if (node.next){
+#ifdef PA_DEBUG_ListenerSet
+        node.next->sanitizer.check_usage();
+#endif
+        node.next->prevs_next = node.prevs_next;
+    }
+
+#ifdef PA_DEBUG_ListenerSet
+    node.sanitizer.check_usage();
+#endif
+
+    m_listeners.erase(iter);
     m_count.store(m_listeners.size(), std::memory_order_relaxed);
     m_lock.unlock_write();
     return true;
@@ -135,9 +220,37 @@ void ListenerSet<ListenerType>::run_method(Function function, Args&&... args){
     if (empty()){
         return;
     }
-    ReadSpinLock lg(m_lock);
-    for (auto& item : m_listeners){
-        (item.first->*function)(std::forward<Args>(args)...);
+    std::exception_ptr err;
+
+    m_lock.acquire_read();
+
+    Node* node = m_list;
+    while (node){
+        {
+            ReadSpinLock lg(node->lock);
+
+#ifdef PA_DEBUG_ListenerSet
+            node->sanitizer.check_usage();
+#endif
+            m_lock.unlock_read();
+            try{
+                (node->listener.*function)(std::forward<Args>(args)...);
+            }catch (...){
+                if (!err){
+                    err = std::current_exception();
+                }
+            }
+            m_lock.acquire_read();
+        }
+#ifdef PA_DEBUG_ListenerSet
+        node->sanitizer.check_usage();
+#endif
+        node = node->next;
+    }
+
+    m_lock.unlock_read();
+    if (err){
+        std::rethrow_exception(err);
     }
 }
 

SerialPrograms/Source/CommonFramework/VideoPipeline/VideoSession.cpp

@@ -27,9 +27,6 @@ void VideoSession::remove_state_listener(StateListener& listener){
     m_state_listeners.remove(listener);
 }
 
-bool VideoSession::try_add_state_listener(StateListener& listener){
-    return m_state_listeners.try_add(listener);
-}
 bool VideoSession::try_remove_state_listener(StateListener& listener){
     return m_state_listeners.try_remove(listener);
 }

SerialPrograms/Source/CommonFramework/VideoPipeline/VideoSession.h

@@ -58,7 +58,6 @@ class VideoSession
     //  Remove the state listener.
     void remove_state_listener(StateListener& listener);
 
-    bool try_add_state_listener(StateListener& listener);
     bool try_remove_state_listener(StateListener& listener);
 
     //  Implements VideoFeed::add_frame_listener().

SerialPrograms/Source/Controllers/SerialPABotBase/Connection/PABotBase.cpp

@@ -958,17 +958,14 @@ BotBaseMessage PABotBase::wait_for_request(uint64_t seqnum, Cancellable* cancell
 
 
 void PABotBase::cv_wait(Cancellable* cancellable, std::unique_lock<std::mutex>& lg){
-    //  This doesn't work yet. Disable it.
-    if (cancellable == nullptr || true){
+    if (cancellable == nullptr){
         m_cv.wait(lg);
         return;
     }
 
-    //  Only wait if we're able to attach the cancel listener.
-    if (cancellable->try_add_cancel_listener(*this)){
-        m_cv.wait(lg);
-        cancellable->remove_cancel_listener(*this);
-    }
+    cancellable->add_cancel_listener(*this);
+    m_cv.wait(lg);
+    cancellable->remove_cancel_listener(*this);
 }