fast journal export

src/Simplex/Messaging/Server/Main.hs

@@ -52,7 +52,7 @@ import Simplex.Messaging.Server.Env.STM
 import Simplex.Messaging.Server.Expiration
 import Simplex.Messaging.Server.Information
 import Simplex.Messaging.Server.Main.Init
-import Simplex.Messaging.Server.MsgStore.Journal (JournalMsgStore (..), QStoreCfg (..), stmQueueStore)
+import Simplex.Messaging.Server.MsgStore.Journal (JournalMsgStore (..), QStoreCfg (..), exportJournalMessages, stmQueueStore)
 import Simplex.Messaging.Server.MsgStore.Types (MsgStoreClass (..), SQSType (..), SMSType (..), newMsgStore)
 import Simplex.Messaging.Server.QueueStore.Postgres.Config
 import Simplex.Messaging.Server.StoreLog.ReadWrite (readQueueStore)
@@ -60,11 +60,11 @@ 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, unlessM)
+import Simplex.Messaging.Util (eitherToMaybe, ifM, tshow, unlessM)
 import System.Directory (createDirectoryIfMissing, doesDirectoryExist, doesFileExist)
 import System.Exit (exitFailure)
 import System.FilePath (combine)
-import System.IO (BufferMode (..), hSetBuffering, stderr, stdout)
+import System.IO (BufferMode (..), IOMode (..), hSetBuffering, stderr, stdout, withFile)
 import Text.Read (readMaybe)
 
 #if defined(dbServerPostgres)
@@ -131,24 +131,26 @@ smpServerCLI_ generateSite serveStaticFiles attachStaticFiles cfgPath logPath =
                 Right (ASType SQSMemory SMSMemory) -> "store_messages set to `memory`, update it to `journal` in INI file"
                 Right (ASType _ SMSJournal) -> "store_messages set to `journal`"
                 Left e -> e <> ", configure storage correctly"
-        SCExport
+        SCExport fast
           | msgsFileExists && msgsDirExists -> exitConfigureMsgStorage
           | msgsFileExists -> do
               putStrLn $ storeMsgsFilePath <> " file already exists."
               exitFailure
+          | fast -> do
+              confirmExport
+              logNote $ "saving messages to file " <> T.pack storeMsgsFilePath
+              ms <- newJournalMsgStore logPath MQStoreCfg
+              total <- withFile storeMsgsFilePath WriteMode $ exportJournalMessages True ms
+              logNote $ "messages saved: " <> tshow total
+              completedExport
           | otherwise -> do
-              confirmOrExit
-                ("WARNING: journal directory " <> storeMsgsJournalDir <> " will be exported to message log file " <> storeMsgsFilePath)
-                "Journal not exported"
+              confirmExport
+              let msType = readStoreType ini
               case readStoreType ini of
-                Right (ASType SQSMemory msType) -> do
+                Right (ASType SQSMemory _) -> do
                   ms <- newJournalMsgStore logPath MQStoreCfg
                   readQueueStore True (mkQueue ms False) storeLogFile $ stmQueueStore ms
                   exportMessages True 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
@@ -158,12 +160,23 @@ smpServerCLI_ generateSite serveStaticFiles attachStaticFiles cfgPath logPath =
                     exitFailure
                   ms <- newJournalMsgStore logPath $ PQStoreCfg PostgresStoreCfg {dbOpts, dbStoreLogPath, confirmMigrations = MCYesUp, deletedTTL = iniDeletedTTL ini}
                   exportMessages True 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
 #endif
-                Left e -> putStrLn $ e <> ", configure storage correctly"
+                Left _ -> pure ()
+              completedExport
+          where
+            confirmExport =
+              confirmOrExit
+                ("WARNING: journal directory " <> storeMsgsJournalDir <> " will be exported to message log file " <> storeMsgsFilePath)
+                "Journal not exported"
+            completedExport = do
+              putStrLn "Export completed"
+              putStrLn $ case readStoreType ini of
+                Right (ASType SQSMemory SMSMemory) -> "store_messages set to `memory`, start the server."
+                Right (ASType SQSMemory SMSJournal) -> "store_messages set to `journal`, update it to `memory` in INI file"
+                Right (ASType SQSPostgres SMSJournal) -> "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`)."
+                Left e -> e <> ", configure storage correctly"
         SCDelete
           | not msgsDirExists -> do
               putStrLn $ storeMsgsJournalDir <> " directory does not exists."
@@ -202,7 +215,7 @@ smpServerCLI_ generateSite serveStaticFiles attachStaticFiles cfgPath logPath =
           where
             setToDbStr :: String
             setToDbStr = "store_queues set to `memory`, update it to `database` in INI file"
-        SCExport
+        SCExport _
           | schemaExists && storeLogExists -> exitConfigureQueueStore connstr schema
           | not schemaExists -> do
               putStrLn $ "Schema " <> B.unpack schema <> " does not exist in PostrgreSQL database: " <> B.unpack connstr
@@ -679,7 +692,7 @@ data CliCommand
   | Journal StoreCmd
   | Database StoreCmd DBOpts
 
-data StoreCmd = SCImport | SCExport | SCDelete
+data StoreCmd = SCImport | SCExport Bool | SCDelete
 
 cliCommandP :: FilePath -> FilePath -> FilePath -> Parser CliCommand
 cliCommandP cfgPath logPath iniFile =
@@ -840,7 +853,8 @@ cliCommandP cfgPath logPath iniFile =
     storeCmdP src dest =
       hsubparser
         ( command "import" (info (pure SCImport) (progDesc $ "Import " <> src <> " into a new " <> dest))
-            <> command "export" (info (pure SCExport) (progDesc $ "Export " <> dest <> " to " <> src))
+            <> command "export" (info (pure $ SCExport False) (progDesc $ "Export " <> dest <> " to " <> src))
+            <> command "fast-export" (info (pure $ SCExport True) (progDesc $ "Fast export of " <> dest <> " to " <> src))
             <> command "delete" (info (pure SCDelete) (progDesc $ "Delete " <> dest))
         )
     parseBasicAuth :: ReadM ServerPassword

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

@@ -47,6 +47,8 @@ module Simplex.Messaging.Server.MsgStore.Journal
 #if defined(dbServerPostgres)
     postgresQueueStore,
 #endif
+    exportJournalMessages,
+    encodeMessages,
   )
 where
 
@@ -56,14 +58,16 @@ import Control.Logger.Simple
 import Control.Monad
 import Control.Monad.Trans.Except
 import qualified Data.Attoparsec.ByteString.Char8 as A
+import qualified Data.ByteString.Builder as BLD
 import Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Char8 as B
-import Data.Either (fromRight)
+import Data.Char (ord)
+import Data.Either (fromRight, partitionEithers)
 import Data.Functor (($>))
 import Data.Int (Int64)
-import Data.List (sort)
+import Data.List (sort, sortBy)
 import qualified Data.Map.Strict as M
-import Data.Maybe (fromMaybe, isJust, isNothing, mapMaybe)
+import Data.Maybe (catMaybes, fromMaybe, isJust, isNothing, mapMaybe)
 import Data.Text (Text)
 import qualified Data.Text as T
 import Data.Text.Encoding (decodeLatin1)
@@ -75,6 +79,7 @@ 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
 import Simplex.Messaging.Server.MsgStore.Journal.SharedLock
 import Simplex.Messaging.Server.MsgStore.Types
 import Simplex.Messaging.Server.QueueStore
@@ -87,8 +92,9 @@ import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
 import Simplex.Messaging.Util (ifM, tshow, whenM, ($>>=), (<$$>))
 import System.Directory
+import System.Exit (exitFailure)
 import System.FilePath (takeFileName, (</>))
-import System.IO (BufferMode (..), Handle, IOMode (..), SeekMode (..))
+import System.IO (BufferMode (..), Handle, IOMode (..), SeekMode (..), stdout)
 import qualified System.IO as IO
 import System.Random (StdGen, genByteString, newStdGen)
 
@@ -1019,3 +1025,83 @@ hClose h =
 
 closeOnException :: Handle -> IO a -> IO a
 closeOnException h a = a `E.onException` hClose h
+
+exportJournalMessages :: Bool -> JournalMsgStore s -> Handle -> IO Int
+exportJournalMessages tty ms@JournalMsgStore {config} h = ifM (doesDirectoryExist storePath) exportStore (pure 0)
+  where
+    exportStore = do
+      (!qCount, !msgCount) <- foldQueues 0 exportQueueMessages (0, 0) ("", storePath)
+      putStrLn $ progress qCount
+      pure msgCount
+    JournalStoreConfig {storePath, pathParts} = config
+    exportQueueMessages :: (Int, Int) -> (String, FilePath) -> IO (Int, Int)
+    exportQueueMessages (!i, !count) (queueId, dir) = do
+      let i' = i + 1
+      when (tty && i' `mod` 100000 == 0) $ putStr (progress i' <> "\r") >> IO.hFlush stdout
+      case strDecode $ B.pack queueId of
+        Right rId -> do
+          let statePath = msgQueueStatePath dir queueId
+          msgs <-
+            readQueueState ms statePath >>= \case
+              (Just MsgQueueState {readState = rs, writeState = ws, size}, _)
+                | size == 0 -> pure []
+                | journalId rs == journalId ws -> do
+                    -- just one journal
+                    let f = journalFilePath dir $ journalId rs
+                    s <- B.readFile f
+                    let (errs, msgs) = parseMsgs $ B.take (bytePos' ws - bytePos' rs) $ B.drop (bytePos' rs) s
+                    unless (null errs) $ do
+                      when tty $ putStrLn $ progress i'
+                      logErr errs f
+                    pure msgs
+                | otherwise -> do
+                    let rf = journalFilePath dir $ journalId rs
+                        wf = journalFilePath dir $ journalId ws
+                    r <- B.readFile rf
+                    w <- B.readFile wf
+                    let (rErrs, rMsgs) = parseMsgs $ B.take (fromIntegral $ byteCount rs) $ B.drop (bytePos' rs) r
+                        (wErrs, wMsgs) = parseMsgs $ B.take (bytePos' ws) w
+                    unless (null rErrs && null wErrs) $ do
+                      when tty $ putStrLn $ progress i'
+                      unless (null rErrs) $ logErr rErrs rf
+                      unless (null wErrs) $ logErr wErrs wf
+                    pure $ rMsgs ++ wMsgs
+                where
+                  bytePos' = fromIntegral . bytePos
+                  parseMsgs = partitionEithers . map strDecode . B.lines
+                  logErr errs f = putStrLn $ "Error reading " <> show (length errs) <> " messages from " <> f
+              _ -> pure []
+          unless (null msgs) $ BLD.hPutBuilder h $ encodeMessages rId msgs
+          pure (i', count + length msgs)
+        Left e -> do
+          logError $ "STORE: exportQueueMessages, message queue directory " <> T.pack dir <> " is invalid, " <> tshow e
+          exitFailure
+    progress i = "Processed: " <> show i <> " queues"
+    foldQueues depth f acc (queueId, path) = do
+      let f' = if depth == pathParts - 1 then f else foldQueues (depth + 1) f
+      listDirs >>= foldM f' acc
+      where
+        listDirs = fmap catMaybes . mapM queuePath . sortBy compareB64 =<< listDirectory path
+        compareB64 [] [] = EQ
+        compareB64 [] _ = LT
+        compareB64 _ [] = GT
+        compareB64 (c : cs) (c' : cs')
+          | c == c' = compareB64 cs cs'
+          | otherwise = compare (charValue c) (charValue c')
+        charValue '-' = 62
+        charValue '_' = 63
+        charValue c
+          | c >= 'A' && c <= 'Z' = ord c - ord 'A'
+          | c >= 'a' && c <= 'z' = 26 + ord c - ord 'a'
+          | c >= '0' && c <= '9' = 52 + ord c - ord '0'
+          | otherwise = -1
+        queuePath dir = do
+          let !path' = path </> dir
+              !queueId' = queueId <> dir
+          ifM
+            (doesDirectoryExist path')
+            (pure $ Just (queueId', path'))
+            (Nothing <$ putStrLn ("Error: path " <> path' <> " is not a directory, skipping"))
+
+encodeMessages :: RecipientId -> [Message] -> BLD.Builder
+encodeMessages rId = mconcat . map (\msg -> BLD.byteString (strEncode $ MLRv3 rId msg) <> BLD.char8 '\n')

tests/CoreTests/MsgStoreTests.hs

@@ -227,6 +227,7 @@ testExportImportStore ms = do
   length <$> listDirectory (msgQueueDirectory ms rId1) `shouldReturn` 2
   length <$> listDirectory (msgQueueDirectory ms rId2) `shouldReturn` 3
   exportMessages False ms testStoreMsgsFile False
+  testFastExport ms testStoreMsgsFile
   closeMsgStore ms
   closeStoreLog sl
   let cfg = (testJournalStoreCfg MQStoreCfg :: JournalStoreConfig 'QSMemory) {storePath = testStoreMsgsDir2}
@@ -238,6 +239,7 @@ testExportImportStore ms = do
   length <$> listDirectory (msgQueueDirectory ms rId1) `shouldReturn` 2
   length <$> listDirectory (msgQueueDirectory ms rId2) `shouldReturn` 3 -- 2 message files
   exportMessages False ms' testStoreMsgsFile2 False
+  testFastExport ms' testStoreMsgsFile2
   (B.readFile testStoreMsgsFile2 `shouldReturn`) =<< B.readFile (testStoreMsgsFile <> ".bak")
   stmStore <- newMsgStore testSMTStoreConfig
   readWriteQueueStore True (mkQueue stmStore True) testStoreLogFile (queueStore stmStore) >>= closeStoreLog
@@ -245,6 +247,11 @@ testExportImportStore ms = do
     importMessages False stmStore testStoreMsgsFile2 Nothing False
   exportMessages False stmStore testStoreMsgsFile False
   (B.sort <$> B.readFile testStoreMsgsFile `shouldReturn`) =<< (B.sort <$> B.readFile (testStoreMsgsFile2 <> ".bak"))
+  where
+    testFastExport ms' f = do
+      void $ withFile (f <> ".fast") WriteMode $ exportJournalMessages False ms'
+      s <- B.readFile f
+      B.readFile (f <> ".fast") `shouldReturn` s
 
 testQueueState :: JournalMsgStore s -> IO ()
 testQueueState ms = do