Make sure to remove post/pre handlers

src/Loop.h

@@ -45,26 +45,16 @@ private:
     static void preCb(us_loop_t *loop) {
         LoopData *loopData = (LoopData *) us_loop_ext(loop);
 
-        if (loopData->preHandler) {
-            loopData->preHandler((Loop *) loop);
-        }
-
-        /* trying this one here */
-        for (auto &f : loopData->postHandlers) {
-            f((Loop *) loop);
+        for (auto &p : loopData->preHandlers) {
+            p.second((Loop *) loop);
         }
     }
 
     static void postCb(us_loop_t *loop) {
         LoopData *loopData = (LoopData *) us_loop_ext(loop);
 
-        /* We should move over to using only these */
-        for (auto &f : loopData->postHandlers) {
-            f((Loop *) loop);
-        }
-
-        if (loopData->postHandler) {
-            loopData->postHandler((Loop *) loop);
+        for (auto &p : loopData->postHandlers) {
+            p.second((Loop *) loop);
         }
     }
 
@@ -119,24 +109,30 @@ public:
         us_loop_free((us_loop_t *) this);
     }
 
-    /* We want to have multiple of these */
-    void addPostHandler(fu2::unique_function<void(Loop *)> &&handler) {
+    void addPostHandler(void *key, fu2::unique_function<void(Loop *)> &&handler) {
         LoopData *loopData = (LoopData *) us_loop_ext((us_loop_t *) this);
 
-        loopData->postHandlers.emplace_back(std::move(handler));
+        loopData->postHandlers.emplace(key, std::move(handler));
     }
 
-    /* Set postCb callback */
-    void setPostHandler(fu2::unique_function<void(Loop *)> &&handler) {
+    /* Bug: what if you remove a handler while iterating them? */
+    void removePostHandler(void *key) {
         LoopData *loopData = (LoopData *) us_loop_ext((us_loop_t *) this);
 
-        loopData->postHandler = std::move(handler);
+        loopData->postHandlers.erase(key);
     }
 
-    void setPreHandler(fu2::unique_function<void(Loop *)> &&handler) {
+    void addPreHandler(void *key, fu2::unique_function<void(Loop *)> &&handler) {
         LoopData *loopData = (LoopData *) us_loop_ext((us_loop_t *) this);
 
-        loopData->preHandler = std::move(handler);
+        loopData->preHandlers.emplace(key, std::move(handler));
+    }
+
+    /* Bug: what if you remove a handler while iterating them? */
+    void removePreHandler(void *key) {
+        LoopData *loopData = (LoopData *) us_loop_ext((us_loop_t *) this);
+
+        loopData->preHandlers.erase(key);
     }
 
     /* Defer this callback on Loop's thread of execution */

src/LoopData.h

@@ -22,6 +22,7 @@
 #include <functional>
 #include <vector>
 #include <mutex>
+#include <map>
 
 #include "PerMessageDeflate.h"
 
@@ -38,10 +39,8 @@ private:
     int currentDeferQueue = 0;
     std::vector<fu2::unique_function<void()>> deferQueues[2];
 
-    fu2::unique_function<void(Loop *)> postHandler, preHandler;
-
-    /* Move over to these later on */
-    std::vector<fu2::unique_function<void(Loop *)>> postHandlers;
+    /* Map from void ptr to handler */
+    std::map<void *, fu2::unique_function<void(Loop *)>> postHandlers, preHandlers;
 
 public:
     ~LoopData() {

src/WebSocketContextData.h

@@ -47,6 +47,12 @@ struct WebSocketContextData {
     /* Each websocket context has a topic tree for pub/sub */
     TopicTree topicTree;
 
+    ~WebSocketContextData() {
+        /* We must unregister any loop post handler here */
+        Loop::get()->removePostHandler(this);
+        Loop::get()->removePreHandler(this);
+    }
+
     WebSocketContextData() : topicTree([this](Subscriber *s, std::string_view data) -> int {
         /* We rely on writing to regular asyncSockets */
         auto *asyncSocket = (AsyncSocket<SSL> *) s->user;
@@ -66,8 +72,13 @@ struct WebSocketContextData {
         /* Reserved, unused */
         return 0;
     }) {
-        /* bug: This should probably happen in both post and pre, esp for libuv */
-        Loop::get()->addPostHandler([this](Loop *loop) {
+        /* We empty for both pre and post just to make sure */
+        Loop::get()->addPostHandler(this, [this](Loop *loop) {
+            /* Commit pub/sub batches every loop iteration */
+            topicTree.drain();
+        });
+
+        Loop::get()->addPreHandler(this, [this](Loop *loop) {
             /* Commit pub/sub batches every loop iteration */
             topicTree.drain();
         });