expire messages concurrently

src/Simplex/Messaging/Server.hs

@@ -1748,7 +1748,7 @@ saveServerMessages drainMsgs =
 exportMessages :: MsgStoreClass s => s -> FilePath -> Bool -> IO ()
 exportMessages ms f drainMsgs = do
   logInfo $ "saving messages to file " <> T.pack f
-  Sum total <- liftIO $ withFile f WriteMode $ withAllMsgQueues ms . saveQueueMsgs
+  Sum total <- liftIO $ withFile f WriteMode $ withAllMsgQueues 1 ms . saveQueueMsgs
   logInfo $ "messages saved: " <> tshow total
   where
     saveQueueMsgs h rId q = getQueueMessages drainMsgs q >>= \msgs -> Sum (length msgs) <$ BLD.hPutBuilder h (encodeMessages rId msgs)
@@ -1764,13 +1764,15 @@ processServerMessages = do
         AMS SMSMemory ms@STMMsgStore {storeConfig = STMStoreConfig {storePath}} -> case storePath of
           Just f -> ifM (doesFileExist f) (importMessages ms f old_) (pure newMessageStats)
           Nothing -> pure newMessageStats
-        AMS SMSJournal ms -> case old_ of
-          Just old -> do
-            logInfo "expiring journal store messages..."
-            withAllMsgQueues ms $ \_ -> processExpireQueue old
-          Nothing -> do
-            logInfo "validating journal store messages..."
-            withAllMsgQueues ms $ \_ -> processValidateQueue
+        AMS SMSJournal ms -> do
+          n <- getNumCapabilities
+          case old_ of
+            Just old -> do
+              logInfo "expiring journal store messages..."
+              withAllMsgQueues n ms $ \_ -> processExpireQueue old
+            Nothing -> do
+              logInfo "validating journal store messages..."
+              withAllMsgQueues n ms $ \_ -> processValidateQueue
         where
           processExpireQueue old q =
             runExceptT expireQueue >>= \case

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

@@ -64,6 +64,7 @@ import System.FilePath ((</>))
 import System.IO (BufferMode (..), Handle, IOMode (..), SeekMode (..), stdout)
 import qualified System.IO as IO
 import System.Random (StdGen, genByteString, newStdGen)
+import UnliftIO.Internals.Async (pooledMapConcurrentlyIO)
 
 data JournalMsgStore = JournalMsgStore
   { config :: JournalStoreConfig,
@@ -214,36 +215,43 @@ instance MsgStoreClass JournalMsgStore where
   -- It is used to export storage to a single file and also to expire messages and validate all queues when server is started.
   -- TODO this function requires case-sensitive file system, because it uses queue directory as recipient ID.
   -- It can be made to support case-insensite FS by supporting more than one queue per directory, by getting recipient ID from state file name.
-  withAllMsgQueues :: Monoid a => JournalMsgStore -> (RecipientId -> JournalMsgQueue -> IO a) -> IO a
-  withAllMsgQueues ms@JournalMsgStore {config} action = ifM (doesDirectoryExist storePath) processStore (pure mempty)
+  withAllMsgQueues :: Monoid a => Int -> JournalMsgStore -> (RecipientId -> JournalMsgQueue -> IO a) -> IO a
+  withAllMsgQueues concurrency ms@JournalMsgStore {config} action = ifM (doesDirectoryExist storePath) processStore (pure mempty)
     where
       processStore = do
         closeMsgStore ms
         lock <- createLockIO -- the same lock is used for all queues
-        (!count, !res) <- foldQueues 0 (processQueue lock) (0, mempty) ("", storePath)
-        progress count
+        count <- newTVarIO (0 :: Int)
+        !r <- foldQueues 0 (processQueue lock count) ("", storePath)
+        progress =<< readTVarIO count
         putStrLn ""
-        pure res
+        pure r
       JournalStoreConfig {storePath, pathParts} = config
-      processQueue queueLock (!i :: Int, !r) (queueId, dir) = do
+      processQueue queueLock count (queueId, dir) = do
+        i <- atomically $ stateTVar count $ \i -> (i, i + 1)
         when (i `mod` 100 == 0) $ progress i
         let statePath = msgQueueStatePath dir queueId
         q <- openMsgQueue ms JMQueue {queueDirectory = dir, queueLock, statePath}
-        r' <- case strDecode $ B.pack queueId of
+        !r <- case strDecode $ B.pack queueId of
           Right rId -> action rId q
           Left e -> do
             putStrLn ("Error: message queue directory " <> dir <> " is invalid: " <> e)
             exitFailure
         closeMsgQueue q
-        pure (i + 1, r <> r')
+        pure r
       progress i = do
         putStr $ "Processed: " <> show i <> " queues\r"
         IO.hFlush stdout
-      foldQueues depth f acc (queueId, path) = do
+      mapM' :: (a -> IO b) -> [a] -> IO [b]
+      mapM' 
+        | concurrency == 1 = mapM
+        | otherwise = pooledMapConcurrentlyIO concurrency
+      foldQueues depth f (queueId, path) = do
         let f' = if depth == pathParts - 1 then f else foldQueues (depth + 1) f
-        listDirs >>= foldM f' acc
+        !r <- fmap mconcat . mapM' f' =<< listDirs
+        pure r
         where
-          listDirs = fmap catMaybes . mapM queuePath =<< listDirectory path
+          listDirs = fmap catMaybes . mapM' queuePath =<< listDirectory path
           queuePath dir = do
             let !path' = path </> dir
                 !queueId' = queueId <> dir

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

@@ -57,7 +57,8 @@ instance MsgStoreClass STMMsgStore where
   activeMsgQueues = msgQueues
   {-# INLINE activeMsgQueues #-}
 
-  withAllMsgQueues = withActiveMsgQueues
+  -- no concurrency here
+  withAllMsgQueues _ = withActiveMsgQueues
   {-# INLINE withAllMsgQueues #-}
 
   logQueueStates _ = pure ()

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

@@ -1,3 +1,4 @@
+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE GADTs #-}
@@ -24,7 +25,7 @@ class Monad (StoreMonad s) => MsgStoreClass s where
   newMsgStore :: MsgStoreConfig s -> IO s
   closeMsgStore :: s -> IO ()
   activeMsgQueues :: s -> TMap RecipientId (MsgQueue s)
-  withAllMsgQueues :: Monoid a => s -> (RecipientId -> MsgQueue s -> IO a) -> IO a
+  withAllMsgQueues :: Monoid a => Int -> s -> (RecipientId -> MsgQueue s -> IO a) -> IO a
   logQueueStates :: s -> IO ()
   logQueueState :: MsgQueue s -> IO ()
   getMsgQueue :: s -> RecipientId -> ExceptT ErrorType IO (MsgQueue s)
@@ -47,7 +48,12 @@ data SMSType :: MSType -> Type where
 data AMSType = forall s. AMSType (SMSType s)
 
 withActiveMsgQueues :: (MsgStoreClass s, Monoid a) => s -> (RecipientId -> MsgQueue s -> IO a) -> IO a
-withActiveMsgQueues st f = readTVarIO (activeMsgQueues st) >>= M.foldrWithKey (\k v -> (>>= \r -> r `seq` (r <>) <$> f k v)) (pure mempty)
+withActiveMsgQueues st f = readTVarIO (activeMsgQueues st) >>= M.foldrWithKey run (pure mempty)
+  where
+    run k v acc = do
+      !r <- acc
+      !r' <- f k v
+      pure $! r <> r'
 
 tryPeekMsg :: MsgStoreClass s => MsgQueue s -> ExceptT ErrorType IO (Maybe Message)
 tryPeekMsg mq = isolateQueue mq "tryPeekMsg" $ tryPeekMsg_ mq