mirror of
https://github.com/simplex-chat/simplexmq.git
synced 2026-07-20 17:52:02 +00:00
@@ -360,8 +360,8 @@ getNetworkConfig = readTVarIO . useNetworkConfig
|
||||
|
||||
reconnectAllServers :: MonadUnliftIO m => AgentClient -> m ()
|
||||
reconnectAllServers c = liftIO $ do
|
||||
closeProtocolServerClients c smpClients
|
||||
closeProtocolServerClients c ntfClients
|
||||
reconnectServerClients c smpClients
|
||||
reconnectServerClients c ntfClients
|
||||
|
||||
-- | Register device notifications token
|
||||
registerNtfToken :: AgentErrorMonad m => AgentClient -> DeviceToken -> NotificationsMode -> m NtfTknStatus
|
||||
|
||||
@@ -28,6 +28,7 @@ module Simplex.Messaging.Agent.Client
|
||||
withInvLock,
|
||||
closeAgentClient,
|
||||
closeProtocolServerClients,
|
||||
reconnectServerClients,
|
||||
closeXFTPServerClient,
|
||||
runSMPServerTest,
|
||||
runXFTPServerTest,
|
||||
@@ -140,6 +141,7 @@ import Data.Bifunctor (bimap, first, second)
|
||||
import Data.ByteString.Base64
|
||||
import Data.ByteString.Char8 (ByteString)
|
||||
import qualified Data.ByteString.Char8 as B
|
||||
import Data.Composition ((.:.))
|
||||
import Data.Either (lefts, partitionEithers)
|
||||
import Data.Functor (($>))
|
||||
import Data.List (deleteFirstsBy, foldl', partition, (\\))
|
||||
@@ -499,11 +501,15 @@ instance ProtocolServerClient XFTPErrorType FileResponse where
|
||||
getSMPServerClient :: forall m. AgentMonad m => AgentClient -> SMPTransportSession -> m SMPClient
|
||||
getSMPServerClient c@AgentClient {active, smpClients, msgQ} tSess@(userId, srv, _) = do
|
||||
unlessM (readTVarIO active) . throwError $ INACTIVE
|
||||
v <- atomically (getTSessVar c tSess smpClients)
|
||||
either newClient (waitForProtocolClient c tSess) v
|
||||
`catchAgentError` \e -> resubscribeSMPSession c tSess >> throwError e
|
||||
atomically (getTSessVar c tSess smpClients)
|
||||
>>= either newClient (waitForProtocolClient c tSess)
|
||||
where
|
||||
newClient = newProtocolClient c tSess smpClients connectClient
|
||||
-- we resubscribe only on newClient error, but not on waitForProtocolClient error,
|
||||
-- as the large number of delivery workers waiting for the client TMVar
|
||||
-- make it expensive to check for pending subscriptions.
|
||||
newClient v =
|
||||
newProtocolClient c tSess smpClients connectClient v
|
||||
`catchAgentError` \e -> resubscribeSMPSession c tSess >> throwError e
|
||||
connectClient :: SMPClientVar -> m SMPClient
|
||||
connectClient v = do
|
||||
cfg <- getClientConfig c smpCfg
|
||||
@@ -515,14 +521,19 @@ getSMPServerClient c@AgentClient {active, smpClients, msgQ} tSess@(userId, srv,
|
||||
removeClientAndSubs >>= serverDown
|
||||
logInfo . decodeUtf8 $ "Agent disconnected from " <> showServer srv
|
||||
where
|
||||
-- we make active subscriptions pending only if the client for tSess was current (in the map) and active,
|
||||
-- because we can have a race condition when a new current client could have already
|
||||
-- made subscriptions active, and the old client would be processing diconnection later.
|
||||
removeClientAndSubs :: IO ([RcvQueue], [ConnId])
|
||||
removeClientAndSubs = atomically $ do
|
||||
removeTSessVar v tSess smpClients
|
||||
qs <- RQ.getDelSessQueues tSess $ activeSubs c
|
||||
mapM_ (`RQ.addQueue` pendingSubs c) qs
|
||||
let cs = S.fromList $ map qConnId qs
|
||||
cs' <- RQ.getConns $ activeSubs c
|
||||
pure (qs, S.toList $ cs `S.difference` cs')
|
||||
removeClientAndSubs = atomically $ ifM currentActiveClient removeSubs $ pure ([], [])
|
||||
where
|
||||
currentActiveClient = (&&) <$> removeTSessVar' v tSess smpClients <*> readTVar active
|
||||
removeSubs = do
|
||||
qs <- RQ.getDelSessQueues tSess $ activeSubs c
|
||||
mapM_ (`RQ.addQueue` pendingSubs c) qs
|
||||
let cs = S.fromList $ map qConnId qs
|
||||
cs' <- RQ.getConns $ activeSubs c
|
||||
pure (qs, S.toList $ cs `S.difference` cs')
|
||||
|
||||
serverDown :: ([RcvQueue], [ConnId]) -> IO ()
|
||||
serverDown (qs, conns) = whenM (readTVarIO active) $ do
|
||||
@@ -648,9 +659,13 @@ getTSessVar c tSess vs = maybe (Left <$> newSessionVar) (pure . Right) =<< TM.lo
|
||||
pure v
|
||||
|
||||
removeTSessVar :: SessionVar a -> TransportSession msg -> TMap (TransportSession msg) (SessionVar a) -> STM ()
|
||||
removeTSessVar v tSess vs =
|
||||
TM.lookup tSess vs
|
||||
>>= mapM_ (\v' -> when (sessionVarId v == sessionVarId v') $ TM.delete tSess vs)
|
||||
removeTSessVar = void .:. removeTSessVar'
|
||||
|
||||
removeTSessVar' :: SessionVar a -> TransportSession msg -> TMap (TransportSession msg) (SessionVar a) -> STM Bool
|
||||
removeTSessVar' v tSess vs =
|
||||
TM.lookup tSess vs >>= \case
|
||||
Just v' | sessionVarId v == sessionVarId v' -> TM.delete tSess vs $> True
|
||||
_ -> pure False
|
||||
|
||||
waitForProtocolClient :: (AgentMonad m, ProtocolTypeI (ProtoType msg)) => AgentClient -> TransportSession msg -> ClientVar msg -> m (Client msg)
|
||||
waitForProtocolClient c (_, srv, _) v = do
|
||||
@@ -738,6 +753,10 @@ closeProtocolServerClients :: ProtocolServerClient err msg => AgentClient -> (Ag
|
||||
closeProtocolServerClients c clientsSel =
|
||||
atomically (clientsSel c `swapTVar` M.empty) >>= mapM_ (forkIO . closeClient_ c)
|
||||
|
||||
reconnectServerClients :: ProtocolServerClient err msg => AgentClient -> (AgentClient -> TMap (TransportSession msg) (ClientVar msg)) -> IO ()
|
||||
reconnectServerClients c clientsSel =
|
||||
readTVarIO (clientsSel c) >>= mapM_ (forkIO . closeClient_ c)
|
||||
|
||||
closeClient :: ProtocolServerClient err msg => AgentClient -> (AgentClient -> TMap (TransportSession msg) (ClientVar msg)) -> TransportSession msg -> IO ()
|
||||
closeClient c clientSel tSess =
|
||||
atomically (TM.lookupDelete tSess $ clientSel c) >>= mapM_ (closeClient_ c)
|
||||
|
||||
@@ -118,7 +118,6 @@ data ProtocolClient err msg = ProtocolClient
|
||||
sessionId :: SessionId,
|
||||
sessionTs :: UTCTime,
|
||||
thVersion :: Version,
|
||||
timeoutPerBlock :: Int,
|
||||
blockSize :: Int,
|
||||
batch :: Bool,
|
||||
client_ :: PClient err msg
|
||||
@@ -151,7 +150,6 @@ clientStub sessionId = do
|
||||
sessionId,
|
||||
sessionTs = undefined,
|
||||
thVersion = 5,
|
||||
timeoutPerBlock = undefined,
|
||||
blockSize = smpBlockSize,
|
||||
batch = undefined,
|
||||
client_ =
|
||||
@@ -314,7 +312,7 @@ getProtocolClient transportSession@(_, srv, _) cfg@ProtocolClientConfig {qSize,
|
||||
`catch` \(e :: IOException) -> pure . Left $ PCEIOError e
|
||||
Left e -> pure $ Left e
|
||||
where
|
||||
NetworkConfig {tcpConnectTimeout, tcpTimeout, tcpTimeoutPerKb, smpPingInterval} = networkConfig
|
||||
NetworkConfig {tcpConnectTimeout, tcpTimeout, smpPingInterval} = networkConfig
|
||||
mkProtocolClient :: TransportHost -> STM (PClient err msg)
|
||||
mkProtocolClient transportHost = do
|
||||
connected <- newTVar False
|
||||
@@ -365,8 +363,7 @@ getProtocolClient transportSession@(_, srv, _) cfg@ProtocolClientConfig {qSize,
|
||||
Left e -> atomically . putTMVar cVar . Left $ PCETransportError e
|
||||
Right th@THandle {sessionId, thVersion, blockSize, batch} -> do
|
||||
sessionTs <- getCurrentTime
|
||||
let timeoutPerBlock = (blockSize * tcpTimeoutPerKb) `div` 1024
|
||||
c' = ProtocolClient {action = Nothing, client_ = c, sessionId, thVersion, sessionTs, timeoutPerBlock, blockSize, batch}
|
||||
let c' = ProtocolClient {action = Nothing, client_ = c, sessionId, thVersion, sessionTs, blockSize, batch}
|
||||
atomically $ do
|
||||
writeTVar (connected c) True
|
||||
putTMVar cVar $ Right c'
|
||||
|
||||
@@ -50,6 +50,7 @@ import qualified Data.ByteString.Char8 as B
|
||||
import Data.Either (fromRight, partitionEithers)
|
||||
import Data.Functor (($>))
|
||||
import Data.Int (Int64)
|
||||
import qualified Data.IntMap.Strict as IM
|
||||
import Data.List (intercalate)
|
||||
import qualified Data.List.NonEmpty as L
|
||||
import qualified Data.Map.Strict as M
|
||||
@@ -157,7 +158,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg} = do
|
||||
updateSubscribers = do
|
||||
(qId, clnt) <- readTQueue $ subQ s
|
||||
let clientToBeNotified c' =
|
||||
if sameClientSession clnt c'
|
||||
if sameClientId clnt c'
|
||||
then pure Nothing
|
||||
else do
|
||||
yes <- readTVar $ connected c'
|
||||
@@ -165,9 +166,12 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg} = do
|
||||
TM.lookupInsert qId clnt (subs s) $>>= clientToBeNotified
|
||||
endPreviousSubscriptions :: (QueueId, Client) -> M (Maybe s)
|
||||
endPreviousSubscriptions (qId, c) = do
|
||||
void . forkIO $ do
|
||||
tId <- atomically $ stateTVar (endThreadSeq c) $ \next -> (next, next + 1)
|
||||
t <- forkIO $ do
|
||||
labelMyThread $ label <> ".endPreviousSubscriptions"
|
||||
atomically $ writeTBQueue (sndQ c) [(CorrId "", qId, END)]
|
||||
atomically $ modifyTVar' (endThreads c) $ IM.delete tId
|
||||
mkWeakThreadId t >>= atomically . modifyTVar' (endThreads c) . IM.insert tId
|
||||
atomically $ TM.lookupDelete qId (clientSubs c)
|
||||
|
||||
expireMessagesThread_ :: ServerConfig -> [M ()]
|
||||
@@ -286,7 +290,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg} = do
|
||||
CPClients -> do
|
||||
active <- unliftIO u (asks clients) >>= readTVarIO
|
||||
hPutStrLn h $ "clientId,sessionId,connected,createdAt,rcvActiveAt,sndActiveAt,age,subscriptions"
|
||||
forM_ (M.toList active) $ \(cid, Client {sessionId, connected, createdAt, rcvActiveAt, sndActiveAt, subscriptions}) -> do
|
||||
forM_ (IM.toList active) $ \(cid, Client {sessionId, connected, createdAt, rcvActiveAt, sndActiveAt, subscriptions}) -> do
|
||||
connected' <- bshow <$> readTVarIO connected
|
||||
rcvActiveAt' <- strEncode <$> readTVarIO rcvActiveAt
|
||||
sndActiveAt' <- strEncode <$> readTVarIO sndActiveAt
|
||||
@@ -327,13 +331,13 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg} = do
|
||||
hPutStrLn h "Sockets: "
|
||||
hPutStrLn h $ "accepted: " <> show accepted
|
||||
hPutStrLn h $ "closed: " <> show closed
|
||||
hPutStrLn h $ "active: " <> show (M.size active)
|
||||
hPutStrLn h $ "leaked: " <> show (accepted - closed - M.size active)
|
||||
hPutStrLn h $ "active: " <> show (IM.size active)
|
||||
hPutStrLn h $ "leaked: " <> show (accepted - closed - IM.size active)
|
||||
CPSocketThreads -> do
|
||||
#if MIN_VERSION_base(4,18,0)
|
||||
(_, _, active') <- unliftIO u $ asks sockets
|
||||
active <- readTVarIO active'
|
||||
forM_ (M.toList active) $ \(sid, tid') ->
|
||||
forM_ (IM.toList active) $ \(sid, tid') ->
|
||||
deRefWeak tid' >>= \case
|
||||
Nothing -> hPutStrLn h $ intercalate "," [show sid, "", "gone", ""]
|
||||
Just tid -> do
|
||||
@@ -376,7 +380,7 @@ runClientTransport th@THandle {thVersion, sessionId} = do
|
||||
nextClientId <- asks clientSeq
|
||||
c <- atomically $ do
|
||||
new@Client {clientId} <- newClient nextClientId q thVersion sessionId ts
|
||||
TM.insert clientId new active
|
||||
modifyTVar' active $ IM.insert clientId new
|
||||
pure new
|
||||
s <- asks server
|
||||
expCfg <- asks $ inactiveClientExpiration . config
|
||||
@@ -389,23 +393,26 @@ runClientTransport th@THandle {thVersion, sessionId} = do
|
||||
noSubscriptions c = atomically $ (&&) <$> TM.null (subscriptions c) <*> TM.null (ntfSubscriptions c)
|
||||
|
||||
clientDisconnected :: Client -> M ()
|
||||
clientDisconnected c@Client {clientId, subscriptions, connected, sessionId} = do
|
||||
clientDisconnected c@Client {clientId, subscriptions, connected, sessionId, endThreads} = do
|
||||
labelMyThread . B.unpack $ "client $" <> encode sessionId <> " disc"
|
||||
atomically $ writeTVar connected False
|
||||
subs <- readTVarIO subscriptions
|
||||
subs <- atomically $ do
|
||||
writeTVar connected False
|
||||
swapTVar subscriptions M.empty
|
||||
liftIO $ mapM_ cancelSub subs
|
||||
atomically $ writeTVar subscriptions M.empty
|
||||
cs <- asks $ subscribers . server
|
||||
atomically . mapM_ (\rId -> TM.update deleteCurrentClient rId cs) $ M.keys subs
|
||||
asks clients >>= atomically . TM.delete clientId
|
||||
srvSubs <- asks $ subscribers . server
|
||||
atomically $ modifyTVar' srvSubs $ \cs ->
|
||||
M.foldrWithKey (\sub _ -> M.update deleteCurrentClient sub) cs subs
|
||||
asks clients >>= atomically . (`modifyTVar'` IM.delete clientId)
|
||||
tIds <- atomically $ swapTVar endThreads IM.empty
|
||||
liftIO $ mapM_ (mapM_ killThread <=< deRefWeak) tIds
|
||||
where
|
||||
deleteCurrentClient :: Client -> Maybe Client
|
||||
deleteCurrentClient c'
|
||||
| sameClientSession c c' = Nothing
|
||||
| sameClientId c c' = Nothing
|
||||
| otherwise = Just c'
|
||||
|
||||
sameClientSession :: Client -> Client -> Bool
|
||||
sameClientSession Client {sessionId} Client {sessionId = s'} = sessionId == s'
|
||||
sameClientId :: Client -> Client -> Bool
|
||||
sameClientId Client {clientId} Client {clientId = cId'} = clientId == cId'
|
||||
|
||||
cancelSub :: TVar Sub -> IO ()
|
||||
cancelSub sub =
|
||||
|
||||
@@ -10,6 +10,8 @@ import Control.Monad.IO.Unlift
|
||||
import Crypto.Random
|
||||
import Data.ByteString.Char8 (ByteString)
|
||||
import Data.Int (Int64)
|
||||
import Data.IntMap.Strict (IntMap)
|
||||
import qualified Data.IntMap.Strict as IM
|
||||
import Data.List.NonEmpty (NonEmpty)
|
||||
import Data.Map.Strict (Map)
|
||||
import qualified Data.Map.Strict as M
|
||||
@@ -107,7 +109,7 @@ data Env = Env
|
||||
serverStats :: ServerStats,
|
||||
sockets :: SocketState,
|
||||
clientSeq :: TVar Int,
|
||||
clients :: TMap Int Client
|
||||
clients :: TVar (IntMap Client)
|
||||
}
|
||||
|
||||
data Server = Server
|
||||
@@ -124,6 +126,8 @@ data Client = Client
|
||||
ntfSubscriptions :: TMap NotifierId (),
|
||||
rcvQ :: TBQueue (NonEmpty (Maybe QueueRec, Transmission Cmd)),
|
||||
sndQ :: TBQueue (NonEmpty (Transmission BrokerMsg)),
|
||||
endThreads :: TVar (IntMap (Weak ThreadId)),
|
||||
endThreadSeq :: TVar Int,
|
||||
thVersion :: Version,
|
||||
sessionId :: ByteString,
|
||||
connected :: TVar Bool,
|
||||
@@ -155,10 +159,12 @@ newClient nextClientId qSize thVersion sessionId createdAt = do
|
||||
ntfSubscriptions <- TM.empty
|
||||
rcvQ <- newTBQueue qSize
|
||||
sndQ <- newTBQueue qSize
|
||||
endThreads <- newTVar IM.empty
|
||||
endThreadSeq <- newTVar 0
|
||||
connected <- newTVar True
|
||||
rcvActiveAt <- newTVar createdAt
|
||||
sndActiveAt <- newTVar createdAt
|
||||
return Client {clientId, subscriptions, ntfSubscriptions, rcvQ, sndQ, thVersion, sessionId, connected, createdAt, rcvActiveAt, sndActiveAt}
|
||||
return Client {clientId, subscriptions, ntfSubscriptions, rcvQ, sndQ, endThreads, endThreadSeq, thVersion, sessionId, connected, createdAt, rcvActiveAt, sndActiveAt}
|
||||
|
||||
newSubscription :: SubscriptionThread -> STM Sub
|
||||
newSubscription subThread = do
|
||||
@@ -178,7 +184,7 @@ newEnv config@ServerConfig {caCertificateFile, certificateFile, privateKeyFile,
|
||||
serverStats <- atomically . newServerStats =<< liftIO getCurrentTime
|
||||
sockets <- atomically newSocketState
|
||||
clientSeq <- newTVarIO 0
|
||||
clients <- atomically TM.empty
|
||||
clients <- newTVarIO mempty
|
||||
return Env {config, server, serverIdentity, queueStore, msgStore, random, storeLog, tlsServerParams, serverStats, sockets, clientSeq, clients}
|
||||
where
|
||||
restoreQueues :: QueueStore -> FilePath -> m (StoreLog 'WriteMode)
|
||||
|
||||
@@ -29,14 +29,14 @@ import Control.Monad.IO.Unlift
|
||||
import qualified Crypto.Store.X509 as SX
|
||||
import Data.Default (def)
|
||||
import Data.List (find)
|
||||
import Data.IntMap.Strict (IntMap)
|
||||
import qualified Data.IntMap.Strict as IM
|
||||
import Data.Maybe (fromJust)
|
||||
import qualified Data.X509 as X
|
||||
import Data.X509.Validation (Fingerprint (..))
|
||||
import qualified Data.X509.Validation as XV
|
||||
import Network.Socket
|
||||
import qualified Network.TLS as T
|
||||
import Simplex.Messaging.TMap (TMap)
|
||||
import qualified Simplex.Messaging.TMap as TM
|
||||
import Simplex.Messaging.Transport
|
||||
import Simplex.Messaging.Util (catchAll_, labelMyThread, tshow)
|
||||
import System.Exit (exitFailure)
|
||||
@@ -110,18 +110,18 @@ runTCPServerSocket (accepted, gracefullyClosed, clients) started getSocket serve
|
||||
forever . E.bracketOnError (accept sock) (close . fst) $ \(conn, _peer) -> do
|
||||
cId <- atomically $ stateTVar accepted $ \cId -> let cId' = cId + 1 in cId `seq` (cId', cId')
|
||||
let closeConn _ = do
|
||||
atomically $ TM.delete cId clients
|
||||
atomically $ modifyTVar' clients $ IM.delete cId
|
||||
gracefulClose conn 5000 `catchAll_` pure () -- catchAll_ is needed here in case the connection was closed earlier
|
||||
atomically $ modifyTVar' gracefullyClosed (+1)
|
||||
tId <- mkWeakThreadId =<< server conn `forkFinally` closeConn
|
||||
atomically $ TM.insert cId tId clients
|
||||
atomically $ modifyTVar' clients $ IM.insert cId tId
|
||||
|
||||
type SocketState = (TVar Int, TVar Int, TMap Int (Weak ThreadId))
|
||||
type SocketState = (TVar Int, TVar Int, TVar (IntMap (Weak ThreadId)))
|
||||
|
||||
newSocketState :: STM SocketState
|
||||
newSocketState = (,,) <$> newTVar 0 <*> newTVar 0 <*> newTVar mempty
|
||||
|
||||
closeServer :: TMVar Bool -> TMap Int (Weak ThreadId) -> Socket -> IO ()
|
||||
closeServer :: TMVar Bool -> TVar (IntMap (Weak ThreadId)) -> Socket -> IO ()
|
||||
closeServer started clients sock = do
|
||||
readTVarIO clients >>= mapM_ (deRefWeak >=> mapM_ killThread)
|
||||
close sock
|
||||
|
||||
Reference in New Issue
Block a user