1 files changed, 14 insertions, 4 deletions
diff --git a/include/sibs/BootstrapConnection.hpp b/include/sibs/BootstrapConnection.hpp
index 9e2dfff..8156bd7 100644
--- a/include/sibs/BootstrapConnection.hpp
+++ b/include/sibs/BootstrapConnection.hpp
@@ -4,6 +4,7 @@
 #include "DirectConnection.hpp"
 #include "PubsubKey.hpp"
 #include <mutex>
+#include <atomic>
 
 namespace sibs
 {
@@ -23,6 +24,13 @@ namespace sibs
     // Return false if you want to stop listening on the key
     using BoostrapConnectionListenCallbackFunc = std::function<bool(const DirectConnectionPeer *peer, const void *data, const usize size)>;
 
+    struct SubscribeData
+    {
+        BoostrapConnectionListenCallbackFunc listenCallbackFunc = nullptr;
+        std::vector<std::shared_ptr<DirectConnectionPeer>> peers;
+        i64 listenStartTimeMs = 0;
+    };
+
     struct ListenHandle
     {
         PubsubKey key;
@@ -34,6 +42,7 @@ namespace sibs
     public:
         // Throws BootstrapConnectionException on error
         BootstrapConnection(const Ipv4 &bootstrapAddress);
+        ~BootstrapConnection();
         
         // If we are already listening on the key @pubsubKey then the callback function is overwritten
         ListenHandle listen(const PubsubKey &pubsubKey, BoostrapConnectionListenCallbackFunc callbackFunc);
@@ -46,14 +55,15 @@ namespace sibs
 
         std::vector<std::shared_ptr<DirectConnectionPeer>> getPeers();
     private:
+        ListenHandle listen(const PubsubKey &pubsubKey, BoostrapConnectionListenCallbackFunc callbackFunc, bool registerCallbackFunc);
         void receiveDataFromServer(std::shared_ptr<DirectConnectionPeer> peer, MessageType messageType, const void *data, const usize size);
         void receiveDataFromPeer(std::shared_ptr<DirectConnectionPeer> peer, MessageType messageType, const void *data, const usize size);
     private:
         DirectConnections connections;
         std::shared_ptr<DirectConnectionPeer> serverPeer;
-        PubsubKeyMap<BoostrapConnectionListenCallbackFunc> listenCallbackFuncs;
-        PubsubKeyMap<std::vector<std::shared_ptr<DirectConnectionPeer>>> subscribedPeers;
-        std::recursive_mutex listenerCallbackFuncMutex;
-        std::recursive_mutex subscribedPeersMutex;
+        PubsubKeyMap<SubscribeData> subscribeData;
+        std::recursive_mutex subscribeDataMutex;
+        bool alive;
+        std::atomic_int putThreadCount;
     };
 }