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

This reverts commit afb338a41a.
This commit is contained in:
Evgeny Poberezkin
2025-04-21 21:01:24 +01:00
parent aa9b93eee5
commit 6b60f8bab6
4 changed files with 51 additions and 79 deletions
+4 -12
View File
@@ -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
@@ -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
@@ -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 =
+5 -2
View File
@@ -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) ""