From adde6134929d0f83d134c701660344d037bf731f Mon Sep 17 00:00:00 2001
From: spaced4ndy <8711996+spaced4ndy@users.noreply.github.com>
Date: Fri, 5 Jul 2024 21:16:12 +0400
Subject: [PATCH] group sessions

---
 src/Simplex/Messaging/Agent/Client.hs  | 10 ++++++++--
 tests/AgentTests/FunctionalAPITests.hs | 11 +++--------
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/src/Simplex/Messaging/Agent/Client.hs b/src/Simplex/Messaging/Agent/Client.hs
index d8da2786f..6340e5c3f 100644
--- a/src/Simplex/Messaging/Agent/Client.hs
+++ b/src/Simplex/Messaging/Agent/Client.hs
@@ -938,9 +938,11 @@ reconnectSMPServerClients c = do
   -- 3. close clients
   mapM_ (liftIO . forkIO . closeClient_ c) clients
   -- 4. resubscribe pending subscriptions
+  mode <- liftIO $ getSessionMode c
   pending <- readTVarIO (getRcvQueues $ pendingSubs c)
-  forM_ (M.toList pending) $ \((userId, srv, rId), _) ->
-    resubscribeSMPSession c (userId, srv, Just rId)
+  -- Group transport sessions to avoid multiple UP events in case session mode is TSMUser
+  let tSessions = queuesToSessions pending mode
+  forM_ tSessions $ \tSess -> resubscribeSMPSession c tSess
   where
     groupConnsByServer :: Map (UserId, SMPServer, RecipientId) RcvQueue -> Map SMPServer [ConnId]
     groupConnsByServer = foldl' insertConnId M.empty
@@ -950,6 +952,10 @@ reconnectSMPServerClients c = do
           M.insertWith (<>) server [connId] acc
     notifyDOWN :: SMPServer -> [ConnId] -> IO ()
     notifyDOWN server connIds = atomically $ writeTBQueue (subQ c) ("", "", AEvt SAENone (DOWN server connIds))
+    queuesToSessions :: Map (UserId, SMPServer, RecipientId) RcvQueue -> TransportSessionMode -> Set SMPTransportSession
+    queuesToSessions qs mode = case mode of
+      TSMEntity -> M.foldrWithKey (\(userId, srv, rId) _ acc -> S.insert (userId, srv, Just rId) acc) S.empty qs
+      TSMUser -> M.foldrWithKey (\(userId, srv, _) _ acc -> S.insert (userId, srv, Nothing) acc) S.empty qs
 
 reconnectSMPServer :: AgentClient -> UserId -> SMPServer -> IO ()
 reconnectSMPServer c userId srv = do
diff --git a/tests/AgentTests/FunctionalAPITests.hs b/tests/AgentTests/FunctionalAPITests.hs
index 9cf7c2793..8516f16a9 100644
--- a/tests/AgentTests/FunctionalAPITests.hs
+++ b/tests/AgentTests/FunctionalAPITests.hs
@@ -2828,6 +2828,7 @@ testTwoUsers = withAgentClients2 $ \a b -> do
     liftIO $ threadDelay 250000
     ("", "", DOWN _ _) <- nGet a
     ("", "", UP _ _) <- nGet a
+    ("", "", UP _ _) <- nGet a
     a `hasClients` 2
 
     exchangeGreetingsMsgId 4 a bId1 b aId1
@@ -2836,8 +2837,6 @@ testTwoUsers = withAgentClients2 $ \a b -> do
     liftIO $ setNetworkConfig a nc {sessionMode = TSMUser}
     liftIO $ threadDelay 250000
     ("", "", DOWN _ _) <- nGet a
-    ("", "", DOWN _ _) <- nGet a
-    ("", "", UP _ _) <- nGet a
     ("", "", UP _ _) <- nGet a
     a `hasClients` 1
 
@@ -2851,7 +2850,8 @@ testTwoUsers = withAgentClients2 $ \a b -> do
     liftIO $ setNetworkConfig a nc {sessionMode = TSMEntity}
     liftIO $ threadDelay 250000
     ("", "", DOWN _ _) <- nGet a
-    ("", "", DOWN _ _) <- nGet a
+    ("", "", UP _ _) <- nGet a
+    ("", "", UP _ _) <- nGet a
     ("", "", UP _ _) <- nGet a
     ("", "", UP _ _) <- nGet a
     a `hasClients` 4
@@ -2863,11 +2863,6 @@ testTwoUsers = withAgentClients2 $ \a b -> do
     liftIO $ setNetworkConfig a nc {sessionMode = TSMUser}
     liftIO $ threadDelay 250000
     ("", "", DOWN _ _) <- nGet a
-    ("", "", DOWN _ _) <- nGet a
-    ("", "", DOWN _ _) <- nGet a
-    ("", "", DOWN _ _) <- nGet a
-    ("", "", UP _ _) <- nGet a
-    ("", "", UP _ _) <- nGet a
     ("", "", UP _ _) <- nGet a
     ("", "", UP _ _) <- nGet a
     a `hasClients` 2