mirror of
https://github.com/simplex-chat/simplexmq.git
synced 2026-07-28 10:00:02 +00:00
refactor agent subscriptions with TMap2 (#517)
* refactor agent subscriptions with TMap2 * refactor * refactor * comment
This commit is contained in:
Evgeny Poberezkin
@@ -82,6 +82,7 @@ library
|
||||
Simplex.Messaging.Server.Stats
|
||||
Simplex.Messaging.Server.StoreLog
|
||||
Simplex.Messaging.TMap
|
||||
Simplex.Messaging.TMap2
|
||||
Simplex.Messaging.Transport
|
||||
Simplex.Messaging.Transport.Client
|
||||
Simplex.Messaging.Transport.HTTP2
|
||||
|
||||
@@ -90,7 +90,6 @@ import Data.Maybe (listToMaybe)
|
||||
import Data.Set (Set)
|
||||
import qualified Data.Set as S
|
||||
import Data.Text.Encoding
|
||||
import Data.Tuple (swap)
|
||||
import Data.Word (Word16)
|
||||
import qualified Database.SQLite.Simple as DB
|
||||
import Simplex.Messaging.Agent.Env.SQLite
|
||||
@@ -131,6 +130,8 @@ import Simplex.Messaging.Protocol
|
||||
import qualified Simplex.Messaging.Protocol as SMP
|
||||
import Simplex.Messaging.TMap (TMap)
|
||||
import qualified Simplex.Messaging.TMap as TM
|
||||
import Simplex.Messaging.TMap2 (TMap2)
|
||||
import qualified Simplex.Messaging.TMap2 as TM2
|
||||
import Simplex.Messaging.Transport.Client (TransportHost)
|
||||
import Simplex.Messaging.Util
|
||||
import Simplex.Messaging.Version
|
||||
@@ -155,10 +156,9 @@ data AgentClient = AgentClient
|
||||
ntfServers :: TVar [NtfServer],
|
||||
ntfClients :: TMap NtfServer NtfClientVar,
|
||||
useNetworkConfig :: TVar NetworkConfig,
|
||||
subscrSrvrs :: TMap SMPServer (TMap ConnId RcvQueue),
|
||||
pendingSubscrSrvrs :: TMap SMPServer (TMap ConnId RcvQueue),
|
||||
subscrConns :: TVar (Set ConnId),
|
||||
activeSubscrConns :: TMap ConnId SMPServer,
|
||||
activeSubs :: TMap2 SMPServer ConnId RcvQueue,
|
||||
pendingSubs :: TMap2 SMPServer ConnId RcvQueue,
|
||||
connMsgsQueued :: TMap ConnId Bool,
|
||||
smpQueueMsgQueues :: TMap (SMPServer, SMP.SenderId) (TQueue InternalId),
|
||||
smpQueueMsgDeliveries :: TMap (SMPServer, SMP.SenderId) (Async ()),
|
||||
@@ -210,10 +210,9 @@ newAgentClient InitialAgentServers {smp, ntf, netCfg} agentEnv = do
|
||||
ntfServers <- newTVar ntf
|
||||
ntfClients <- TM.empty
|
||||
useNetworkConfig <- newTVar netCfg
|
||||
subscrSrvrs <- TM.empty
|
||||
pendingSubscrSrvrs <- TM.empty
|
||||
subscrConns <- newTVar S.empty
|
||||
activeSubscrConns <- TM.empty
|
||||
activeSubs <- TM2.empty
|
||||
pendingSubs <- TM2.empty
|
||||
connMsgsQueued <- TM.empty
|
||||
smpQueueMsgQueues <- TM.empty
|
||||
smpQueueMsgDeliveries <- TM.empty
|
||||
@@ -231,7 +230,7 @@ newAgentClient InitialAgentServers {smp, ntf, netCfg} agentEnv = do
|
||||
asyncClients <- newTVar []
|
||||
clientId <- stateTVar (clientCounter agentEnv) $ \i -> let i' = i + 1 in (i', i')
|
||||
lock <- newTMVar ()
|
||||
return AgentClient {active, rcvQ, subQ, msgQ, smpServers, smpClients, ntfServers, ntfClients, useNetworkConfig, subscrSrvrs, pendingSubscrSrvrs, subscrConns, activeSubscrConns, connMsgsQueued, smpQueueMsgQueues, smpQueueMsgDeliveries, connCmdsQueued, asyncCmdQueues, asyncCmdProcesses, ntfNetworkOp, rcvNetworkOp, msgDeliveryOp, sndNetworkOp, databaseOp, agentState, getMsgLocks, reconnections, asyncClients, clientId, agentEnv, lock}
|
||||
return AgentClient {active, rcvQ, subQ, msgQ, smpServers, smpClients, ntfServers, ntfClients, useNetworkConfig, subscrConns, activeSubs, pendingSubs, connMsgsQueued, smpQueueMsgQueues, smpQueueMsgDeliveries, connCmdsQueued, asyncCmdQueues, asyncCmdProcesses, ntfNetworkOp, rcvNetworkOp, msgDeliveryOp, sndNetworkOp, databaseOp, agentState, getMsgLocks, reconnections, asyncClients, clientId, agentEnv, lock}
|
||||
|
||||
agentDbPath :: AgentClient -> FilePath
|
||||
agentDbPath AgentClient {agentEnv = Env {store = SQLiteStore {dbFilePath}}} = dbFilePath
|
||||
@@ -270,26 +269,18 @@ getSMPServerClient c@AgentClient {active, smpClients, msgQ} srv = do
|
||||
removeClientAndSubs :: IO (Maybe (Map ConnId RcvQueue))
|
||||
removeClientAndSubs = atomically $ do
|
||||
TM.delete srv smpClients
|
||||
TM.lookupDelete srv (subscrSrvrs c) >>= mapM updateSubs
|
||||
TM2.lookupDelete1 srv (activeSubs c) >>= mapM updateSubs
|
||||
where
|
||||
updateSubs cVar = do
|
||||
cs <- readTVar cVar
|
||||
modifyTVar' (activeSubscrConns c) (`M.withoutKeys` M.keysSet cs)
|
||||
addPendingSubs cVar cs
|
||||
pure cs
|
||||
|
||||
addPendingSubs cVar cs = do
|
||||
let ps = pendingSubscrSrvrs c
|
||||
TM.lookup srv ps >>= \case
|
||||
Just v -> TM.union cs v
|
||||
_ -> TM.insert srv cVar ps
|
||||
TM2.insert1 srv cVar $ pendingSubs c
|
||||
readTVar cVar
|
||||
|
||||
serverDown :: Map ConnId RcvQueue -> IO ()
|
||||
serverDown cs = unless (M.null cs) $
|
||||
whenM (readTVarIO active) $ do
|
||||
let conns = M.keys cs
|
||||
notifySub "" $ hostEvent DISCONNECT client
|
||||
unless (null conns) . notifySub "" $ DOWN srv conns
|
||||
serverDown cs = whenM (readTVarIO active) $ do
|
||||
notifySub "" $ hostEvent DISCONNECT client
|
||||
let conns = M.keys cs
|
||||
unless (null conns) $ do
|
||||
notifySub "" $ DOWN srv conns
|
||||
atomically $ mapM_ (releaseGetLock c) cs
|
||||
unliftIO u reconnectServer
|
||||
|
||||
@@ -307,7 +298,7 @@ getSMPServerClient c@AgentClient {active, smpClients, msgQ} srv = do
|
||||
reconnectClient :: m ()
|
||||
reconnectClient =
|
||||
withAgentLock c $
|
||||
atomically (TM.lookup srv (pendingSubscrSrvrs c) >>= mapM readTVar)
|
||||
atomically (TM2.lookup1 srv (pendingSubs c) >>= mapM readTVar)
|
||||
>>= mapM_ resubscribe
|
||||
where
|
||||
resubscribe :: Map ConnId RcvQueue -> m ()
|
||||
@@ -413,10 +404,9 @@ closeAgentClient c = liftIO $ do
|
||||
cancelActions $ reconnections c
|
||||
cancelActions $ asyncClients c
|
||||
cancelActions $ smpQueueMsgDeliveries c
|
||||
clear subscrSrvrs
|
||||
clear pendingSubscrSrvrs
|
||||
atomically . TM2.clear $ activeSubs c
|
||||
atomically . TM2.clear $ pendingSubs c
|
||||
clear subscrConns
|
||||
clear activeSubscrConns
|
||||
clear connMsgsQueued
|
||||
clear smpQueueMsgQueues
|
||||
clear getMsgLocks
|
||||
@@ -520,17 +510,17 @@ subscribeQueue c rq@RcvQueue {server, rcvPrivateKey, rcvId} connId = do
|
||||
whenM (atomically . TM.member (server, rcvId) $ getMsgLocks c) . throwError $ CMD PROHIBITED
|
||||
atomically $ do
|
||||
modifyTVar (subscrConns c) $ S.insert connId
|
||||
addPendingSubscription c rq connId
|
||||
TM2.insert server connId rq $ pendingSubs c
|
||||
withLogClient c server rcvId "SUB" $ \smp ->
|
||||
liftIO (runExceptT (subscribeSMPQueue smp rcvPrivateKey rcvId) >>= processSubResult c rq connId)
|
||||
>>= either throwError pure
|
||||
|
||||
processSubResult :: AgentClient -> RcvQueue -> ConnId -> Either ProtocolClientError () -> IO (Either ProtocolClientError ())
|
||||
processSubResult c rq@RcvQueue {server} connId r = do
|
||||
processSubResult c rq connId r = do
|
||||
case r of
|
||||
Left e ->
|
||||
atomically . unless (temporaryClientError e) $
|
||||
removePendingSubscription c server connId
|
||||
TM2.delete connId (pendingSubs c)
|
||||
_ -> addSubscription c rq connId
|
||||
pure r
|
||||
|
||||
@@ -550,9 +540,9 @@ temporaryAgentError = \case
|
||||
subscribeQueues :: AgentMonad m => AgentClient -> SMPServer -> Map ConnId RcvQueue -> m (Maybe SMPClient, Map ConnId (Either AgentErrorType ()))
|
||||
subscribeQueues c srv qs = do
|
||||
(errs, qs_) <- partitionEithers <$> mapM checkQueue (M.assocs qs)
|
||||
forM_ qs_ $ \q -> atomically $ do
|
||||
modifyTVar (subscrConns c) . S.insert $ fst q
|
||||
uncurry (addPendingSubscription c) $ swap q
|
||||
forM_ qs_ $ \(connId, rq@RcvQueue {server}) -> atomically $ do
|
||||
modifyTVar (subscrConns c) $ S.insert connId
|
||||
TM2.insert server connId rq $ pendingSubs c
|
||||
case L.nonEmpty qs_ of
|
||||
Just qs' -> do
|
||||
smp_ <- tryError (getSMPServerClient c srv)
|
||||
@@ -574,35 +564,18 @@ subscribeQueues c srv qs = do
|
||||
|
||||
addSubscription :: MonadIO m => AgentClient -> RcvQueue -> ConnId -> m ()
|
||||
addSubscription c rq@RcvQueue {server} connId = atomically $ do
|
||||
TM.insert connId server $ activeSubscrConns c
|
||||
modifyTVar (subscrConns c) $ S.insert connId
|
||||
addSubs_ (subscrSrvrs c) rq connId
|
||||
removePendingSubscription c server connId
|
||||
TM2.insert server connId rq $ activeSubs c
|
||||
TM2.delete connId $ pendingSubs c
|
||||
|
||||
hasActiveSubscription :: AgentClient -> ConnId -> STM Bool
|
||||
hasActiveSubscription c connId = TM.member connId (activeSubscrConns c)
|
||||
|
||||
addPendingSubscription :: AgentClient -> RcvQueue -> ConnId -> STM ()
|
||||
addPendingSubscription = addSubs_ . pendingSubscrSrvrs
|
||||
|
||||
addSubs_ :: TMap SMPServer (TMap ConnId RcvQueue) -> RcvQueue -> ConnId -> STM ()
|
||||
addSubs_ ss rq@RcvQueue {server} connId =
|
||||
TM.lookup server ss >>= \case
|
||||
Just m -> TM.insert connId rq m
|
||||
_ -> TM.singleton connId rq >>= \m -> TM.insert server m ss
|
||||
hasActiveSubscription c connId = TM2.member connId $ activeSubs c
|
||||
|
||||
removeSubscription :: AgentClient -> ConnId -> STM ()
|
||||
removeSubscription c connId = do
|
||||
modifyTVar (subscrConns c) $ S.delete connId
|
||||
server_ <- TM.lookupDelete connId $ activeSubscrConns c
|
||||
mapM_ (\server -> removeSubs_ (subscrSrvrs c) server connId) server_
|
||||
|
||||
removePendingSubscription :: AgentClient -> SMPServer -> ConnId -> STM ()
|
||||
removePendingSubscription = removeSubs_ . pendingSubscrSrvrs
|
||||
|
||||
removeSubs_ :: TMap SMPServer (TMap ConnId RcvQueue) -> SMPServer -> ConnId -> STM ()
|
||||
removeSubs_ ss server connId =
|
||||
TM.lookup server ss >>= mapM_ (TM.delete connId)
|
||||
TM2.delete connId $ activeSubs c
|
||||
TM2.delete connId $ pendingSubs c
|
||||
|
||||
getSubscriptions :: AgentClient -> STM (Set ConnId)
|
||||
getSubscriptions = readTVar . subscrConns
|
||||
|
||||
@@ -2,6 +2,7 @@ module Simplex.Messaging.TMap
|
||||
( TMap,
|
||||
empty,
|
||||
singleton,
|
||||
clear,
|
||||
Simplex.Messaging.TMap.null,
|
||||
Simplex.Messaging.TMap.lookup,
|
||||
member,
|
||||
@@ -31,6 +32,10 @@ singleton :: k -> a -> STM (TMap k a)
|
||||
singleton k v = newTVar $ M.singleton k v
|
||||
{-# INLINE singleton #-}
|
||||
|
||||
clear :: TMap k a -> STM ()
|
||||
clear m = writeTVar m M.empty
|
||||
{-# INLINE clear #-}
|
||||
|
||||
null :: TMap k a -> STM Bool
|
||||
null m = M.null <$> readTVar m
|
||||
{-# INLINE null #-}
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
{-# LANGUAGE LambdaCase #-}
|
||||
{-# LANGUAGE NamedFieldPuns #-}
|
||||
|
||||
module Simplex.Messaging.TMap2
|
||||
( TMap2,
|
||||
empty,
|
||||
clear,
|
||||
Simplex.Messaging.TMap2.lookup,
|
||||
lookup1,
|
||||
member,
|
||||
insert,
|
||||
insert1,
|
||||
delete,
|
||||
lookupDelete1,
|
||||
)
|
||||
where
|
||||
|
||||
import Control.Concurrent.STM
|
||||
import Control.Monad (forM_, (>=>))
|
||||
import qualified Data.Map.Strict as M
|
||||
import Simplex.Messaging.TMap (TMap)
|
||||
import qualified Simplex.Messaging.TMap as TM
|
||||
import Simplex.Messaging.Util (whenM, ($>>=))
|
||||
|
||||
-- | this type is designed for k2 being unique in the whole data, and k1 grouping multiple values with k2 keys.
|
||||
-- It allows direct access via k1 to a group of k2 values and via k2 to one value
|
||||
data TMap2 k1 k2 a = TMap2
|
||||
{ _m1 :: TMap k1 (TMap k2 a),
|
||||
_m2 :: TMap k2 k1
|
||||
}
|
||||
|
||||
empty :: STM (TMap2 k1 k2 a)
|
||||
empty = TMap2 <$> TM.empty <*> TM.empty
|
||||
|
||||
clear :: TMap2 k1 k2 a -> STM ()
|
||||
clear TMap2 {_m1, _m2} = TM.clear _m1 >> TM.clear _m2
|
||||
|
||||
lookup :: (Ord k1, Ord k2) => k2 -> TMap2 k1 k2 a -> STM (Maybe a)
|
||||
lookup k2 TMap2 {_m1, _m2} = do
|
||||
TM.lookup k2 _m2 $>>= (`TM.lookup` _m1) $>>= TM.lookup k2
|
||||
|
||||
lookup1 :: Ord k1 => k1 -> TMap2 k1 k2 a -> STM (Maybe (TMap k2 a))
|
||||
lookup1 k1 TMap2 {_m1} = TM.lookup k1 _m1
|
||||
{-# INLINE lookup1 #-}
|
||||
|
||||
member :: Ord k2 => k2 -> TMap2 k1 k2 a -> STM Bool
|
||||
member k2 TMap2 {_m2} = TM.member k2 _m2
|
||||
{-# INLINE member #-}
|
||||
|
||||
insert :: (Ord k1, Ord k2) => k1 -> k2 -> a -> TMap2 k1 k2 a -> STM ()
|
||||
insert k1 k2 v TMap2 {_m1, _m2} =
|
||||
TM.lookup k2 _m2 >>= \case
|
||||
Just k1'
|
||||
| k1 == k1' -> _insert1
|
||||
| otherwise -> _delete1 k1' k2 _m1 >> _insert2
|
||||
_ -> _insert2
|
||||
where
|
||||
_insert1 =
|
||||
TM.lookup k1 _m1 >>= \case
|
||||
Just m -> TM.insert k2 v m
|
||||
_ -> TM.singleton k2 v >>= \m -> TM.insert k1 m _m1
|
||||
_insert2 = TM.insert k2 k1 _m2 >> _insert1
|
||||
|
||||
insert1 :: (Ord k1, Ord k2) => k1 -> TMap k2 a -> TMap2 k1 k2 a -> STM ()
|
||||
insert1 k1 m' TMap2 {_m1, _m2} =
|
||||
TM.lookup k1 _m1 >>= \case
|
||||
Just m -> readTVar m' >>= (`TM.union` m)
|
||||
_ -> TM.insert k1 m' _m1
|
||||
|
||||
delete :: (Ord k1, Ord k2) => k2 -> TMap2 k1 k2 a -> STM ()
|
||||
delete k2 TMap2 {_m1, _m2} = TM.lookupDelete k2 _m2 >>= mapM_ (\k1 -> _delete1 k1 k2 _m1)
|
||||
|
||||
_delete1 :: (Ord k1, Ord k2) => k1 -> k2 -> TMap k1 (TMap k2 a) -> STM ()
|
||||
_delete1 k1 k2 m1 =
|
||||
TM.lookup k1 m1
|
||||
>>= mapM_ (\m -> TM.delete k2 m >> whenM (TM.null m) (TM.delete k1 m1))
|
||||
|
||||
lookupDelete1 :: (Ord k1, Ord k2) => k1 -> TMap2 k1 k2 a -> STM (Maybe (TMap k2 a))
|
||||
lookupDelete1 k1 TMap2 {_m1, _m2} = do
|
||||
m_ <- TM.lookupDelete k1 _m1
|
||||
forM_ m_ $ readTVar >=> modifyTVar' _m2 . flip M.withoutKeys . M.keysSet
|
||||
pure m_
|
||||
Reference in New Issue
Block a user