From 15a17f3c136a22fa26d1624c270c0576a4fb2611 Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <2769109+epoberezkin@users.noreply.github.com>
Date: Sat, 2 Jul 2022 10:13:06 +0100
Subject: [PATCH] core: subscribe to all connections concurrently (#770)

---
 src/Simplex/Chat.hs            | 23 ++++++++++++-----------
 src/Simplex/Chat/Controller.hs |  2 +-
 2 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/src/Simplex/Chat.hs b/src/Simplex/Chat.hs
index ee7337367a..01d24ee4cc 100644
--- a/src/Simplex/Chat.hs
+++ b/src/Simplex/Chat.hs
@@ -146,24 +146,26 @@ newChatController chatStore user cfg@ChatConfig {agentConfig = aCfg, tbqSize, de
           pure ss {smp = fromMaybe defaultSMPServers $ nonEmpty userSmpServers}
         _ -> pure ss
 
-runChatController :: (MonadUnliftIO m, MonadReader ChatController m) => User -> Bool -> m ()
-runChatController user subConns = race_ notificationSubscriber $ agentSubscriber user subConns
-
 startChatController :: (MonadUnliftIO m, MonadReader ChatController m) => User -> Bool -> m (Async ())
 startChatController user subConns = do
   asks smpAgent >>= resumeAgentClient
   s <- asks agentAsync
-  readTVarIO s >>= maybe (start s) pure
+  readTVarIO s >>= maybe (start s) (pure . fst)
   where
     start s = do
-      a <- async $ runChatController user subConns
-      atomically . writeTVar s $ Just a
-      pure a
+      a1 <- async $ race_ notificationSubscriber agentSubscriber
+      a2 <-
+        if subConns
+          then Just <$> async (subscribeUserConnections subscribeConnection user)
+          else pure Nothing
+      atomically . writeTVar s $ Just (a1, a2)
+      pure a1
 
 stopChatController :: MonadUnliftIO m => ChatController -> m ()
 stopChatController ChatController {smpAgent, agentAsync = s} = do
   disconnectAgentClient smpAgent
-  readTVarIO s >>= mapM_ uninterruptibleCancel >> atomically (writeTVar s Nothing)
+  readTVarIO s >>= mapM_ (\(a1, a2) -> uninterruptibleCancel a1 >> mapM_ uninterruptibleCancel a2)
+  atomically (writeTVar s Nothing)
 
 withLock :: MonadUnliftIO m => TMVar () -> m a -> m a
 withLock lock =
@@ -1000,11 +1002,10 @@ acceptContactRequest User {userId, profile} UserContactRequest {agentInvitationI
   connId <- withAgent $ \a -> acceptContact a invId . directMessage $ XInfo profile
   withStore' $ \db -> createAcceptedContact db userId connId cName profileId p userContactLinkId xContactId
 
-agentSubscriber :: (MonadUnliftIO m, MonadReader ChatController m) => User -> Bool -> m ()
-agentSubscriber user subConns = do
+agentSubscriber :: (MonadUnliftIO m, MonadReader ChatController m) => m ()
+agentSubscriber = do
   q <- asks $ subQ . smpAgent
   l <- asks chatLock
-  when subConns $ subscribeUserConnections subscribeConnection user
   forever $ do
     (_, connId, msg) <- atomically $ readTBQueue q
     u <- readTVarIO =<< asks currentUser
diff --git a/src/Simplex/Chat/Controller.hs b/src/Simplex/Chat/Controller.hs
index 1a3980cc93..b3760fcae7 100644
--- a/src/Simplex/Chat/Controller.hs
+++ b/src/Simplex/Chat/Controller.hs
@@ -73,7 +73,7 @@ data ChatController = ChatController
     activeTo :: TVar ActiveTo,
     firstTime :: Bool,
     smpAgent :: AgentClient,
-    agentAsync :: TVar (Maybe (Async ())),
+    agentAsync :: TVar (Maybe (Async (), Maybe (Async ()))),
     chatStore :: SQLiteStore,
     chatStoreChanged :: TVar Bool, -- if True, chat should be fully restarted
     idsDrg :: TVar ChaChaDRG,