diff --git a/src/Simplex/Messaging/Server.hs b/src/Simplex/Messaging/Server.hs index a9dd33250..2e09a1dff 100644 --- a/src/Simplex/Messaging/Server.hs +++ b/src/Simplex/Messaging/Server.hs @@ -256,7 +256,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt forall sub. String -> Server s -> (Server s -> ServerSubscribers s) -> - (Client s -> TMap SubKey sub) -> + (Client s -> TMap QueueId sub) -> (Client s -> TVar (Int64, IdsHash)) -> Maybe (sub -> IO ()) -> M s () @@ -278,17 +278,17 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt Nothing -> updateSubDisconnected where updateSubConnected c = case clntSub of - CSClient qId k prevServiceId serviceId_ -> do + CSClient qId prevServiceId serviceId_ -> do modifyTVar' subClients $ IS.insert clntId -- add ID to server's subscribed cients as'' <- if prevServiceId == serviceId_ then pure [] else endServiceSub prevServiceId qId END case serviceId_ of Just serviceId -> do modifyTVar' totalServiceSubs $ addServiceSubs (1, queueIdHash qId) -- server count and IDs hash for all services as <- endQueueSub qId END - as' <- cancelServiceSubs serviceId =<< upsertSubscribedClient (subKey serviceId) c serviceSubscribers + as' <- cancelServiceSubs serviceId =<< upsertSubscribedClient serviceId c serviceSubscribers pure $ as ++ as' ++ as'' Nothing -> do - as <- prevSub qId END (CSAEndSub qId) =<< upsertSubscribedClient k c queueSubscribers + as <- prevSub qId END (CSAEndSub qId) =<< upsertSubscribedClient qId c queueSubscribers pure $ as ++ as'' CSDeleted qId serviceId -> do removeWhenNoSubs c @@ -298,10 +298,10 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt CSService serviceId changedSubs -> do modifyTVar' subClients $ IS.insert clntId -- add ID to server's subscribed cients modifyTVar' totalServiceSubs $ addServiceSubs changedSubs -- server count and IDs hash for all services - cancelServiceSubs serviceId =<< upsertSubscribedClient (subKey serviceId) c serviceSubscribers + cancelServiceSubs serviceId =<< upsertSubscribedClient serviceId c serviceSubscribers updateSubDisconnected = case clntSub of -- do not insert client if it is already disconnected, but send END/DELD to any other client subscribed to this queue or service - CSClient qId _ prevServiceId serviceId -> do + CSClient qId prevServiceId serviceId -> do as <- endQueueSub qId END as' <- endServiceSub serviceId qId END as'' <- if prevServiceId == serviceId then pure [] else endServiceSub prevServiceId qId END @@ -345,7 +345,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt unsubPrev :: Maybe sub -> IO () unsubPrev s_ = sequence_ (unsub_ <*> s_) endSub :: Client s -> QueueId -> STM (Maybe sub) - endSub c qId = TM.lookupDelete (subKey qId) (clientSubs c) >>= (removeWhenNoSubs c $>) + endSub c qId = TM.lookupDelete qId (clientSubs c) >>= (removeWhenNoSubs c $>) endServiceQueueSub :: Client s -> QueueId -> STM (Maybe sub) endServiceQueueSub c qId = do modifyTVar' (clientServiceSubs c) decrease @@ -387,7 +387,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt whenM (currentClient readTVarIO c) $ do subs <- readTVarIO ntfSubscriptions unless (M.null subs) $ do - let ntfs'' = filter (\(nId, _) -> M.member (subKey nId) subs) ntfs' + let ntfs'' = filter (\(nId, _) -> M.member nId subs) ntfs' tryAny (atomically $ flushSubscribedNtfs ntfs'' c) >>= updateNtfStats c deliverServiceNtfs ntfs' cv = readTVarIO cv >>= mapM_ deliver where @@ -946,9 +946,9 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt clnts' <- readTVarIO subClients hPutStrLn h $ protoName <> " subscribed clients count 2: " <> show (IS.size clnts') <> (if showIds then " " <> show clnts' else "") where - countSubClients :: Map SubKey (TVar (Maybe (Client s))) -> IO IS.IntSet + countSubClients :: Map QueueId (TVar (Maybe (Client s))) -> IO IS.IntSet countSubClients = foldM (\ !s c -> maybe s ((`IS.insert` s) . clientId) <$> readTVarIO c) IS.empty - countClientSubs :: (Client s -> TMap SubKey a) -> Maybe (Map SubKey a -> IO (Int, Int, Int, Int)) -> IM.IntMap (Client s) -> IO (Int, (Int, Int, Int, Int), Int, (Natural, Natural, Natural)) + countClientSubs :: (Client s -> TMap QueueId a) -> Maybe (Map QueueId a -> IO (Int, Int, Int, Int)) -> IM.IntMap (Client s) -> IO (Int, (Int, Int, Int, Int), Int, (Natural, Natural, Natural)) countClientSubs subSel countSubs_ = foldM addSubs (0, (0, 0, 0, 0), 0, (0, 0, 0)) where addSubs :: (Int, (Int, Int, Int, Int), Int, (Natural, Natural, Natural)) -> Client s -> IO (Int, (Int, Int, Int, Int), Int, (Natural, Natural, Natural)) @@ -973,7 +973,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt sl <- atomically $ lengthTBQueue sndQ ml <- atomically $ lengthTBQueue msgQ pure (rl, sl, ml) - countSMPSubs :: Map SubKey Sub -> IO (Int, Int, Int, Int) + countSMPSubs :: Map QueueId Sub -> IO (Int, Int, Int, Int) countSMPSubs = foldM countSubs (0, 0, 0, 0) where countSubs (c1, c2, c3, c4) Sub {subThread} = case subThread of @@ -1126,13 +1126,13 @@ clientDisconnected c@Client {clientId, subscriptions, ntfSubscriptions, serviceS tIds <- atomically $ swapTVar endThreads IM.empty liftIO $ mapM_ (mapM_ killThread <=< deRefWeak) tIds where - updateSubscribers :: Map SubKey a -> ServerSubscribers s -> IO () + updateSubscribers :: Map QueueId a -> ServerSubscribers s -> IO () updateSubscribers subs ServerSubscribers {queueSubscribers, subClients} = do mapM_ (\qId -> deleteSubcribedClient qId c queueSubscribers) (M.keys subs) atomically $ modifyTVar' subClients $ IS.delete clientId updateServiceSubs :: ServiceId -> TVar (Int64, IdsHash) -> ServerSubscribers s -> IO () updateServiceSubs serviceId subsCount ServerSubscribers {totalServiceSubs, serviceSubscribers} = do - deleteSubcribedClient (subKey serviceId) c serviceSubscribers + deleteSubcribedClient serviceId c serviceSubscribers atomically . modifyTVar' totalServiceSubs . subtractServiceSubs =<< readTVarIO subsCount cancelSub :: Sub -> IO () @@ -1684,7 +1684,7 @@ client subscribeQueueAndDeliver :: Maybe Message -> StoreQueue s -> QueueRec -> M s ResponseAndMessage subscribeQueueAndDeliver msg_ q qr@QueueRec {rcvServiceId} = - liftIO (TM.lookupIO (subKey entId) $ subscriptions clnt) >>= \case + liftIO (TM.lookupIO entId $ subscriptions clnt) >>= \case Nothing -> deliver =<< sharedSubscribeQueue q rcvServiceId subscribers subscriptions serviceSubsCount (newSubscription NoSub) rcvServices Just s@Sub {subThread} -> do @@ -1712,27 +1712,26 @@ client getSub :: STM Sub getSub = - TM.lookup (subKey entId) (subscriptions clnt) >>= \case + TM.lookup entId (subscriptions clnt) >>= \case Just sub -> pure sub Nothing -> do sub <- newSubscription NoSub - TM.insert (subKey entId) sub $ subscriptions clnt + TM.insert entId sub $ subscriptions clnt pure sub subscribeNewQueue :: RecipientId -> QueueRec -> M s () subscribeNewQueue rId QueueRec {rcvServiceId} = do - let k = subKey rId case rcvServiceId of Just _ -> atomically $ modifyTVar' (serviceSubsCount clnt) $ addServiceSubs (1, queueIdHash rId) Nothing -> do sub <- atomically $ newSubscription NoSub - atomically $ TM.insert k sub $ subscriptions clnt - atomically $ writeTQueue (subQ subscribers) (CSClient rId k rcvServiceId rcvServiceId, clientId) + atomically $ TM.insert rId sub $ subscriptions clnt + atomically $ writeTQueue (subQ subscribers) (CSClient rId rcvServiceId rcvServiceId, clientId) -- clients that use GET are not added to server subscribers getMessage :: StoreQueue s -> QueueRec -> M s (Transmission BrokerMsg) getMessage q qr = do - atomically (TM.lookup (subKey entId) $ subscriptions clnt) >>= \case + atomically (TM.lookup entId $ subscriptions clnt) >>= \case Nothing -> atomically newSub >>= (`getMessage_` Nothing) Just s@Sub {subThread} -> @@ -1749,7 +1748,7 @@ client newSub :: STM Sub newSub = do s <- newProhibitedSub - TM.insert (subKey entId) s $ subscriptions clnt + TM.insert entId s $ subscriptions clnt -- Here we don't account for this client as subscribed in the server -- and don't notify other subscribed clients. -- This is tracked as "subscription" in the client to prevent these @@ -1798,7 +1797,7 @@ client StoreQueue s -> Maybe ServiceId -> ServerSubscribers s -> - (Client s -> TMap SubKey sub) -> + (Client s -> TMap QueueId sub) -> (Client s -> TVar (Int64, IdsHash)) -> STM sub -> (ServerStats -> ServiceStats) -> @@ -1806,8 +1805,7 @@ client sharedSubscribeQueue q queueServiceId srvSubscribers clientSubs clientServiceSubs mkSub servicesSel = do stats <- asks serverStats let incSrvStat sel = incStat $ sel $ servicesSel stats - k = subKey entId - writeSub = writeTQueue (subQ srvSubscribers) (CSClient entId k queueServiceId clntServiceId, clientId) + writeSub = writeTQueue (subQ srvSubscribers) (CSClient entId queueServiceId clntServiceId, clientId) liftIO $ case clntServiceId of Just serviceId | queueServiceId == Just serviceId -> do @@ -1845,12 +1843,12 @@ client unless hasSub $ atomically writeSub pure r where - getSubscription = TM.lookup k $ clientSubs clnt + getSubscription = TM.lookup entId $ clientSubs clnt newSub = \case Just sub -> pure (True, Just sub) Nothing -> do sub <- mkSub - TM.insert k sub $ clientSubs clnt + TM.insert entId sub $ clientSubs clnt pure (False, Just sub) subscribeServiceMessages :: ServiceId -> (Int64, IdsHash) -> M s BrokerMsg @@ -1889,13 +1887,13 @@ client atomically $ writeTBQueue msgQ [(NoCorrId, rId, MSG (encryptMsg qr msg))] pure (qCnt + 1, msgCnt + 1, dupCnt, evts) getSubscription rId = - TM.lookup (subKey rId) (subscriptions clnt) >>= \case + TM.lookup rId (subscriptions clnt) >>= \case -- If delivery subscription already exists, then there is no need to deliver message. -- It may have been created when the message is sent after service subscription is created. Just _sub -> pure Nothing Nothing -> do sub <- newSubscription NoSub - TM.insert (subKey rId) sub $ subscriptions clnt + TM.insert rId sub $ subscriptions clnt pure $ Just sub subscribeServiceNotifications :: ServiceId -> (Int64, IdsHash) -> M s BrokerMsg @@ -1928,7 +1926,7 @@ client acknowledgeMsg :: MsgId -> StoreQueue s -> QueueRec -> M s (Transmission BrokerMsg) acknowledgeMsg msgId q qr = - liftIO (TM.lookupIO (subKey entId) $ subscriptions clnt) >>= \case + liftIO (TM.lookupIO entId $ subscriptions clnt) >>= \case Nothing -> pure $ err NO_MSG Just sub -> atomically (getDelivered sub) >>= \case @@ -2058,7 +2056,7 @@ client -- and delivery is cancelled - -- the new client will receive message in response to SUB. readTVar rcv - $>>= \rc@Client {subscriptions = subs, sndQ = sndQ'} -> TM.lookup (subKey rId) subs + $>>= \rc@Client {subscriptions = subs, sndQ = sndQ'} -> TM.lookup rId subs >>= maybe (newServiceDeliverySub subs) (pure . Just) $>>= \s@Sub {subThread, delivered} -> case subThread of ProhibitSub -> pure Nothing @@ -2075,7 +2073,7 @@ client newServiceDeliverySub subs | isJust (rcvServiceId qr) = do sub <- newSubscription NoSub - TM.insert (subKey rId) sub subs + TM.insert rId sub subs pure $ Just sub | otherwise = pure Nothing deliver sndQ' s ts = do @@ -2206,7 +2204,7 @@ client -- Possibly, the same should be done if the queue is suspended, but currently we do not use it -- queue is usually deleted by the same client that is currently subscribed, -- we delete subscription here, so the client with no subscriptions can be disconnected. - sub <- atomically $ TM.lookupDelete (subKey entId) $ subscriptions clnt + sub <- atomically $ TM.lookupDelete entId $ subscriptions clnt liftIO $ mapM_ cancelSub sub when (isJust rcvServiceId) $ atomically $ modifyTVar' (serviceSubsCount clnt) $ subtractServiceSubs (1, queueIdHash (recipientId q)) atomically $ writeTQueue (subQ subscribers) (CSDeleted entId rcvServiceId, clientId) @@ -2224,7 +2222,7 @@ client getQueueInfo :: StoreQueue s -> QueueRec -> M s BrokerMsg getQueueInfo q QueueRec {senderKey, notifier} = do fmap (either ERR INFO) $ liftIO $ runExceptT $ do - qiSub <- liftIO $ TM.lookupIO (subKey entId) (subscriptions clnt) >>= mapM mkQSub + qiSub <- liftIO $ TM.lookupIO entId (subscriptions clnt) >>= mapM mkQSub qiSize <- getQueueSize ms q qiMsg <- toMsgInfo <$$> tryPeekMsg ms q let info = QueueInfo {qiSnd = isJust senderKey, qiNtf = isJust notifier, qiSub, qiSize, qiMsg} diff --git a/src/Simplex/Messaging/Server/Env/STM.hs b/src/Simplex/Messaging/Server/Env/STM.hs index 703a2c85c..b4333959f 100644 --- a/src/Simplex/Messaging/Server/Env/STM.hs +++ b/src/Simplex/Messaging/Server/Env/STM.hs @@ -29,8 +29,6 @@ module Simplex.Messaging.Server.Env.STM Server (..), ServerSubscribers (..), SubscribedClients, - SubKey, - subKey, ProxyAgent (..), Client (..), ClientId, @@ -87,7 +85,6 @@ import Control.Logger.Simple import Control.Monad import qualified Crypto.PubKey.RSA as RSA import Crypto.Random -import Data.ByteString.Short (ShortByteString) import Data.Int (Int64) import Data.IntMap.Strict (IntMap) import qualified Data.IntMap.Strict as IM @@ -391,53 +388,45 @@ data ServerSubscribers s = ServerSubscribers -- any STM transaction that reads subscribed client will re-evaluate in this case. -- The subscriptions that were made at any point are not removed - -- this is a better trade-off with intermittently connected mobile clients. -data SubscribedClients s = SubscribedClients (TMap SubKey (TVar (Maybe (Client s)))) +data SubscribedClients s = SubscribedClients (TMap EntityId (TVar (Maybe (Client s)))) --- Subscription maps use unpinned keys. A pinned key keeps its whole pinned block alive, and with it --- the weak pointers and finalizers of dead crypto keys allocated in that block (~5 KB per subscription). -type SubKey = ShortByteString - -subKey :: EntityId -> SubKey -subKey = unEntityId -{-# INLINE subKey #-} - -getSubscribedClients :: SubscribedClients s -> IO (Map SubKey (TVar (Maybe (Client s)))) +getSubscribedClients :: SubscribedClients s -> IO (Map EntityId (TVar (Maybe (Client s)))) getSubscribedClients (SubscribedClients cs) = readTVarIO cs getSubscribedClient :: EntityId -> SubscribedClients s -> IO (Maybe (TVar (Maybe (Client s)))) -getSubscribedClient entId (SubscribedClients cs) = TM.lookupIO (subKey entId) cs +getSubscribedClient entId (SubscribedClients cs) = TM.lookupIO entId cs {-# INLINE getSubscribedClient #-} -- insert subscribed and current client, return previously subscribed client if it is different -upsertSubscribedClient :: SubKey -> Client s -> SubscribedClients s -> STM (Maybe (Client s)) -upsertSubscribedClient k c (SubscribedClients cs) = - TM.lookup k cs >>= \case - Nothing -> Nothing <$ TM.insertM k (newTVar (Just c)) cs +upsertSubscribedClient :: EntityId -> Client s -> SubscribedClients s -> STM (Maybe (Client s)) +upsertSubscribedClient entId c (SubscribedClients cs) = + TM.lookup entId cs >>= \case + Nothing -> Nothing <$ TM.insertM entId (newTVar (Just c)) cs Just cv -> readTVar cv >>= \case Just c' | sameClientId c c' -> pure Nothing c_ -> c_ <$ writeTVar cv (Just c) lookupSubscribedClient :: EntityId -> SubscribedClients s -> STM (Maybe (Client s)) -lookupSubscribedClient entId (SubscribedClients cs) = TM.lookup (subKey entId) cs $>>= readTVar +lookupSubscribedClient entId (SubscribedClients cs) = TM.lookup entId cs $>>= readTVar {-# INLINE lookupSubscribedClient #-} -- lookup and delete currently subscribed client lookupDeleteSubscribedClient :: EntityId -> SubscribedClients s -> STM (Maybe (Client s)) lookupDeleteSubscribedClient entId (SubscribedClients cs) = - TM.lookupDelete (subKey entId) cs $>>= (`swapTVar` Nothing) + TM.lookupDelete entId cs $>>= (`swapTVar` Nothing) {-# INLINE lookupDeleteSubscribedClient #-} -deleteSubcribedClient :: SubKey -> Client s -> SubscribedClients s -> IO () -deleteSubcribedClient k c (SubscribedClients cs) = +deleteSubcribedClient :: EntityId -> Client s -> SubscribedClients s -> IO () +deleteSubcribedClient entId c (SubscribedClients cs) = -- lookup of the subscribed client TVar can be in separate transaction, -- as long as the client is read in the same transaction - -- it prevents removing the next subscribed client and also avoids STM contention for the Map. - TM.lookupIO k cs >>= mapM_ (\cv -> atomically $ whenM (sameClient c cv) $ delete cv) + TM.lookupIO entId cs >>= mapM_ (\cv -> atomically $ whenM (sameClient c cv) $ delete cv) where delete cv = do writeTVar cv Nothing - TM.delete k cs + TM.delete entId cs sameClientId :: Client s -> (Client s) -> Bool sameClientId c c' = clientId c == clientId c' @@ -448,7 +437,7 @@ sameClient c cv = maybe False (sameClientId c) <$> readTVar cv {-# INLINE sameClient #-} data ClientSub - = CSClient QueueId SubKey (Maybe ServiceId) (Maybe ServiceId) -- includes the key in client subscriptions, previous and new associated service IDs + = CSClient QueueId (Maybe ServiceId) (Maybe ServiceId) -- includes previous and new associated service IDs | CSDeleted QueueId (Maybe ServiceId) -- includes previously associated service IDs | CSService ServiceId (Int64, IdsHash) -- only send END to idividual client subs on message delivery, not of SSUB/NSSUB @@ -460,8 +449,8 @@ type ClientId = Int data Client s = Client { clientId :: ClientId, - subscriptions :: TMap SubKey Sub, - ntfSubscriptions :: TMap SubKey (), + subscriptions :: TMap RecipientId Sub, + ntfSubscriptions :: TMap NotifierId (), serviceSubscribed :: TVar Bool, -- set independently of serviceSubsCount, to track whether service subscription command was received ntfServiceSubscribed :: TVar Bool, serviceSubsCount :: TVar (Int64, IdsHash), -- only one service can be subscribed, based on its certificate, this is subscription count