smp server: do not cache all queues from database while processing expirations (#1483)

* smp server: expire only active queues

* version

* do not cache all queues while processing expirations

* refactor

* foldWithOptions_

* version

* use shared lock when expiring all queues

* use TMVar

* comment

* rename

* remove fold options

* do not create locks in the Map for temporarily loaded queues

* fix

* revert version

Author: epoberezkin

simplexmq.cabal

@@ -228,6 +228,7 @@ library
           Simplex.Messaging.Server.Main
           Simplex.Messaging.Server.MsgStore
           Simplex.Messaging.Server.MsgStore.Journal
+          Simplex.Messaging.Server.MsgStore.Journal.SharedLock
           Simplex.Messaging.Server.MsgStore.STM
           Simplex.Messaging.Server.MsgStore.Types
           Simplex.Messaging.Server.NtfStore

src/Simplex/Messaging/Agent/Lock.hs

@@ -6,6 +6,7 @@ module Simplex.Messaging.Agent.Lock
     withLock',
     withGetLock,
     withGetLocks,
+    getPutLock,
   )
 where
 

src/Simplex/Messaging/Server.hs

@@ -97,7 +97,7 @@ import Simplex.Messaging.Server.Control
 import Simplex.Messaging.Server.Env.STM as Env
 import Simplex.Messaging.Server.Expiration
 import Simplex.Messaging.Server.MsgStore
-import Simplex.Messaging.Server.MsgStore.Journal (JournalMsgStore, JournalQueue, closeMsgQueue)
+import Simplex.Messaging.Server.MsgStore.Journal (JournalMsgStore, JournalQueue)
 import Simplex.Messaging.Server.MsgStore.STM
 import Simplex.Messaging.Server.MsgStore.Types
 import Simplex.Messaging.Server.NtfStore
@@ -404,11 +404,11 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt
           old <- expireBeforeEpoch expCfg
           now <- systemSeconds <$> getSystemTime
           msgStats@MessageStats {storedMsgsCount = stored, expiredMsgsCount = expired} <-
-            withAllMsgQueues False ms $ expireQueueMsgs now ms old
+            withAllMsgQueues False "idleDeleteExpiredMsgs" ms $ expireQueueMsgs now ms old
           atomicWriteIORef (msgCount stats) stored
           atomicModifyIORef'_ (msgExpired stats) (+ expired)
           printMessageStats "STORE: messages" msgStats
-        expireQueueMsgs now ms old q = fmap (fromRight newMessageStats) . runExceptT $ do
+        expireQueueMsgs now ms old q = do
           (expired_, stored) <- idleDeleteExpiredMsgs now ms q old
           pure MessageStats {storedMsgsCount = stored, expiredMsgsCount = fromMaybe 0 expired_, storedQueues = 1}
 
@@ -1806,19 +1806,18 @@ exportMessages :: MsgStoreClass s => Bool -> s -> FilePath -> Bool -> IO ()
 exportMessages tty ms f drainMsgs = do
   logInfo $ "saving messages to file " <> T.pack f
   liftIO $ withFile f WriteMode $ \h ->
-    tryAny (withAllMsgQueues tty ms $ saveQueueMsgs h) >>= \case
+    tryAny (unsafeWithAllMsgQueues tty ms $ saveQueueMsgs h) >>= \case
       Right (Sum total) -> logInfo $ "messages saved: " <> tshow total
       Left e -> do
         logError $ "error exporting messages: " <> tshow e
         exitFailure
   where
     saveQueueMsgs h q = do
-      let rId = recipientId q
-      runExceptT (getQueueMessages drainMsgs ms q) >>= \case
-        Right msgs -> Sum (length msgs) <$ BLD.hPutBuilder h (encodeMessages rId msgs)
-        Left e -> do
-          logError $ "STORE: saveQueueMsgs, error exporting messages from queue " <> decodeLatin1 (strEncode rId) <> ", " <> tshow e
-          exitFailure
+      msgs <-
+        unsafeRunStore q "saveQueueMsgs" $
+          getQueueMessages_ drainMsgs q =<< getMsgQueue ms q False
+      BLD.hPutBuilder h $ encodeMessages (recipientId q) msgs
+      pure $ Sum $ length msgs
     encodeMessages rId = mconcat . map (\msg -> BLD.byteString (strEncode $ MLRv3 rId msg) <> BLD.char8 '\n')
 
 processServerMessages :: StartOptions -> M (Maybe MessageStats)
@@ -1838,33 +1837,23 @@ processServerMessages StartOptions {skipWarnings} = do
         | expire = Just <$> case old_ of
             Just old -> do
               logInfo "expiring journal store messages..."
-              withAllMsgQueues False ms $ processExpireQueue old
+              run $ processExpireQueue old
             Nothing -> do
               logInfo "validating journal store messages..."
-              withAllMsgQueues False ms $ processValidateQueue
+              run processValidateQueue
         | otherwise = logWarn "skipping message expiration" $> Nothing
         where
+          run a = unsafeWithAllMsgQueues False ms a `catchAny` \_ -> exitFailure
           processExpireQueue :: Int64 -> JournalQueue s -> IO MessageStats
-          processExpireQueue old q =
-            runExceptT expireQueue >>= \case
-              Right (storedMsgsCount, expiredMsgsCount) ->
-                pure MessageStats {storedMsgsCount, expiredMsgsCount, storedQueues = 1}
-              Left e -> do
-                logError $ "STORE: processExpireQueue, failed expiring messages in queue, " <> tshow e
-                exitFailure
-            where
-              expireQueue = do
-                expired'' <- deleteExpiredMsgs ms q old
-                stored'' <- getQueueSize ms q
-                liftIO $ closeMsgQueue q
-                pure (stored'', expired'')
+          processExpireQueue old q = unsafeRunStore q "processExpireQueue" $ do
+            mq <- getMsgQueue ms q False
+            expiredMsgsCount <- deleteExpireMsgs_ old q mq
+            storedMsgsCount <- getQueueSize_ mq
+            pure MessageStats {storedMsgsCount, expiredMsgsCount, storedQueues = 1}
           processValidateQueue :: JournalQueue s -> IO MessageStats
-          processValidateQueue q =
-            runExceptT (getQueueSize ms q) >>= \case
-              Right storedMsgsCount -> pure newMessageStats {storedMsgsCount, storedQueues = 1}
-              Left e -> do
-                logError $ "STORE: processValidateQueue, failed opening message queue, " <> tshow e
-                exitFailure
+          processValidateQueue q = unsafeRunStore q "processValidateQueue" $ do
+            storedMsgsCount <- getQueueSize_ =<< getMsgQueue ms q False
+            pure newMessageStats {storedMsgsCount, storedQueues = 1}
 
 importMessages :: forall s. MsgStoreClass s => Bool -> s -> FilePath -> Maybe Int64 -> Bool -> IO MessageStats
 importMessages tty ms f old_ skipWarnings  = do

src/Simplex/Messaging/Server/MsgStore/Journal.hs

@@ -55,6 +55,7 @@ import Control.Monad.Trans.Except
 import qualified Data.Attoparsec.ByteString.Char8 as A
 import Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Char8 as B
+import Data.Either (fromRight)
 import Data.Functor (($>))
 import Data.Int (Int64)
 import Data.List (intercalate, sort)
@@ -65,10 +66,11 @@ import Data.Time.Clock (NominalDiffTime, UTCTime, addUTCTime, getCurrentTime)
 import Data.Time.Clock.System (SystemTime (..), getSystemTime)
 import Data.Time.Format.ISO8601 (iso8601Show, iso8601ParseM)
 import GHC.IO (catchAny)
-import Simplex.Messaging.Agent.Client (getMapLock, withLockMap)
+import Simplex.Messaging.Agent.Client (getMapLock)
 import Simplex.Messaging.Agent.Lock
 import Simplex.Messaging.Encoding.String
 import Simplex.Messaging.Protocol
+import Simplex.Messaging.Server.MsgStore.Journal.SharedLock
 import Simplex.Messaging.Server.MsgStore.Types
 import Simplex.Messaging.Server.QueueStore
 import Simplex.Messaging.Server.QueueStore.Postgres
@@ -87,6 +89,7 @@ data JournalMsgStore s = JournalMsgStore
   { config :: JournalStoreConfig s,
     random :: TVar StdGen,
     queueLocks :: TMap RecipientId Lock,
+    sharedLock :: TMVar RecipientId,
     queueStore_ :: QStore s,
     expireBackupsBefore :: UTCTime
   }
@@ -138,6 +141,7 @@ data QStoreCfg s where
 data JournalQueue (s :: QSType) = JournalQueue
   { recipientId' :: RecipientId,
     queueLock :: Lock,
+    sharedLock :: TMVar RecipientId,
     -- To avoid race conditions and errors when restoring queues,
     -- Nothing is written to TVar when queue is deleted.
     queueRec' :: TVar (Maybe QueueRec),
@@ -276,7 +280,8 @@ instance StoreQueueClass (JournalQueue s) where
   msgQueue = msgQueue'
   {-# INLINE msgQueue #-}
   withQueueLock :: JournalQueue s -> String -> IO a -> IO a
-  withQueueLock = withLock' . queueLock
+  withQueueLock JournalQueue {recipientId', queueLock, sharedLock} =
+    withLockWaitShared recipientId' queueLock sharedLock
   {-# INLINE withQueueLock #-}
 
 instance QueueStoreClass (JournalQueue s) (QStore s) where
@@ -316,6 +321,27 @@ instance QueueStoreClass (JournalQueue s) (QStore s) where
   deleteStoreQueue = withQS deleteStoreQueue
   {-# INLINE deleteStoreQueue #-}
 
+mkTempQueue :: JournalMsgStore s -> RecipientId -> QueueRec -> IO (JournalQueue s)
+mkTempQueue ms rId qr = createLockIO >>= makeQueue_ ms rId qr
+{-# INLINE mkTempQueue #-}
+
+makeQueue_ :: JournalMsgStore s -> RecipientId -> QueueRec -> Lock -> IO (JournalQueue s)
+makeQueue_ JournalMsgStore {sharedLock} rId qr queueLock = do
+  queueRec' <- newTVarIO $ Just qr
+  msgQueue' <- newTVarIO Nothing
+  activeAt <- newTVarIO 0
+  queueState <- newTVarIO Nothing
+  pure $
+    JournalQueue
+      { recipientId' = rId,
+        queueLock,
+        sharedLock,
+        queueRec',
+        msgQueue',
+        activeAt,
+        queueState
+      }
+
 instance MsgStoreClass (JournalMsgStore s) where
   type StoreMonad (JournalMsgStore s) = StoreIO s
   type QueueStore (JournalMsgStore s) = QStore s
@@ -326,9 +352,10 @@ instance MsgStoreClass (JournalMsgStore s) where
   newMsgStore config@JournalStoreConfig {queueStoreCfg} = do
     random <- newTVarIO =<< newStdGen
     queueLocks <- TM.emptyIO
+    sharedLock <- newEmptyTMVarIO
     queueStore_ <- newQueueStore @(JournalQueue s) queueStoreCfg
     expireBackupsBefore <- addUTCTime (- expireBackupsAfter config) <$> getCurrentTime
-    pure JournalMsgStore {config, random, queueLocks, queueStore_, expireBackupsBefore}
+    pure JournalMsgStore {config, random, queueLocks, sharedLock, queueStore_, expireBackupsBefore}
 
   closeMsgStore :: JournalMsgStore s -> IO ()
   closeMsgStore ms = do
@@ -341,10 +368,32 @@ instance MsgStoreClass (JournalMsgStore s) where
   withActiveMsgQueues :: Monoid a => JournalMsgStore s -> (JournalQueue s -> IO a) -> IO a
   withActiveMsgQueues = withQS withLoadedQueues . queueStore_
 
-  withAllMsgQueues :: Monoid a => Bool -> JournalMsgStore s -> (JournalQueue s -> IO a) -> IO a
-  withAllMsgQueues tty ms action = case queueStore_ ms of
-    MQStore st -> withLoadedQueues st action
-    PQStore st -> foldQueues tty st (mkQueue ms) action
+  -- This function can only be used in server CLI commands or before server is started.
+  -- It does not cache queues and is NOT concurrency safe.
+  unsafeWithAllMsgQueues :: Monoid a => Bool -> JournalMsgStore s -> (JournalQueue s -> IO a) -> IO a
+  unsafeWithAllMsgQueues tty ms action = case queueStore_ ms of
+    MQStore st -> withLoadedQueues st run 
+    PQStore st -> foldQueueRecs tty st $ uncurry (mkTempQueue ms) >=> run
+    where
+      run q = do
+        r <- action q
+        closeMsgQueue q
+        pure r
+
+  -- This function is concurrency safe, it is used to expire queues.
+  withAllMsgQueues :: forall a. Monoid a => Bool -> String -> JournalMsgStore s -> (JournalQueue s -> StoreIO s a) -> IO a
+  withAllMsgQueues tty op ms@JournalMsgStore {queueLocks, sharedLock} action = case queueStore_ ms of
+    MQStore st ->
+      withLoadedQueues st $ \q ->
+        run $ isolateQueue q op $ action q
+    PQStore st ->
+      foldQueueRecs tty st $ \(rId, qr) -> do
+        q <- mkTempQueue ms rId qr
+        withSharedWaitLock rId queueLocks sharedLock $
+          run $ tryStore' op rId $ unStoreIO $ action q
+    where
+      run :: ExceptT ErrorType IO a -> IO a
+      run = fmap (fromRight mempty) . runExceptT
 
   logQueueStates :: JournalMsgStore s -> IO ()
   logQueueStates ms = withActiveMsgQueues ms $ unStoreIO . logQueueState
@@ -361,20 +410,11 @@ instance MsgStoreClass (JournalMsgStore s) where
 
   mkQueue :: JournalMsgStore s -> RecipientId -> QueueRec -> IO (JournalQueue s)
   mkQueue ms rId qr = do
-    queueLock <- atomically $ getMapLock (queueLocks ms) rId
-    queueRec' <- newTVarIO $ Just qr
-    msgQueue' <- newTVarIO Nothing
-    activeAt <- newTVarIO 0
-    queueState <- newTVarIO Nothing
-    pure $
-      JournalQueue
-        { recipientId' = rId,
-          queueLock,
-          queueRec',
-          msgQueue',
-          activeAt,
-          queueState
-        }
+    lock <- atomically $ getMapLock (queueLocks ms) rId
+    makeQueue_ ms rId qr lock
+
+  getLoadedQueue :: JournalMsgStore s -> JournalQueue s -> StoreIO s (JournalQueue s)
+  getLoadedQueue ms sq = StoreIO $ fromMaybe sq <$> TM.lookupIO (recipientId sq) (loadedQueues $ queueStore_ ms)
 
   getMsgQueue :: JournalMsgStore s -> JournalQueue s -> Bool -> StoreIO s (JournalMsgQueue s)
   getMsgQueue ms@JournalMsgStore {random} q'@JournalQueue {recipientId' = rId, msgQueue'} forWrite =
@@ -546,8 +586,11 @@ instance MsgStoreClass (JournalMsgStore s) where
         $>>= \hs -> updateReadPos q mq logState len hs $> Just ()
 
   isolateQueue :: JournalQueue s -> String -> StoreIO s a -> ExceptT ErrorType IO a
-  isolateQueue JournalQueue {recipientId' = rId, queueLock} op a =
-    tryStore' op rId $ withLock' queueLock op $ unStoreIO a
+  isolateQueue sq op = tryStore' op (recipientId' sq) . withQueueLock sq op . unStoreIO
+
+  unsafeRunStore :: JournalQueue s -> String -> StoreIO s a -> IO a
+  unsafeRunStore sq op a =
+    unStoreIO a `E.catch` \e -> storeError op (recipientId' sq) e >> E.throwIO e
 
 updateActiveAt :: JournalQueue s -> IO ()
 updateActiveAt q = atomically . writeTVar (activeAt q) . systemSeconds =<< getSystemTime
@@ -556,15 +599,16 @@ tryStore' :: String -> RecipientId -> IO a -> ExceptT ErrorType IO a
 tryStore' op rId = tryStore op rId . fmap Right
 
 tryStore :: forall a. String -> RecipientId -> IO (Either ErrorType a) -> ExceptT ErrorType IO a
-tryStore op rId a = ExceptT $ E.mask_ $ E.try a >>= either storeErr pure
-  where
-    storeErr :: E.SomeException -> IO (Either ErrorType a)
-    storeErr e =
-      let e' = intercalate ", " [op, B.unpack $ strEncode rId, show e]
-       in logError ("STORE: " <> T.pack e') $> Left (STORE e')
+tryStore op rId a = ExceptT $ E.mask_ $ a `E.catch` storeError op rId
+
+storeError :: String -> RecipientId -> E.SomeException -> IO (Either ErrorType a)
+storeError op rId e =
+  let e' = intercalate ", " [op, B.unpack $ strEncode rId, show e]
+    in logError ("STORE: " <> T.pack e') $> Left (STORE e')
 
 isolateQueueId :: String -> JournalMsgStore s -> RecipientId -> IO (Either ErrorType a) -> ExceptT ErrorType IO a
-isolateQueueId op ms rId = tryStore op rId . withLockMap (queueLocks ms) rId op
+isolateQueueId op JournalMsgStore {queueLocks, sharedLock} rId =
+  tryStore op rId . withLockMapWaitShared rId queueLocks sharedLock op
 
 openMsgQueue :: JournalMsgStore s -> JMQueue -> Bool -> IO (JournalMsgQueue s)
 openMsgQueue ms@JournalMsgStore {config} q@JMQueue {queueDirectory = dir, statePath} forWrite = do

src/Simplex/Messaging/Server/MsgStore/Journal/SharedLock.hs

@@ -0,0 +1,43 @@
+module Simplex.Messaging.Server.MsgStore.Journal.SharedLock
+  ( withLockWaitShared,
+    withLockMapWaitShared,
+    withSharedWaitLock,
+  )
+where
+
+import Control.Concurrent.STM
+import qualified Control.Exception as E
+import Control.Monad
+import Simplex.Messaging.Agent.Lock
+import Simplex.Messaging.Agent.Client (getMapLock)
+import Simplex.Messaging.Protocol (RecipientId)
+import Simplex.Messaging.TMap (TMap)
+import qualified Simplex.Messaging.TMap as TM
+import Simplex.Messaging.Util (($>>), ($>>=))
+
+-- wait until shared lock with passed ID is released and take lock
+withLockWaitShared :: RecipientId -> Lock -> TMVar RecipientId -> String -> IO a -> IO a
+withLockWaitShared rId lock shared name =
+  E.bracket_
+    (atomically $ waitShared rId shared >> putTMVar lock name)
+    (void $ atomically $ takeTMVar lock)
+
+-- wait until shared lock with passed ID is released and take lock from Map for this ID
+withLockMapWaitShared :: RecipientId -> TMap RecipientId Lock -> TMVar RecipientId -> String -> IO a -> IO a
+withLockMapWaitShared rId locks shared name a =
+  E.bracket
+    (atomically $ waitShared rId shared >> getPutLock (getMapLock locks) rId name)
+    (atomically . takeTMVar)
+    (const a)
+
+waitShared :: RecipientId -> TMVar RecipientId -> STM ()
+waitShared rId shared = tryReadTMVar shared >>= mapM_ (\rId' -> when (rId == rId') retry)
+
+-- wait until lock with passed ID in Map is released and take shared lock for this ID
+withSharedWaitLock :: RecipientId -> TMap RecipientId Lock -> TMVar RecipientId -> IO a -> IO a
+withSharedWaitLock rId locks shared =
+  E.bracket_
+    (atomically $ waitLock >> putTMVar shared rId)
+    (atomically $ takeTMVar shared)
+  where
+    waitLock = TM.lookup rId locks $>>= tryReadTMVar $>> retry

src/Simplex/Messaging/Server/MsgStore/STM.hs

@@ -80,7 +80,9 @@ instance MsgStoreClass STMMsgStore where
   {-# INLINE closeMsgStore #-}
   withActiveMsgQueues = withLoadedQueues . queueStore_
   {-# INLINE withActiveMsgQueues #-}
-  withAllMsgQueues _ = withLoadedQueues . queueStore_
+  unsafeWithAllMsgQueues _ = withLoadedQueues . queueStore_
+  {-# INLINE unsafeWithAllMsgQueues #-}
+  withAllMsgQueues _tty _op ms action = withLoadedQueues (queueStore_ ms) $ atomically . action
   {-# INLINE withAllMsgQueues #-}
   logQueueStates _ = pure ()
   {-# INLINE logQueueStates #-}
@@ -92,6 +94,10 @@ instance MsgStoreClass STMMsgStore where
   mkQueue _ rId qr = STMQueue rId <$> newTVarIO (Just qr) <*> newTVarIO Nothing
   {-# INLINE mkQueue #-}
 
+  getLoadedQueue :: STMMsgStore -> STMQueue -> STM STMQueue
+  getLoadedQueue _ = pure
+  {-# INLINE getLoadedQueue #-}
+
   getMsgQueue :: STMMsgStore -> STMQueue -> Bool -> STM STMMsgQueue
   getMsgQueue _ STMQueue {msgQueue'} _ = readTVar msgQueue' >>= maybe newQ pure
     where
@@ -168,3 +174,8 @@ instance MsgStoreClass STMMsgStore where
 
   isolateQueue :: STMQueue -> String -> STM a -> ExceptT ErrorType IO a
   isolateQueue _ _ = liftIO . atomically
+  {-# INLINE isolateQueue #-}
+
+  unsafeRunStore :: STMQueue -> String -> STM a -> IO a
+  unsafeRunStore _ _ = atomically
+  {-# INLINE unsafeRunStore #-}