From c605156302b78a8411a1a083c52e7d5506f11f9c Mon Sep 17 00:00:00 2001
From: spaced4ndy <8711996+spaced4ndy@users.noreply.github.com>
Date: Mon, 22 Jul 2024 18:43:38 +0400
Subject: [PATCH] agent: getAgentSubsTotal api (#1238)

* agent: getAgentSubsTotal api

* move

* export

* count sessions

* simplify

* simplify 2

---------

Co-authored-by: Evgeny Poberezkin <evgeny@poberezkin.com>
---
 src/Simplex/Messaging/Agent.hs        |  1 +
 src/Simplex/Messaging/Agent/Client.hs | 18 ++++++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/src/Simplex/Messaging/Agent.hs b/src/Simplex/Messaging/Agent.hs
index f26c46925..0bd95d1dc 100644
--- a/src/Simplex/Messaging/Agent.hs
+++ b/src/Simplex/Messaging/Agent.hs
@@ -106,6 +106,7 @@ module Simplex.Messaging.Agent
     rcConnectHost,
     rcConnectCtrl,
     rcDiscoverCtrl,
+    getAgentSubsTotal,
     getAgentServersSummary,
     resetAgentServersStats,
     foregroundAgent,
diff --git a/src/Simplex/Messaging/Agent/Client.hs b/src/Simplex/Messaging/Agent/Client.hs
index 3a7616499..8aacd8836 100644
--- a/src/Simplex/Messaging/Agent/Client.hs
+++ b/src/Simplex/Messaging/Agent/Client.hs
@@ -93,6 +93,7 @@ module Simplex.Messaging.Agent.Client
     AgentServersSummary (..),
     ServerSessions (..),
     SMPServerSubs (..),
+    getAgentSubsTotal,
     getAgentServersSummary,
     getAgentSubscriptions,
     slowNetworkConfig,
@@ -2017,6 +2018,23 @@ data ServerSessions = ServerSessions
   }
   deriving (Show)
 
+getAgentSubsTotal :: AgentClient -> [UserId] -> IO (SMPServerSubs, Bool)
+getAgentSubsTotal c userIds = do
+  ssActive <- getSubsCount activeSubs
+  ssPending <- getSubsCount pendingSubs
+  sess <- hasSession . M.toList =<< readTVarIO (smpClients c)
+  pure (SMPServerSubs {ssActive, ssPending}, sess)
+  where
+    getSubsCount subs = M.foldrWithKey' addSub 0 <$> readTVarIO (getRcvQueues $ subs c)
+    addSub (userId, _, _) _ cnt = if userId `elem` userIds then cnt + 1 else cnt
+    hasSession :: [(SMPTransportSession, SMPClientVar)] -> IO Bool
+    hasSession = \case
+      [] -> pure False
+      (s : ss) -> ifM (isConnected s) (pure True) (hasSession ss)
+    isConnected ((userId, _, _), SessionVar {sessionVar})
+      | userId `elem` userIds = atomically $ maybe False isRight <$> tryReadTMVar sessionVar
+      | otherwise = pure False
+
 getAgentServersSummary :: AgentClient -> IO AgentServersSummary
 getAgentServersSummary c@AgentClient {smpServersStats, xftpServersStats, ntfServersStats, srvStatsStartedAt, agentEnv} = do
   sss <- mapM getAgentSMPServerStats =<< readTVarIO smpServersStats