From c1aee1f22c9174b559a66d2c9568cbb8772b61e0 Mon Sep 17 00:00:00 2001 From: Evgeny Poberezkin Date: Wed, 23 Oct 2024 21:10:40 +0100 Subject: [PATCH] concurrently 2 --- src/Simplex/Messaging/Server.hs | 2 +- .../Messaging/Server/MsgStore/Journal.hs | 59 +++++++++++-------- .../Messaging/Server/MsgStore/Types.hs | 10 ++-- src/Simplex/Messaging/Util.hs | 6 -- 4 files changed, 39 insertions(+), 38 deletions(-) diff --git a/src/Simplex/Messaging/Server.hs b/src/Simplex/Messaging/Server.hs index 1ceb77069..93503b7db 100644 --- a/src/Simplex/Messaging/Server.hs +++ b/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..." diff --git a/src/Simplex/Messaging/Server/MsgStore/Journal.hs b/src/Simplex/Messaging/Server/MsgStore/Journal.hs index 17100a812..b50aef0b3 100644 --- a/src/Simplex/Messaging/Server/MsgStore/Journal.hs +++ b/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 diff --git a/src/Simplex/Messaging/Server/MsgStore/Types.hs b/src/Simplex/Messaging/Server/MsgStore/Types.hs index 094151aae..12be7042c 100644 --- a/src/Simplex/Messaging/Server/MsgStore/Types.hs +++ b/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 diff --git a/src/Simplex/Messaging/Util.hs b/src/Simplex/Messaging/Util.hs index ea0322b74..55e65bfca 100644 --- a/src/Simplex/Messaging/Util.hs +++ b/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_ #-}