smp server: faster export (#1626)

* smp server: faster export

* flush

* export messages with PostgreSQL database

* remove flush

* batch message writes

* Revert "batch message writes"

This reverts commit 61fb3c3.

* remove $!

* fast journal export

* another approach

* Revert "smp server: remove dependency of message size on the version (#1627)"

This reverts commit 8fea152.

* style

* faster?

* cleanup

* cleanup

* refactor

* refactor

* concurrent read messages

* Revert "concurrent read messages"

This reverts commit 05a32e6.

* concurrent read/write

* parameter to export to/import from another message log file

* Revert "parameter to export to/import from another message log file"

This reverts commit 4e88b03.

* Revert "concurrent read/write"

This reverts commit a8eab1f.

Author: epoberezkin

src/Simplex/Messaging/Server.hs

@@ -105,7 +105,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)
+import Simplex.Messaging.Server.MsgStore.Journal (JournalMsgStore, JournalQueue, getJournalQueueMessages)
 import Simplex.Messaging.Server.MsgStore.STM
 import Simplex.Messaging.Server.MsgStore.Types
 import Simplex.Messaging.Server.NtfStore
@@ -2104,27 +2104,36 @@ randomId = fmap EntityId . randomId'
 {-# INLINE randomId #-}
 
 saveServerMessages :: Bool -> MsgStore s -> IO ()
-saveServerMessages drainMsgs = \case
-  StoreMemory ms@STMMsgStore {storeConfig = STMStoreConfig {storePath}} -> case storePath of
+saveServerMessages drainMsgs ms = case ms of
+  StoreMemory STMMsgStore {storeConfig = STMStoreConfig {storePath}} -> case storePath of
     Just f -> exportMessages False ms f drainMsgs
     Nothing -> logNote "undelivered messages are not saved"
   StoreJournal _ -> logNote "closed journal message storage"
 
-exportMessages :: MsgStoreClass s => Bool -> s -> FilePath -> Bool -> IO ()
-exportMessages tty ms f drainMsgs = do
+exportMessages :: forall s. MsgStoreClass s => Bool -> MsgStore s -> FilePath -> Bool -> IO ()
+exportMessages tty st f drainMsgs = do
   logNote $ "saving messages to file " <> T.pack f
-  liftIO $ withFile f WriteMode $ \h ->
-    tryAny (unsafeWithAllMsgQueues tty True ms $ saveQueueMsgs h) >>= \case
-      Right (Sum total) -> logNote $ "messages saved: " <> tshow total
+  run $ case st of
+    StoreMemory ms -> exportMessages_ ms $ getMsgs ms
+    StoreJournal ms -> exportMessages_ ms $ getJournalMsgs ms
+  where
+    exportMessages_ ms get = fmap (\(Sum n) -> n) . unsafeWithAllMsgQueues tty ms . saveQueueMsgs get
+    run :: (Handle -> IO Int) -> IO ()
+    run a = liftIO $ withFile f WriteMode $ tryAny . a >=> \case
+      Right n -> logNote $ "messages saved: " <> tshow n
       Left e -> do
         logError $ "error exporting messages: " <> tshow e
         exitFailure
-  where
-    saveQueueMsgs h q = do
-      msgs <-
-        unsafeRunStore q "saveQueueMsgs" $
-          getQueueMessages_ drainMsgs q =<< getMsgQueue ms q False
-      BLD.hPutBuilder h $ encodeMessages (recipientId q) msgs
+    getJournalMsgs ms q =
+      readTVarIO (msgQueue q) >>= \case
+        Just _ -> getMsgs ms q
+        Nothing -> getJournalQueueMessages ms q
+    getMsgs :: MsgStoreClass s' => s' -> StoreQueue s' -> IO [Message]
+    getMsgs ms q = unsafeRunStore q "saveQueueMsgs" $ getQueueMessages_ drainMsgs q =<< getMsgQueue ms q False
+    saveQueueMsgs :: (StoreQueue s -> IO [Message]) -> Handle -> StoreQueue s -> IO (Sum Int)
+    saveQueueMsgs get h q = do
+      msgs <- get q
+      unless (null msgs) $ 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')
 
@@ -2151,7 +2160,7 @@ processServerMessages StartOptions {skipWarnings} = do
               run processValidateQueue
         | otherwise = logWarn "skipping message expiration" $> Nothing
         where
-          run a = unsafeWithAllMsgQueues False False ms a `catchAny` \_ -> exitFailure
+          run a = unsafeWithAllMsgQueues False ms a `catchAny` \_ -> exitFailure
           processExpireQueue :: Int64 -> JournalQueue s -> IO MessageStats
           processExpireQueue old q = unsafeRunStore q "processExpireQueue" $ do
             mq <- getMsgQueue ms q False

src/Simplex/Messaging/Server/Main.hs

@@ -60,7 +60,7 @@ import Simplex.Messaging.Transport (supportedProxyClientSMPRelayVRange, alpnSupp
 import Simplex.Messaging.Transport.Client (TransportHost (..), defaultSocksProxy)
 import Simplex.Messaging.Transport.HTTP2 (httpALPN)
 import Simplex.Messaging.Transport.Server (ServerCredentials (..), mkTransportServerConfig)
-import Simplex.Messaging.Util (eitherToMaybe, ifM)
+import Simplex.Messaging.Util (eitherToMaybe, ifM, unlessM)
 import System.Directory (createDirectoryIfMissing, doesDirectoryExist, doesFileExist)
 import System.Exit (exitFailure)
 import System.FilePath (combine)
@@ -140,16 +140,25 @@ smpServerCLI_ generateSite serveStaticFiles attachStaticFiles cfgPath logPath =
               confirmOrExit
                 ("WARNING: journal directory " <> storeMsgsJournalDir <> " will be exported to message log file " <> storeMsgsFilePath)
                 "Journal not exported"
-              ms <- newJournalMsgStore logPath MQStoreCfg
-              -- TODO [postgres] in case postgres configured, queues must be read from database
-              readQueueStore True (mkQueue ms False) storeLogFile $ stmQueueStore ms
-              exportMessages True ms storeMsgsFilePath False
-              putStrLn "Export completed"
               case readStoreType ini of
-                Right (ASType SQSMemory SMSMemory) -> putStrLn "store_messages set to `memory`, start the server."
-                Right (ASType SQSMemory SMSJournal) -> putStrLn "store_messages set to `journal`, update it to `memory` in INI file"
-                Right (ASType SQSPostgres SMSJournal) ->
+                Right (ASType SQSMemory msType) -> do
+                  ms <- newJournalMsgStore logPath MQStoreCfg
+                  readQueueStore True (mkQueue ms False) storeLogFile $ stmQueueStore ms
+                  exportMessages True (StoreJournal ms) storeMsgsFilePath False
+                  putStrLn "Export completed"
+                  putStrLn $ case msType of
+                    SMSMemory -> "store_messages set to `memory`, start the server."
+                    SMSJournal -> "store_messages set to `journal`, update it to `memory` in INI file"
+                Right (ASType SQSPostgres SMSJournal) -> do
 #if defined(dbServerPostgres)
+                  let dbStoreLogPath = enableDbStoreLog' ini $> storeLogFilePath
+                      dbOpts@DBOpts {connstr, schema} = iniDBOptions ini defaultDBOpts
+                  unlessM (checkSchemaExists connstr schema) $ do
+                    putStrLn $ "Schema " <> B.unpack schema <> " does not exist in PostrgreSQL database: " <> B.unpack connstr
+                    exitFailure
+                  ms <- newJournalMsgStore logPath $ PQStoreCfg PostgresStoreCfg {dbOpts, dbStoreLogPath, confirmMigrations = MCYesUp, deletedTTL = iniDeletedTTL ini}
+                  exportMessages True (StoreJournal ms) storeMsgsFilePath False
+                  putStrLn "Export completed"
                   putStrLn "store_messages set to `journal`, store_queues is set to `database`.\nExport queues to store log to use memory storage for messages (`smp-server database export`)."
 #else
                   noPostgresExit
@@ -582,7 +591,7 @@ exportDatabaseToStoreLog logPath dbOpts storeLogFilePath = do
   ps <- newJournalMsgStore logPath $ PQStoreCfg storeCfg
   sl <- openWriteStoreLog False storeLogFilePath
   Sum sCnt <- foldServiceRecs (postgresQueueStore ps) $ \sr -> logNewService sl sr $> Sum (1 :: Int)
-  Sum qCnt <- foldQueueRecs True True (postgresQueueStore ps) Nothing $ \(rId, qr) -> logCreateQueue sl rId qr $> Sum (1 :: Int)
+  Sum qCnt <- foldQueueRecs True True (postgresQueueStore ps) $ \(rId, qr) -> logCreateQueue sl rId qr $> Sum (1 :: Int)
   closeStoreLog sl
   pure (sCnt, qCnt)
 #endif

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

@@ -38,6 +38,7 @@ module Simplex.Messaging.Server.MsgStore.Journal
     msgQueueStatePath,
     readQueueState,
     newMsgQueueState,
+    getJournalQueueMessages,
     newJournalId,
     appendState,
     queueLogFileName,
@@ -58,7 +59,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.Either (fromRight, partitionEithers)
 import Data.Functor (($>))
 import Data.Int (Int64)
 import Data.List (sort)
@@ -405,11 +406,11 @@ instance MsgStoreClass (JournalMsgStore s) where
 
   -- 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 -> Bool -> JournalMsgStore s -> (JournalQueue s -> IO a) -> IO a
-  unsafeWithAllMsgQueues tty withData ms action = case queueStore_ ms of
+  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
 #if defined(dbServerPostgres)
-    PQStore st -> foldQueueRecs tty withData st Nothing $ uncurry (mkQueue ms False) >=> run
+    PQStore st -> foldQueueRecs False tty st $ uncurry (mkQueue ms False) >=> run
 #endif
     where
       run q = do
@@ -429,7 +430,7 @@ instance MsgStoreClass (JournalMsgStore s) where
 #if defined(dbServerPostgres)
     PQStore st -> do
       let JournalMsgStore {queueLocks, sharedLock} = ms
-      foldQueueRecs tty False st (Just veryOld) $ \(rId, qr) -> do
+      foldRecentQueueRecs veryOld tty st $ \(rId, qr) -> do
         q <- mkQueue ms False rId qr
         withSharedWaitLock rId queueLocks sharedLock $ run $ tryStore' "deleteExpiredMsgs" rId $
           getLoadedQueue q >>= unStoreIO . expireQueueMsgs ms now old
@@ -485,7 +486,7 @@ instance MsgStoreClass (JournalMsgStore s) where
     where
       newQ = do
         let dir = msgQueueDirectory ms rId
-            statePath = msgQueueStatePath dir $ B.unpack (strEncode rId)
+            statePath = msgQueueStatePath dir rId
             queue = JMQueue {queueDirectory = dir, statePath}
         q <- ifM (doesDirectoryExist dir) (openMsgQueue ms queue forWrite) (createQ queue)
         atomically $ writeTVar msgQueue' $ Just q
@@ -563,8 +564,9 @@ instance MsgStoreClass (JournalMsgStore s) where
     where
       getSize = maybe (pure (-1)) (fmap size . readTVarIO . state)
 
+   -- drainMsgs is never True with Journal storage
   getQueueMessages_ :: Bool -> JournalQueue s -> JournalMsgQueue s -> StoreIO s [Message]
-  getQueueMessages_ drainMsgs q' q = StoreIO (run [])
+  getQueueMessages_ drainMsgs q' q = StoreIO $ if drainMsgs then run [] else readTVarIO (state q) >>= runFast
     where
       run msgs = readTVarIO (handles q) >>= maybe (pure []) (getMsg msgs)
       getMsg msgs hs = chooseReadJournal q' q drainMsgs hs >>= maybe (pure msgs) readMsg
@@ -573,6 +575,16 @@ instance MsgStoreClass (JournalMsgStore s) where
             (msg, len) <- hGetMsgAt h $ bytePos rs
             updateReadPos q' q drainMsgs len hs
             (msg :) <$> run msgs
+      runFast MsgQueueState {writeState = ws, readState = rs, size}
+        | size > 0 =
+            readTVarIO (handles q) >>= \case
+              Just (MsgQueueHandles _ rh wh_) -> do
+                msgs <- getJournalRange rh (bytePos rs) (byteCount rs)
+                case wh_ of
+                  Just wh -> (msgs ++) <$> getJournalRange wh 0 (bytePos ws)
+                  Nothing -> pure msgs
+              Nothing -> pure []
+        | otherwise = pure []
 
   writeMsg :: JournalMsgStore s -> JournalQueue s -> Bool -> Message -> ExceptT ErrorType IO (Maybe (Message, Bool))
   writeMsg ms q' logState msg = isolateQueue q' "writeMsg" $ do
@@ -795,8 +807,8 @@ msgQueueDirectory JournalMsgStore {config = JournalStoreConfig {storePath, pathP
       let (seg, s') = B.splitAt 2 s
        in seg : splitSegments (n - 1) s'
 
-msgQueueStatePath :: FilePath -> String -> FilePath
-msgQueueStatePath dir queueId = dir </> (queueLogFileName <> "." <> queueId <> logFileExt)
+msgQueueStatePath :: FilePath -> RecipientId -> FilePath
+msgQueueStatePath dir rId = dir </> (queueLogFileName <> "." <> B.unpack (strEncode rId) <> logFileExt)
 
 createNewJournal :: FilePath -> ByteString -> IO Handle
 createNewJournal dir journalId = do
@@ -1019,3 +1031,33 @@ hClose h =
 
 closeOnException :: Handle -> IO a -> IO a
 closeOnException h a = a `E.onException` hClose h
+
+getJournalQueueMessages :: JournalMsgStore s -> JournalQueue s -> IO [Message]
+getJournalQueueMessages ms q =
+  readQueueState ms (msgQueueStatePath dir rId) >>= \case
+    (Just MsgQueueState {readState = rs, writeState = ws, size}, _) | size > 0 -> do
+      msgs <- getMsgs (journalId rs) (bytePos rs) (byteCount rs)
+      if journalId rs == journalId ws
+        then pure msgs
+        else (msgs ++) <$> getMsgs (journalId ws) 0 (bytePos ws)
+    _ -> pure []
+  where
+    rId = recipientId' q
+    dir = msgQueueDirectory ms rId
+    getMsgs jId from to =
+      IO.withFile (journalFilePath dir jId) ReadWriteMode $ \h' ->
+        getJournalRange h' from to
+
+getJournalRange :: Handle -> Int64 -> Int64 -> IO [Message]
+getJournalRange h from to
+  | to > from = do
+      IO.hSeek h AbsoluteSeek $ fromIntegral from
+      parseMsgs =<< B.hGet h (fromIntegral $ to - from)
+  | otherwise = pure []
+  where
+    parseMsgs s = do
+      let (errs, msgs) = partitionEithers $ map strDecode $ B.lines s
+      unless (null errs) $ do
+        f <- IO.hShow h
+        putStrLn $ "Error reading " <> show (length errs) <> " messages from " <> f
+      pure msgs

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

@@ -82,7 +82,7 @@ instance MsgStoreClass STMMsgStore where
   {-# INLINE closeMsgStore #-}
   withActiveMsgQueues = withLoadedQueues . queueStore_
   {-# INLINE withActiveMsgQueues #-}
-  unsafeWithAllMsgQueues _ _ = withLoadedQueues . queueStore_
+  unsafeWithAllMsgQueues _ = withLoadedQueues . queueStore_
   {-# INLINE unsafeWithAllMsgQueues #-}
 
   expireOldMessages :: Bool -> STMMsgStore -> Int64 -> Int64 -> IO MessageStats

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

@@ -40,8 +40,8 @@ class (Monad (StoreMonad s), QueueStoreClass (StoreQueue s) (QueueStore s)) => M
   closeMsgStore :: s -> IO ()
   withActiveMsgQueues :: Monoid a => s -> (StoreQueue s -> IO a) -> IO a
   -- This function can only be used in server CLI commands or before server is started.
-  -- tty, withData, store
-  unsafeWithAllMsgQueues :: Monoid a => Bool -> Bool -> s -> (StoreQueue s -> IO a) -> IO a
+  -- tty, store
+  unsafeWithAllMsgQueues :: Monoid a => Bool -> s -> (StoreQueue s -> IO a) -> IO a
   -- tty, store, now, ttl
   expireOldMessages :: Bool -> s -> Int64 -> Int64 -> IO MessageStats
   logQueueStates :: s -> IO ()
@@ -84,6 +84,7 @@ data MessageStats = MessageStats
     expiredMsgsCount :: Int,
     storedQueues :: Int
   }
+  deriving (Show)
 
 instance Monoid MessageStats where
   mempty = MessageStats 0 0 0

src/Simplex/Messaging/Server/QueueStore/Postgres.hs

@@ -25,6 +25,7 @@ module Simplex.Messaging.Server.QueueStore.Postgres
     batchInsertQueues,
     foldServiceRecs,
     foldQueueRecs,
+    foldRecentQueueRecs,
     handleDuplicate,
     withLog_,
     withDB',
@@ -545,8 +546,28 @@ foldServiceRecs st f =
     DB.fold_ db "SELECT service_id, service_role, service_cert, service_cert_hash, created_at FROM services" mempty $
       \ !acc -> fmap (acc <>) . f . rowToServiceRec
 
-foldQueueRecs :: forall a q. Monoid a => Bool -> Bool -> PostgresQueueStore q -> Maybe Int64 -> ((RecipientId, QueueRec) -> IO a) -> IO a
-foldQueueRecs tty withData st skipOld_ f = do
+foldQueueRecs :: Monoid a => Bool -> Bool -> PostgresQueueStore q -> ((RecipientId, QueueRec) -> IO a) -> IO a
+foldQueueRecs withData = foldQueueRecs_ foldRecs
+  where
+    foldRecs db acc f'
+      | withData = DB.fold_ db (queueRecQueryWithData <> cond) acc $ \acc' -> f' acc' . rowToQueueRecWithData
+      | otherwise = DB.fold_ db (queueRecQuery <> cond) acc $ \acc' -> f' acc' . rowToQueueRec
+    cond = " WHERE deleted_at IS NULL ORDER BY recipient_id ASC"
+
+foldRecentQueueRecs :: Monoid a => Int64 -> Bool -> PostgresQueueStore q -> ((RecipientId, QueueRec) -> IO a) -> IO a
+foldRecentQueueRecs old = foldQueueRecs_ foldRecs
+  where
+    foldRecs db acc f' = DB.fold db (queueRecQuery <> cond) (Only old) acc $ \acc' -> f' acc' . rowToQueueRec
+    cond = " WHERE deleted_at IS NULL AND updated_at > ? ORDER BY recipient_id ASC"
+
+foldQueueRecs_ ::
+  Monoid a =>
+  (DB.Connection -> (Int, a) -> ((Int, a) -> (RecipientId, QueueRec) -> IO (Int, a)) -> IO (Int, a)) ->
+  Bool ->
+  PostgresQueueStore q ->
+  ((RecipientId, QueueRec) -> IO a) ->
+  IO a
+foldQueueRecs_ foldRecs tty st f = do
   (n, r) <- withTransaction (dbStore st) $ \db ->
     foldRecs db (0 :: Int, mempty) $ \(i, acc) qr -> do
       r <- f qr
@@ -557,13 +578,6 @@ foldQueueRecs tty withData st skipOld_ f = do
   when tty $ putStrLn $ progress n
   pure r
   where
-    foldRecs db acc f' = case skipOld_ of
-      Nothing
-        | withData -> DB.fold_ db (queueRecQueryWithData <> " WHERE deleted_at IS NULL") acc $ \acc' -> f' acc' . rowToQueueRecWithData
-        | otherwise -> DB.fold_ db (queueRecQuery <> " WHERE deleted_at IS NULL") acc $ \acc' -> f' acc' . rowToQueueRec
-      Just old
-        | withData -> DB.fold db (queueRecQueryWithData <> " WHERE deleted_at IS NULL AND updated_at > ?") (Only old) acc $ \acc' -> f' acc' . rowToQueueRecWithData
-        | otherwise -> DB.fold db (queueRecQuery <> " WHERE deleted_at IS NULL AND updated_at > ?") (Only old) acc $ \acc' -> f' acc' . rowToQueueRec
     progress i = "Processed: " <> show i <> " records"
 
 queueRecQuery :: Query