concurrently 2

src/Simplex/Messaging/Server.hs

@@ -1765,7 +1765,7 @@ processServerMessages = do
           Just f -> ifM (doesFileExist f) (importMessages ms f old_) (pure newMessageStats)
           Nothing -> pure newMessageStats
         AMS SMSJournal ms -> do
-          n <- getNumCapabilities
+          n <- (8 *) <$> getNumCapabilities
           case old_ of
             Just old -> do
               logInfo "expiring journal store messages..."

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

@@ -57,14 +57,14 @@ import Simplex.Messaging.Protocol (ErrorType (..), Message (..), RecipientId)
 import Simplex.Messaging.Server.MsgStore.Types
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
-import Simplex.Messaging.Util (ifM, tshow, ($>>=))
+import Simplex.Messaging.Util (ifM, tshow, unlessM, ($>>=))
 import System.Directory
 import System.Exit
 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)
+import UnliftIO.Internals.Async (concurrently_, pooledReplicateConcurrentlyN_)
 
 data JournalMsgStore = JournalMsgStore
   { config :: JournalStoreConfig,
@@ -216,42 +216,49 @@ instance MsgStoreClass JournalMsgStore where
   -- 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 => Int -> JournalMsgStore -> (RecipientId -> JournalMsgQueue -> IO a) -> IO a
-  withAllMsgQueues concurrency ms@JournalMsgStore {config} action = ifM (doesDirectoryExist storePath) processStore (pure mempty)
+  withAllMsgQueues n 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
+        queues <- newTBQueueIO 1000
         count <- newTVarIO (0 :: Int)
-        !r <- foldQueues 0 (processQueue lock count) ("", storePath)
+        res <- newTVarIO mempty
+        completed <- newTVarIO False
+        concurrently_
+          (listQueues queues 0 ("", storePath) >> atomically (writeTBQueue queues Nothing))
+          (pooledReplicateConcurrentlyN_ n n $ processQueue' lock queues count res completed)
         progress =<< readTVarIO count
         putStrLn ""
-        pure r
+        readTVarIO res
       JournalStoreConfig {storePath, pathParts} = config
-      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
-          Right rId -> action rId q
-          Left e -> do
-            putStrLn ("Error: message queue directory " <> dir <> " is invalid: " <> e)
-            exitFailure
-        closeMsgQueue q
-        pure r
+      processQueue' queueLock queues count res completed = loop
+        where
+          loop = 
+            unlessM (readTVarIO completed) $
+              atomically (readTBQueue queues `orElse` ifM (readTVar completed) (pure Nothing) retry) >>= \case
+                Nothing -> unlessM (readTVarIO completed) $ atomically $ writeTVar completed True
+                Just (queueId, dir) -> case strDecode $ B.pack queueId of
+                  Right rId -> 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}
+                    action rId q >>= atomically . modifyTVar' res . (<>)
+                    closeMsgQueue q
+                    loop
+                  Left e -> do
+                    putStrLn ("Error: message queue directory " <> dir <> " is invalid: " <> e)
+                    exitFailure
       progress i = do
         putStr $ "Processed: " <> show i <> " queues\r"
         IO.hFlush stdout
-      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
-        !r <- fmap mconcat . mapM' f' =<< listDirs
-        pure r
+      listQueues queues depth (queueId, path) = listDirs >>= mapM_ go
         where
-          listDirs = fmap catMaybes . mapM' queuePath =<< listDirectory path
+          go
+            | depth == pathParts - 1 = atomically . writeTBQueue queues . Just
+            | otherwise = listQueues queues (depth + 1)
+          listDirs = fmap catMaybes . mapM queuePath =<< listDirectory path
           queuePath dir = do
             let !path' = path </> dir
                 !queueId' = queueId <> dir

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

@@ -10,6 +10,7 @@
 module Simplex.Messaging.Server.MsgStore.Types where
 
 import Control.Concurrent.STM
+import Control.Monad (foldM)
 import Control.Monad.Trans.Except
 import Data.Int (Int64)
 import Data.Kind
@@ -48,12 +49,11 @@ 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 run (pure mempty)
+withActiveMsgQueues st f = readTVarIO (activeMsgQueues st) >>= foldM run mempty . M.assocs
   where
-    run k v acc = do
-      !r <- acc
-      !r' <- f k v
-      pure $! r <> r'
+    run !acc (k, v) = do
+      r <- f k v
+      pure $ acc <> r
 
 tryPeekMsg :: MsgStoreClass s => MsgQueue s -> ExceptT ErrorType IO (Maybe Message)
 tryPeekMsg mq = isolateQueue mq "tryPeekMsg" $ tryPeekMsg_ mq

src/Simplex/Messaging/Util.hs

@@ -19,8 +19,6 @@ import Data.Int (Int64)
 import Data.List (groupBy, sortOn)
 import Data.List.NonEmpty (NonEmpty)
 import qualified Data.List.NonEmpty as L
-import Data.Map.Strict (Map)
-import qualified Data.Map.Strict as M
 import Data.Text (Text)
 import qualified Data.Text as T
 import Data.Text.Encoding (decodeUtf8With, encodeUtf8)
@@ -187,7 +185,3 @@ encodeJSON = safeDecodeUtf8 . LB.toStrict . J.encode
 
 decodeJSON :: FromJSON a => Text -> Maybe a
 decodeJSON = J.decode . LB.fromStrict . encodeUtf8
-
-traverseWithKey_ :: Monad m => (k -> v -> m ()) -> Map k v -> m ()
-traverseWithKey_ f = M.foldrWithKey (\k v -> (f k v >>)) (pure ())
-{-# INLINE traverseWithKey_ #-}