diff --git a/src/Loop.h b/src/Loop.h index b2793fc..3681dad 100644 --- a/src/Loop.h +++ b/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 &&handler) { + void addPostHandler(void *key, fu2::unique_function &&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 &&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 &&handler) { + void addPreHandler(void *key, fu2::unique_function &&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 */ diff --git a/src/LoopData.h b/src/LoopData.h index 5dd4abc..112bae8 100644 --- a/src/LoopData.h +++ b/src/LoopData.h @@ -22,6 +22,7 @@ #include #include #include +#include #include "PerMessageDeflate.h" @@ -38,10 +39,8 @@ private: int currentDeferQueue = 0; std::vector> deferQueues[2]; - fu2::unique_function postHandler, preHandler; - - /* Move over to these later on */ - std::vector> postHandlers; + /* Map from void ptr to handler */ + std::map> postHandlers, preHandlers; public: ~LoopData() { diff --git a/src/WebSocketContextData.h b/src/WebSocketContextData.h index da4419e..a88686f 100644 --- a/src/WebSocketContextData.h +++ b/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 *) 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(); });