Revert "ntf server: optimize in-memory storage (#1516)"

This reverts commit afb338a.

Author: epoberezkin

src/Simplex/Messaging/Client/Agent.hs

@@ -358,7 +358,7 @@ smpSubscribeQueues party ca smp srv subs = do
       pending <- maybe (pure M.empty) readTVar =<< TM.lookup srv (pendingSrvSubs ca)
       let acc@(_, _, oks, notPending) = foldr (groupSub pending) (False, [], [], []) (L.zip subs rs)
       unless (null oks) $ addSubscriptions ca srv party oks
-      unless (null notPending) $ removePendingSubs ca srv party $ S.fromList notPending
+      unless (null notPending) $ removePendingSubs ca srv party notPending
       pure acc
     sessId = sessionId $ thParams smp
     groupSub :: Map SMPSub C.APrivateAuthKey -> ((QueueId, C.APrivateAuthKey), Either SMPClientError ()) -> (Bool, [(QueueId, SMPClientError)], [(QueueId, (SessionId, C.APrivateAuthKey))], [QueueId]) -> (Bool, [(QueueId, SMPClientError)], [(QueueId, (SessionId, C.APrivateAuthKey))], [QueueId])
@@ -412,22 +412,14 @@ removeSubscription :: SMPClientAgent -> SMPServer -> SMPSub -> STM ()
 removeSubscription = removeSub_ . srvSubs
 {-# INLINE removeSubscription #-}
 
-removePendingSub :: SMPClientAgent -> SMPServer -> SMPSub -> STM ()
-removePendingSub = removeSub_ . pendingSrvSubs
-{-# INLINE removePendingSub #-}
-
 removeSub_ :: TMap SMPServer (TMap SMPSub s) -> SMPServer -> SMPSub -> STM ()
 removeSub_ subs srv s = TM.lookup srv subs >>= mapM_ (TM.delete s)
 
-removeSubscriptions :: SMPClientAgent -> SMPServer -> SMPSubParty -> Set QueueId -> STM ()
-removeSubscriptions = removeSubs_ . srvSubs
-{-# INLINE removeSubscriptions #-}
-
-removePendingSubs :: SMPClientAgent -> SMPServer -> SMPSubParty -> Set QueueId -> STM ()
+removePendingSubs :: SMPClientAgent -> SMPServer -> SMPSubParty -> [QueueId] -> STM ()
 removePendingSubs = removeSubs_ . pendingSrvSubs
 {-# INLINE removePendingSubs #-}
 
-removeSubs_ :: TMap SMPServer (TMap SMPSub s) -> SMPServer -> SMPSubParty -> Set QueueId -> STM ()
+removeSubs_ :: TMap SMPServer (TMap SMPSub C.APrivateAuthKey) -> SMPServer -> SMPSubParty -> [QueueId] -> STM ()
 removeSubs_ subs srv party qs = TM.lookup srv subs >>= mapM_ (`modifyTVar'` (`M.withoutKeys` ss))
   where
-    ss = S.map (party,) qs
+    ss = S.fromList $ map (party,) qs

src/Simplex/Messaging/Notifications/Server.hs

@@ -691,10 +691,9 @@ client NtfServerClient {rcvQ, sndQ} NtfSubscriber {newSubQ, smpAgent = ca} NtfPu
           TDEL -> do
             logDebug "TDEL"
             st <- asks store
-            ss <- atomically $ deleteNtfToken st tknId
-            forM_ (M.assocs ss) $ \(smpServer, nIds) -> do
-              atomically $ removeSubscriptions ca smpServer SPNotifier nIds
-              atomically $ removePendingSubs ca smpServer SPNotifier nIds
+            qs <- atomically $ deleteNtfToken st tknId
+            forM_ qs $ \SMPQueueNtf {smpServer, notifierId} ->
+              atomically $ removeSubscription ca smpServer (SPNotifier, notifierId)
             cancelInvervalNotifications tknId
             withNtfLog (`logDeleteToken` tknId)
             incNtfStatT token tknDeleted
@@ -733,10 +732,9 @@ client NtfServerClient {rcvQ, sndQ} NtfSubscriber {newSubQ, smpAgent = ca} NtfPu
         subId <- getId
         sub <- atomically $ mkNtfSubData subId newSub
         resp <-
-          ifM
-            (atomically $ addNtfSubscription st subId sub)
-            (atomically (writeTBQueue newSubQ [NtfSub sub]) $> NRSubId subId)
-            (pure $ NRErr AUTH)
+          atomically (addNtfSubscription st subId sub) >>= \case
+            Just _ -> atomically (writeTBQueue newSubQ [NtfSub sub]) $> NRSubId subId
+            _ -> pure $ NRErr AUTH
         withNtfLog (`logCreateSubscription` sub)
         incNtfStat subCreated
         pure (corrId, NoEntity, resp)
@@ -758,7 +756,6 @@ client NtfServerClient {rcvQ, sndQ} NtfSubscriber {newSubQ, smpAgent = ca} NtfPu
             st <- asks store
             atomically $ deleteNtfSubscription st subId
             atomically $ removeSubscription ca smpServer (SPNotifier, notifierId)
-            atomically $ removePendingSub ca smpServer (SPNotifier, notifierId)
             withNtfLog (`logDeleteSubscription` subId)
             incNtfStat subDeleted
             pure NROk

src/Simplex/Messaging/Notifications/Server/Store.hs

@@ -8,7 +8,6 @@
 {-# LANGUAGE OverloadedLists #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE StandaloneDeriving #-}
 
 module Simplex.Messaging.Notifications.Server.Store where
 
@@ -17,16 +16,15 @@ import Control.Monad
 import Data.ByteString.Char8 (ByteString)
 import Data.Functor (($>))
 import Data.List.NonEmpty (NonEmpty (..), (<|))
-import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as M
-import Data.Maybe (isNothing)
+import Data.Maybe (catMaybes)
 import Data.Set (Set)
 import qualified Data.Set as S
 import Data.Word (Word16)
 import qualified Simplex.Messaging.Crypto as C
 import Simplex.Messaging.Encoding.String
 import Simplex.Messaging.Notifications.Protocol
-import Simplex.Messaging.Protocol (NotifierId, NtfPrivateAuthKey, NtfPublicAuthKey, SMPServer)
+import Simplex.Messaging.Protocol (NtfPrivateAuthKey, NtfPublicAuthKey, SMPServer)
 import Simplex.Messaging.Server.QueueStore (RoundedSystemTime)
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
@@ -37,11 +35,8 @@ data NtfStore = NtfStore
     -- multiple registrations exist to protect from malicious registrations if token is compromised
     tokenRegistrations :: TMap DeviceToken (TMap ByteString NtfTokenId),
     subscriptions :: TMap NtfSubscriptionId NtfSubData,
-    -- the first set is used to delete from `subscriptions` when token is deleted, the second - to cancel SMP subsriptions.
-    -- TODO [notifications] it can be simplified once NtfSubData is fully removed.
-    tokenSubscriptions :: TMap NtfTokenId (TMap SMPServer (TVar (Set NtfSubscriptionId), TVar (Set NotifierId))),
-    -- TODO [notifications] for subscriptions that "migrated" to server subscription, we may replace NtfSubData with NtfTokenId here (Either NtfSubData NtfTokenId).
-    subscriptionLookup :: TMap SMPServer (TMap NotifierId NtfSubData),
+    tokenSubscriptions :: TMap NtfTokenId (TVar (Set NtfSubscriptionId)),
+    subscriptionLookup :: TMap SMPQueueNtf NtfSubscriptionId,
     tokenLastNtfs :: TMap NtfTokenId (TVar (NonEmpty PNMessageData))
   }
 
@@ -139,7 +134,7 @@ removeTokenRegistration st NtfTknData {ntfTknId = tId, token, tknVerifyKey} =
         >>= mapM_ (\tId' -> when (tId == tId') $ TM.delete k regs)
     k = C.toPubKey C.pubKeyBytes tknVerifyKey
 
-deleteNtfToken :: NtfStore -> NtfTokenId -> STM (Map SMPServer (Set NotifierId))
+deleteNtfToken :: NtfStore -> NtfTokenId -> STM [SMPQueueNtf]
 deleteNtfToken st tknId = do
   void $
     TM.lookupDelete tknId (tokens st) $>>= \NtfTknData {token, tknVerifyKey} ->
@@ -152,25 +147,25 @@ deleteNtfToken st tknId = do
     regs = tokenRegistrations st
     regKey = C.toPubKey C.pubKeyBytes
 
-deleteTokenSubs :: NtfStore -> NtfTokenId -> STM (Map SMPServer (Set NotifierId))
-deleteTokenSubs st tknId =
-  TM.lookupDelete tknId (tokenSubscriptions st)
-    >>= maybe (pure M.empty) (readTVar >=> deleteSrvSubs)
+deleteTokenSubs :: NtfStore -> NtfTokenId -> STM [SMPQueueNtf]
+deleteTokenSubs st tknId = do
+  qs <-
+    TM.lookupDelete tknId (tokenSubscriptions st)
+      >>= mapM (readTVar >=> mapM deleteSub . S.toList)
+  pure $ maybe [] catMaybes qs
   where
-    deleteSrvSubs :: Map SMPServer (TVar (Set NtfSubscriptionId), TVar (Set NotifierId)) -> STM (Map SMPServer (Set NotifierId))
-    deleteSrvSubs = M.traverseWithKey $ \smpServer (sVar, nVar) -> do
-      sIds <- readTVar sVar
-      modifyTVar' (subscriptions st) (`M.withoutKeys` sIds)
-      nIds <- readTVar nVar
-      TM.lookup smpServer (subscriptionLookup st) >>= mapM_ (`modifyTVar'` (`M.withoutKeys` nIds))
-      pure nIds    
+    deleteSub subId = do
+      TM.lookupDelete subId (subscriptions st)
+        $>>= \NtfSubData {smpQueue} ->
+          TM.delete smpQueue (subscriptionLookup st) $> Just smpQueue
 
 getNtfSubscriptionIO :: NtfStore -> NtfSubscriptionId -> IO (Maybe NtfSubData)
 getNtfSubscriptionIO st subId = TM.lookupIO subId (subscriptions st)
 
 findNtfSubscription :: NtfStore -> SMPQueueNtf -> STM (Maybe NtfSubData)
-findNtfSubscription st SMPQueueNtf {smpServer, notifierId} =
-  TM.lookup smpServer (subscriptionLookup st) $>>= TM.lookup notifierId 
+findNtfSubscription st smpQueue = do
+  TM.lookup smpQueue (subscriptionLookup st)
+    $>>= \subId -> TM.lookup subId (subscriptions st)
 
 findNtfSubscriptionToken :: NtfStore -> SMPQueueNtf -> STM (Maybe NtfTknData)
 findNtfSubscriptionToken st smpQueue = do
@@ -188,45 +183,30 @@ mkNtfSubData ntfSubId (NewNtfSub tokenId smpQueue notifierKey) = do
   subStatus <- newTVar NSNew
   pure NtfSubData {ntfSubId, smpQueue, tokenId, subStatus, notifierKey}
 
--- returns False if subscription existed before
-addNtfSubscription :: NtfStore -> NtfSubscriptionId -> NtfSubData -> STM Bool
-addNtfSubscription st subId sub@NtfSubData {smpQueue = SMPQueueNtf {smpServer, notifierId}, tokenId} =
-  TM.lookup tokenId (tokenSubscriptions st)
-    >>= maybe newTokenSubs pure
-    >>= \ts -> TM.lookup smpServer ts
-    >>= maybe (newTokenSrvSubs ts) pure
-    >>= insertSub
+addNtfSubscription :: NtfStore -> NtfSubscriptionId -> NtfSubData -> STM (Maybe ())
+addNtfSubscription st subId sub@NtfSubData {smpQueue, tokenId} =
+  TM.lookup tokenId (tokenSubscriptions st) >>= maybe newTokenSub pure >>= insertSub
   where
-    newTokenSubs = do
-      ts <- newTVar M.empty
+    newTokenSub = do
+      ts <- newTVar S.empty
       TM.insert tokenId ts $ tokenSubscriptions st
       pure ts
-    newTokenSrvSubs ts = do
-      tss <- (,) <$> newTVar S.empty <*> newTVar S.empty
-      TM.insert smpServer tss ts
-      pure tss
-    insertSub  :: (TVar (Set NtfSubscriptionId), TVar (Set NotifierId)) -> STM Bool
-    insertSub (sIds, nIds) = do
-      modifyTVar' sIds $ S.insert subId
-      modifyTVar' nIds $ S.insert notifierId
+    insertSub ts = do
+      modifyTVar' ts $ S.insert subId
       TM.insert subId sub $ subscriptions st
-      TM.lookup smpServer (subscriptionLookup st)
-        >>= maybe newSubs pure
-        >>= fmap isNothing . TM.lookupInsert notifierId sub
-    newSubs = do
-      ss <- newTVar M.empty
-      TM.insert smpServer ss $ subscriptionLookup st
-      pure ss
+      TM.insert smpQueue subId (subscriptionLookup st)
+      -- return Nothing if subscription existed before
+      pure $ Just ()
 
 deleteNtfSubscription :: NtfStore -> NtfSubscriptionId -> STM ()
-deleteNtfSubscription st subId = TM.lookupDelete subId (subscriptions st) >>= mapM_ deleteSubIndices
-  where
-    deleteSubIndices NtfSubData {smpQueue = SMPQueueNtf {smpServer, notifierId}, tokenId} = do
-      TM.lookup smpServer (subscriptionLookup st) >>= mapM_ (TM.delete notifierId)
-      tss_ <- TM.lookup tokenId (tokenSubscriptions st) $>>= TM.lookup smpServer
-      forM_ tss_ $ \(sIds, nIds) -> do
-        modifyTVar' sIds $ S.delete subId
-        modifyTVar' nIds $ S.delete notifierId
+deleteNtfSubscription st subId = do
+  TM.lookupDelete subId (subscriptions st)
+    >>= mapM_
+      ( \NtfSubData {smpQueue, tokenId} -> do
+          TM.delete smpQueue $ subscriptionLookup st
+          ts_ <- TM.lookup tokenId (tokenSubscriptions st)
+          forM_ ts_ $ \ts -> modifyTVar' ts $ S.delete subId
+      )
 
 addTokenLastNtf :: NtfStore -> NtfTokenId -> PNMessageData -> IO (NonEmpty PNMessageData)
 addTokenLastNtf st tknId newNtf =

tests/ServerTests/SchemaDump.hs

@@ -18,7 +18,7 @@ import Simplex.Messaging.Server.QueueStore.Postgres.Migrations (serverMigrations
 import Simplex.Messaging.Util (ifM)
 import System.Directory (doesFileExist, removeFile)
 import System.Environment (lookupEnv)
-import System.Process (readCreateProcess, shell)
+import System.Process (readCreateProcess, readCreateProcessWithExitCode, shell)
 import Test.Hspec
 
 testDBSchema :: B.ByteString
@@ -87,7 +87,10 @@ getSchema schemaPath = do
         ("pg_dump " <> B.unpack testServerDBConnstr <> " --schema " <> B.unpack testDBSchema)
           <> " --schema-only --no-owner --no-privileges --no-acl --no-subscriptions --no-tablespaces > "
           <> schemaPath
-  void $ readCreateProcess (shell cmd) ""
+  (code, out, err) <- readCreateProcessWithExitCode (shell cmd) ""
+  print code
+  putStrLn $ "out: " <> out
+  putStrLn $ "err: " <> err
   threadDelay 20000
   let sed = (if ci then "sed -i" else "sed -i ''")
   void $ readCreateProcess (shell $ sed <> " '/^--/d' " <> schemaPath) ""