Merge pull request #1006 from simplex-chat/v5521

5.5.2.1
This commit is contained in:
Evgeny Poberezkin
2024-02-21 20:07:34 +00:00
committed by GitHub
6 changed files with 76 additions and 47 deletions
+2 -2
View File
@@ -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
+33 -14
View File
@@ -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)
+2 -5
View File
@@ -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'
+24 -17
View File
@@ -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 =
+9 -3
View File
@@ -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)
+6 -6
View File
@@ -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