refactor/optimize server queue/message store (#340)

* refactor/optimize server queue/message store

* change fst to pattern match

* server store - wrap QueueRec into TVar
This commit is contained in:
Evgeny Poberezkin
2022-03-28 10:29:21 +01:00
committed by GitHub
parent c380431b94
commit 6ef6bedc03
5 changed files with 77 additions and 102 deletions
+6 -9
View File
@@ -111,15 +111,12 @@ newEnv config@ServerConfig {caCertificateFile, certificateFile, privateKeyFile}
return Env {config, server, serverIdentity, queueStore, msgStore, idsDrg, storeLog = s', tlsServerParams}
where
restoreQueues :: QueueStore -> StoreLog 'ReadMode -> m (StoreLog 'WriteMode)
restoreQueues queueStore s = do
(queues, s') <- liftIO $ readWriteStoreLog s
atomically $
modifyTVar' queueStore $ \d ->
d
{ queues,
senders = M.foldr' addSender M.empty queues,
notifiers = M.foldr' addNotifier M.empty queues
}
restoreQueues QueueStore {queues, senders, notifiers} s = do
(qs, s') <- liftIO $ readWriteStoreLog s
atomically $ do
writeTVar (TM.tVar queues) =<< mapM newTVar qs
writeTVar (TM.tVar senders) $ M.foldr' addSender M.empty qs
writeTVar (TM.tVar notifiers) $ M.foldr' addNotifier M.empty qs
pure s'
addSender :: QueueRec -> Map SenderId RecipientId -> Map SenderId RecipientId
addSender q = M.insert (senderId q) (recipientId q)
+8 -13
View File
@@ -6,36 +6,31 @@
module Simplex.Messaging.Server.MsgStore.STM where
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as M
import Numeric.Natural
import Simplex.Messaging.Protocol (RecipientId)
import Simplex.Messaging.Server.MsgStore
import Simplex.Messaging.TMap (TMap)
import qualified Simplex.Messaging.TMap as TM
import UnliftIO.STM
newtype MsgQueue = MsgQueue {msgQueue :: TBQueue Message}
newtype MsgStoreData = MsgStoreData {messages :: Map RecipientId MsgQueue}
type STMMsgStore = TVar MsgStoreData
type STMMsgStore = TMap RecipientId MsgQueue
newMsgStore :: STM STMMsgStore
newMsgStore = newTVar $ MsgStoreData M.empty
newMsgStore = TM.empty
instance MonadMsgStore STMMsgStore MsgQueue STM where
getMsgQueue :: STMMsgStore -> RecipientId -> Natural -> STM MsgQueue
getMsgQueue store rId quota = do
m <- messages <$> readTVar store
maybe (newQ m) return $ M.lookup rId m
getMsgQueue st rId quota = maybe newQ pure =<< TM.lookup rId st
where
newQ m' = do
newQ = do
q <- MsgQueue <$> newTBQueue quota
writeTVar store . MsgStoreData $ M.insert rId q m'
TM.insert rId q st
return q
delMsgQueue :: STMMsgStore -> RecipientId -> STM ()
delMsgQueue store rId =
modifyTVar' store $ MsgStoreData . M.delete rId . messages
delMsgQueue st rId = TM.delete rId st
instance MonadMsgQueue MsgQueue STM where
isFull :: MsgQueue -> STM Bool
@@ -15,6 +15,7 @@ data QueueRec = QueueRec
notifier :: Maybe (NotifierId, NtfPublicVerifyKey),
status :: QueueStatus
}
deriving (Eq, Show)
data QueueStatus = QueueActive | QueueOff deriving (Eq, Show)
+57 -80
View File
@@ -3,6 +3,7 @@
{-# LANGUAGE GADTs #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE RankNTypes #-}
@@ -11,107 +12,83 @@
module Simplex.Messaging.Server.QueueStore.STM where
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as M
import Control.Monad
import Data.Functor (($>))
import Simplex.Messaging.Protocol
import Simplex.Messaging.Server.QueueStore
import Simplex.Messaging.TMap (TMap)
import qualified Simplex.Messaging.TMap as TM
import Simplex.Messaging.Util (ifM)
import UnliftIO.STM
data QueueStoreData = QueueStoreData
{ queues :: Map RecipientId QueueRec,
senders :: Map SenderId RecipientId,
notifiers :: Map NotifierId RecipientId
data QueueStore = QueueStore
{ queues :: TMap RecipientId (TVar QueueRec),
senders :: TMap SenderId RecipientId,
notifiers :: TMap NotifierId RecipientId
}
type QueueStore = TVar QueueStoreData
newQueueStore :: STM QueueStore
newQueueStore = newTVar QueueStoreData {queues = M.empty, senders = M.empty, notifiers = M.empty}
newQueueStore = do
queues <- TM.empty
senders <- TM.empty
notifiers <- TM.empty
pure QueueStore {queues, senders, notifiers}
instance MonadQueueStore QueueStore STM where
addQueue :: QueueStore -> QueueRec -> STM (Either ErrorType ())
addQueue store qRec@QueueRec {recipientId = rId, senderId = sId} = do
cs@QueueStoreData {queues, senders} <- readTVar store
if M.member rId queues || M.member sId senders
then return $ Left DUPLICATE_
else do
writeTVar store $
cs
{ queues = M.insert rId qRec queues,
senders = M.insert sId rId senders
}
return $ Right ()
addQueue QueueStore {queues, senders} q@QueueRec {recipientId = rId, senderId = sId} = do
ifM hasId (pure $ Left DUPLICATE_) $ do
qVar <- newTVar q
TM.insert rId qVar queues
TM.insert sId rId senders
pure $ Right ()
where
hasId = (||) <$> TM.member rId queues <*> TM.member sId senders
getQueue :: QueueStore -> SParty p -> QueueId -> STM (Either ErrorType QueueRec)
getQueue st party qId = do
cs <- readTVar st
pure $ case party of
SRecipient -> getRcpQueue cs qId
SSender -> getPartyQueue cs senders
SNotifier -> getPartyQueue cs notifiers
getQueue QueueStore {queues, senders, notifiers} party qId =
toResult <$> (mapM readTVar =<< getVar)
where
getPartyQueue ::
QueueStoreData ->
(QueueStoreData -> Map QueueId RecipientId) ->
Either ErrorType QueueRec
getPartyQueue cs recipientIds =
case M.lookup qId $ recipientIds cs of
Just rId -> getRcpQueue cs rId
Nothing -> Left AUTH
getVar = case party of
SRecipient -> TM.lookup qId queues
SSender -> TM.lookup qId senders >>= get
SNotifier -> TM.lookup qId notifiers >>= get
get = fmap join . mapM (`TM.lookup` queues)
secureQueue :: QueueStore -> RecipientId -> SndPublicVerifyKey -> STM (Either ErrorType QueueRec)
secureQueue store rId sKey =
updateQueues store rId $ \cs c ->
case senderKey c of
Just _ -> (Left AUTH, cs)
_ -> (Right c, cs {queues = M.insert rId c {senderKey = Just sKey} (queues cs)})
secureQueue QueueStore {queues} rId sKey =
withQueue rId queues $ \qVar ->
readTVar qVar >>= \q -> case senderKey q of
Just _ -> pure Nothing
_ -> writeTVar qVar q {senderKey = Just sKey} $> Just q
addQueueNotifier :: QueueStore -> RecipientId -> NotifierId -> NtfPublicVerifyKey -> STM (Either ErrorType QueueRec)
addQueueNotifier store rId nId nKey = do
cs@QueueStoreData {queues, notifiers} <- readTVar store
if M.member nId notifiers
then pure $ Left DUPLICATE_
else case M.lookup rId queues of
Nothing -> pure $ Left AUTH
Just q -> case notifier q of
Just _ -> pure $ Left AUTH
addQueueNotifier QueueStore {queues, notifiers} rId nId nKey = do
ifM (TM.member nId notifiers) (pure $ Left DUPLICATE_) $
withQueue rId queues $ \qVar ->
readTVar qVar >>= \q -> case notifier q of
Just _ -> pure Nothing
_ -> do
writeTVar store $
cs
{ queues = M.insert rId q {notifier = Just (nId, nKey)} queues,
notifiers = M.insert nId rId notifiers
}
pure $ Right q
writeTVar qVar q {notifier = Just (nId, nKey)}
TM.insert nId rId notifiers
pure $ Just q
suspendQueue :: QueueStore -> RecipientId -> STM (Either ErrorType ())
suspendQueue store rId =
updateQueues store rId $ \cs c ->
(Right (), cs {queues = M.insert rId c {status = QueueOff} (queues cs)})
suspendQueue QueueStore {queues} rId =
withQueue rId queues $ \qVar -> modifyTVar' qVar (\q -> q {status = QueueOff}) $> Just ()
deleteQueue :: QueueStore -> RecipientId -> STM (Either ErrorType ())
deleteQueue store rId =
updateQueues store rId $ \cs c ->
( Right (),
cs
{ queues = M.delete rId (queues cs),
senders = M.delete (senderId c) (senders cs)
}
)
deleteQueue QueueStore {queues, senders, notifiers} rId = do
TM.lookupDelete rId queues >>= \case
Just qVar ->
readTVar qVar >>= \q -> do
TM.delete (senderId q) senders
forM_ (notifier q) $ \(nId, _) -> TM.delete nId notifiers
pure $ Right ()
_ -> pure $ Left AUTH
updateQueues ::
QueueStore ->
RecipientId ->
(QueueStoreData -> QueueRec -> (Either ErrorType a, QueueStoreData)) ->
STM (Either ErrorType a)
updateQueues store rId update = do
cs <- readTVar store
let conn = getRcpQueue cs rId
either (return . Left) (_update cs) conn
where
_update cs c = do
let (res, cs') = update cs c
writeTVar store cs'
return res
toResult :: Maybe a -> Either ErrorType a
toResult = maybe (Left AUTH) Right
getRcpQueue :: QueueStoreData -> RecipientId -> Either ErrorType QueueRec
getRcpQueue cs rId = maybe (Left AUTH) Right . M.lookup rId $ queues cs
withQueue :: RecipientId -> TMap RecipientId (TVar QueueRec) -> (TVar QueueRec -> STM (Maybe a)) -> STM (Either ErrorType a)
withQueue rId queues f = toResult <$> (TM.lookup rId queues >>= fmap join . mapM f)
+5
View File
@@ -2,6 +2,7 @@ module Simplex.Messaging.TMap
( TMap (..),
empty,
Simplex.Messaging.TMap.lookup,
member,
insert,
delete,
lookupInsert,
@@ -26,6 +27,10 @@ lookup :: Ord k => k -> TMap k a -> STM (Maybe a)
lookup k (TMap m) = M.lookup k <$> readTVar m
{-# INLINE lookup #-}
member :: Ord k => k -> TMap k a -> STM Bool
member k (TMap m) = M.member k <$> readTVar m
{-# INLINE member #-}
insert :: Ord k => k -> a -> TMap k a -> STM ()
insert k v (TMap m) = modifyTVar' m $ M.insert k v
{-# INLINE insert #-}