smp-server: key subscriptions by EntityId

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}

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