mirror of
https://github.com/simplex-chat/simplexmq.git
synced 2026-09-25 15:34:37 +00:00
smp server: store messages in PostgreSQL (#1622)
* smp server: store messages in PostgreSQL * stored procedures to write and to expire messages * function to export messages * move all message functions to PostgreSQL, remove delete trigger * comments * import messages to db * fix message import, add export * fix export * fix export * fix compilation flags * import messages line by line * fix server start with database storage * fix compilation * comments
This commit is contained in:
Evgeny
@@ -0,0 +1,325 @@
|
||||
{-# LANGUAGE BangPatterns #-}
|
||||
{-# LANGUAGE DerivingStrategies #-}
|
||||
{-# LANGUAGE DuplicateRecordFields #-}
|
||||
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
|
||||
{-# LANGUAGE InstanceSigs #-}
|
||||
{-# LANGUAGE LambdaCase #-}
|
||||
{-# LANGUAGE MultiWayIf #-}
|
||||
{-# LANGUAGE NamedFieldPuns #-}
|
||||
{-# LANGUAGE OverloadedStrings #-}
|
||||
{-# LANGUAGE QuasiQuotes #-}
|
||||
{-# LANGUAGE ScopedTypeVariables #-}
|
||||
{-# LANGUAGE TupleSections #-}
|
||||
{-# LANGUAGE TypeApplications #-}
|
||||
{-# LANGUAGE TypeFamilies #-}
|
||||
|
||||
module Simplex.Messaging.Server.MsgStore.Postgres
|
||||
( PostgresMsgStore,
|
||||
PostgresMsgStoreCfg (..),
|
||||
PostgresQueue,
|
||||
exportDbMessages,
|
||||
getDbMessageStats,
|
||||
batchInsertMessages,
|
||||
)
|
||||
where
|
||||
|
||||
import Control.Concurrent.STM
|
||||
import Control.Monad
|
||||
import Control.Monad.Reader
|
||||
import Control.Monad.Trans.Except
|
||||
import qualified Data.ByteString as B
|
||||
import qualified Data.ByteString.Builder as BB
|
||||
import qualified Data.ByteString.Lazy as LB
|
||||
import Data.Functor (($>))
|
||||
import Data.IORef
|
||||
import Data.Int (Int64)
|
||||
import Data.List (intersperse)
|
||||
import qualified Data.Map.Strict as M
|
||||
import Data.Text (Text)
|
||||
import Data.Time.Clock.System (SystemTime (..))
|
||||
import Database.PostgreSQL.Simple (Binary (..), Only (..), (:.) (..))
|
||||
import qualified Database.PostgreSQL.Simple as DB
|
||||
import qualified Database.PostgreSQL.Simple.Copy as DB
|
||||
import Database.PostgreSQL.Simple.SqlQQ (sql)
|
||||
import Database.PostgreSQL.Simple.ToField (ToField (..))
|
||||
import Simplex.Messaging.Agent.Store.Postgres.Common
|
||||
import qualified Simplex.Messaging.Crypto as C
|
||||
import Simplex.Messaging.Protocol
|
||||
import Simplex.Messaging.Server.MsgStore
|
||||
import Simplex.Messaging.Server.MsgStore.Types
|
||||
import Simplex.Messaging.Server.QueueStore
|
||||
import Simplex.Messaging.Server.QueueStore.Postgres
|
||||
import Simplex.Messaging.Server.QueueStore.Types
|
||||
import Simplex.Messaging.Server.StoreLog (foldLogLines)
|
||||
import Simplex.Messaging.Encoding.String
|
||||
import Simplex.Messaging.Util (maybeFirstRow, maybeFirstRow', (<$$>))
|
||||
import System.IO (Handle, hFlush, stdout)
|
||||
|
||||
data PostgresMsgStore = PostgresMsgStore
|
||||
{ config :: PostgresMsgStoreCfg,
|
||||
queueStore_ :: PostgresQueueStore'
|
||||
}
|
||||
|
||||
data PostgresMsgStoreCfg = PostgresMsgStoreCfg
|
||||
{ queueStoreCfg :: PostgresStoreCfg,
|
||||
quota :: Int
|
||||
}
|
||||
|
||||
type PostgresQueueStore' = PostgresQueueStore PostgresQueue
|
||||
|
||||
data PostgresQueue = PostgresQueue
|
||||
{ recipientId' :: RecipientId,
|
||||
queueRec' :: TVar (Maybe QueueRec)
|
||||
}
|
||||
|
||||
instance StoreQueueClass PostgresQueue where
|
||||
recipientId = recipientId'
|
||||
{-# INLINE recipientId #-}
|
||||
queueRec = queueRec'
|
||||
{-# INLINE queueRec #-}
|
||||
withQueueLock PostgresQueue {} _ = id -- TODO [messages] maybe it's just transaction?
|
||||
{-# INLINE withQueueLock #-}
|
||||
|
||||
newtype DBTransaction = DBTransaction {dbConn :: DB.Connection}
|
||||
|
||||
type DBStoreIO a = ReaderT DBTransaction IO a
|
||||
|
||||
instance MsgStoreClass PostgresMsgStore where
|
||||
type StoreMonad PostgresMsgStore = ReaderT DBTransaction IO
|
||||
type MsgQueue PostgresMsgStore = ()
|
||||
type QueueStore PostgresMsgStore = PostgresQueueStore'
|
||||
type StoreQueue PostgresMsgStore = PostgresQueue
|
||||
type MsgStoreConfig PostgresMsgStore = PostgresMsgStoreCfg
|
||||
|
||||
newMsgStore :: PostgresMsgStoreCfg -> IO PostgresMsgStore
|
||||
newMsgStore config = do
|
||||
queueStore_ <- newQueueStore @PostgresQueue $ queueStoreCfg config
|
||||
pure PostgresMsgStore {config, queueStore_}
|
||||
|
||||
closeMsgStore :: PostgresMsgStore -> IO ()
|
||||
closeMsgStore = closeQueueStore @PostgresQueue . queueStore_
|
||||
|
||||
withActiveMsgQueues _ _ = error "withActiveMsgQueues not used"
|
||||
|
||||
unsafeWithAllMsgQueues _ _ _ = error "unsafeWithAllMsgQueues not used"
|
||||
|
||||
expireOldMessages :: Bool -> PostgresMsgStore -> Int64 -> Int64 -> IO MessageStats
|
||||
expireOldMessages _tty ms now ttl =
|
||||
maybeFirstRow' newMessageStats toMessageStats $ withConnection st $ \db ->
|
||||
DB.query db "CALL expire_old_messages(?,?,0,0,0)" (now, ttl)
|
||||
where
|
||||
st = dbStore $ queueStore_ ms
|
||||
toMessageStats (expiredMsgsCount, storedMsgsCount, storedQueues) =
|
||||
MessageStats {expiredMsgsCount, storedMsgsCount, storedQueues}
|
||||
|
||||
logQueueStates _ = error "logQueueStates not used"
|
||||
|
||||
logQueueState _ = error "logQueueState not used"
|
||||
|
||||
queueStore = queueStore_
|
||||
{-# INLINE queueStore #-}
|
||||
|
||||
loadedQueueCounts :: PostgresMsgStore -> IO LoadedQueueCounts
|
||||
loadedQueueCounts ms = do
|
||||
loadedQueueCount <- M.size <$> readTVarIO queues
|
||||
loadedNotifierCount <- M.size <$> readTVarIO notifiers
|
||||
notifierLockCount <- M.size <$> readTVarIO notifierLocks
|
||||
pure LoadedQueueCounts {loadedQueueCount, loadedNotifierCount, openJournalCount = 0, queueLockCount = 0, notifierLockCount}
|
||||
where
|
||||
PostgresQueueStore {queues, notifiers, notifierLocks} = queueStore_ ms
|
||||
|
||||
mkQueue :: PostgresMsgStore -> Bool -> RecipientId -> QueueRec -> IO PostgresQueue
|
||||
mkQueue _ _keepLock rId qr = PostgresQueue rId <$> newTVarIO (Just qr)
|
||||
{-# INLINE mkQueue #-}
|
||||
|
||||
getMsgQueue _ _ _ = pure ()
|
||||
{-# INLINE getMsgQueue #-}
|
||||
|
||||
getPeekMsgQueue :: PostgresMsgStore -> PostgresQueue -> DBStoreIO (Maybe ((), Message))
|
||||
getPeekMsgQueue _ q = ((),) <$$> tryPeekMsg_ q ()
|
||||
|
||||
withIdleMsgQueue :: Int64 -> PostgresMsgStore -> PostgresQueue -> (() -> DBStoreIO a) -> DBStoreIO (Maybe a, Int)
|
||||
withIdleMsgQueue _ _ _ _ = error "withIdleMsgQueue not used"
|
||||
|
||||
deleteQueue :: PostgresMsgStore -> PostgresQueue -> IO (Either ErrorType QueueRec)
|
||||
deleteQueue ms q = deleteStoreQueue (queueStore_ ms) q
|
||||
{-# INLINE deleteQueue #-}
|
||||
|
||||
deleteQueueSize :: PostgresMsgStore -> PostgresQueue -> IO (Either ErrorType (QueueRec, Int))
|
||||
deleteQueueSize ms q = runExceptT $ do
|
||||
size <- getQueueSize ms q
|
||||
qr <- ExceptT $ deleteStoreQueue (queueStore_ ms) q
|
||||
pure (qr, size)
|
||||
|
||||
getQueueMessages_ _ _ _ = error "getQueueMessages_ not used"
|
||||
|
||||
writeMsg :: PostgresMsgStore -> PostgresQueue -> Bool -> Message -> ExceptT ErrorType IO (Maybe (Message, Bool))
|
||||
writeMsg ms q _ msg =
|
||||
withDB' "writeMsg" (queueStore_ ms) $ \db -> do
|
||||
let (msgQuota, ntf, body) = case msg of
|
||||
Message {msgFlags = MsgFlags ntf', msgBody = C.MaxLenBS body'} -> (False, ntf', body')
|
||||
MessageQuota {} -> (True, False, B.empty)
|
||||
toResult <$>
|
||||
DB.query
|
||||
db
|
||||
"SELECT quota_written, was_empty FROM write_message(?,?,?,?,?,?,?)"
|
||||
(recipientId' q, Binary (messageId msg), systemSeconds (messageTs msg), msgQuota, ntf, Binary body, quota)
|
||||
where
|
||||
toResult = \case
|
||||
((msgQuota, wasEmpty) : _) -> if msgQuota then Nothing else Just (msg, wasEmpty)
|
||||
[] -> Nothing
|
||||
PostgresMsgStore {config = PostgresMsgStoreCfg {quota}} = ms
|
||||
|
||||
setOverQuota_ :: PostgresQueue -> IO () -- can ONLY be used while restoring messages, not while server running
|
||||
setOverQuota_ _ = error "TODO setOverQuota_" -- TODO [messages]
|
||||
|
||||
getQueueSize_ :: () -> DBStoreIO Int
|
||||
getQueueSize_ _ = error "getQueueSize_ not used"
|
||||
|
||||
getQueueSize :: PostgresMsgStore -> PostgresQueue -> ExceptT ErrorType IO Int
|
||||
getQueueSize ms q =
|
||||
withDB' "getQueueSize" (queueStore_ ms) $ \db ->
|
||||
maybeFirstRow' 0 fromOnly $
|
||||
DB.query db "SELECT msg_queue_size FROM msg_queues WHERE recipient_id = ? AND deleted_at IS NULL" (Only (recipientId' q))
|
||||
|
||||
tryPeekMsg_ :: PostgresQueue -> () -> DBStoreIO (Maybe Message)
|
||||
tryPeekMsg_ q _ = do
|
||||
db <- asks dbConn
|
||||
liftIO $ maybeFirstRow toMessage $
|
||||
DB.query
|
||||
db
|
||||
[sql|
|
||||
SELECT msg_id, msg_ts, msg_quota, msg_ntf_flag, msg_body
|
||||
FROM messages
|
||||
WHERE recipient_id = ?
|
||||
ORDER BY message_id ASC LIMIT 1
|
||||
|]
|
||||
(Only (recipientId' q))
|
||||
|
||||
tryDeleteMsg_ :: PostgresQueue -> () -> Bool -> DBStoreIO ()
|
||||
tryDeleteMsg_ _q _ _ = error "tryDeleteMsg_ not used" -- do
|
||||
|
||||
isolateQueue :: PostgresMsgStore -> PostgresQueue -> Text -> DBStoreIO a -> ExceptT ErrorType IO a
|
||||
isolateQueue ms _q op a = withDB' op (queueStore_ ms) $ runReaderT a . DBTransaction
|
||||
|
||||
unsafeRunStore _ _ _ = error "unsafeRunStore not used"
|
||||
|
||||
tryPeekMsg :: PostgresMsgStore -> PostgresQueue -> ExceptT ErrorType IO (Maybe Message)
|
||||
tryPeekMsg ms q = isolateQueue ms q "tryPeekMsg" $ tryPeekMsg_ q ()
|
||||
{-# INLINE tryPeekMsg #-}
|
||||
|
||||
tryDelMsg :: PostgresMsgStore -> PostgresQueue -> MsgId -> ExceptT ErrorType IO (Maybe Message)
|
||||
tryDelMsg ms q msgId =
|
||||
withDB' "tryDelMsg" (queueStore_ ms) $ \db ->
|
||||
maybeFirstRow toMessage $
|
||||
DB.query db "SELECT r_msg_id, r_msg_ts, r_msg_quota, r_msg_ntf_flag, r_msg_body FROM try_del_msg(?, ?)" (recipientId' q, Binary msgId)
|
||||
|
||||
tryDelPeekMsg :: PostgresMsgStore -> PostgresQueue -> MsgId -> ExceptT ErrorType IO (Maybe Message, Maybe Message)
|
||||
tryDelPeekMsg ms q msgId =
|
||||
withDB' "tryDelPeekMsg" (queueStore_ ms) $ \db ->
|
||||
toResult . map toMessage
|
||||
<$> DB.query db "SELECT r_msg_id, r_msg_ts, r_msg_quota, r_msg_ntf_flag, r_msg_body FROM try_del_peek_msg(?, ?)" (recipientId' q, Binary msgId)
|
||||
where
|
||||
toResult = \case
|
||||
[] -> (Nothing, Nothing)
|
||||
[msg]
|
||||
| messageId msg == msgId -> (Just msg, Nothing)
|
||||
| otherwise -> (Nothing, Just msg)
|
||||
deleted : next : _ -> (Just deleted, Just next)
|
||||
|
||||
deleteExpiredMsgs :: PostgresMsgStore -> PostgresQueue -> Int64 -> ExceptT ErrorType IO Int
|
||||
deleteExpiredMsgs ms q old =
|
||||
maybeFirstRow' 0 (fromIntegral @Int64 . fromOnly) $ withDB' "deleteExpiredMsgs" (queueStore_ ms) $ \db ->
|
||||
DB.query db "SELECT delete_expired_msgs(?, ?)" (recipientId' q, old)
|
||||
|
||||
toMessage :: (Binary MsgId, Int64, Bool, Bool, Binary MsgBody) -> Message
|
||||
toMessage (Binary msgId, ts, msgQuota, ntf, Binary body)
|
||||
| msgQuota = MessageQuota {msgId, msgTs}
|
||||
| otherwise = Message {msgId, msgTs, msgFlags = MsgFlags ntf, msgBody = C.unsafeMaxLenBS body} -- TODO [messages] unsafeMaxLenBS?
|
||||
where
|
||||
msgTs = MkSystemTime ts 0
|
||||
|
||||
exportDbMessages :: Bool -> PostgresMsgStore -> Handle -> IO Int
|
||||
exportDbMessages tty ms h = do
|
||||
rows <- newIORef []
|
||||
n <- withConnection st $ \db -> DB.foldWithOptions_ opts db query 0 $ \i r -> do
|
||||
let i' = i + 1
|
||||
if i' `mod` 1000 > 0
|
||||
then modifyIORef rows (r :)
|
||||
else do
|
||||
readIORef rows >>= writeMessages . (r :)
|
||||
writeIORef rows []
|
||||
when tty $ putStr (progress i' <> "\r") >> hFlush stdout
|
||||
pure i'
|
||||
readIORef rows >>= \rs -> unless (null rs) $ writeMessages rs
|
||||
when tty $ putStrLn $ progress n
|
||||
pure n
|
||||
where
|
||||
st = dbStore $ queueStore_ ms
|
||||
opts = DB.defaultFoldOptions {DB.fetchQuantity = DB.Fixed 1000}
|
||||
query =
|
||||
[sql|
|
||||
SELECT recipient_id, msg_id, msg_ts, msg_quota, msg_ntf_flag, msg_body
|
||||
FROM messages
|
||||
ORDER BY recipient_id, message_id ASC
|
||||
|]
|
||||
writeMessages = BB.hPutBuilder h . encodeMessages . reverse
|
||||
encodeMessages = mconcat . map (\(Only rId :. msg) -> BB.byteString (strEncode $ MLRv3 rId $ toMessage msg) <> BB.char8 '\n')
|
||||
progress i = "Processed: " <> show i <> " records"
|
||||
|
||||
getDbMessageStats :: PostgresMsgStore -> IO MessageStats
|
||||
getDbMessageStats ms =
|
||||
maybeFirstRow' newMessageStats toMessageStats $ withConnection st $ \db ->
|
||||
DB.query_
|
||||
db
|
||||
[sql|
|
||||
SELECT
|
||||
(SELECT COUNT (1) FROM msg_queues WHERE deleted_at IS NULL),
|
||||
(SELECT COUNT (1) FROM messages m JOIN msg_queues q USING recipient_id WHERE deleted_at IS NULL)
|
||||
|]
|
||||
where
|
||||
st = dbStore $ queueStore_ ms
|
||||
toMessageStats (storedQueues, storedMsgsCount) =
|
||||
MessageStats {storedQueues, storedMsgsCount, expiredMsgsCount = 0}
|
||||
|
||||
-- TODO [messages] update counts
|
||||
batchInsertMessages :: StoreQueueClass q => Bool -> FilePath -> PostgresQueueStore q -> IO Int64
|
||||
batchInsertMessages tty f toStore = do
|
||||
putStrLn "Importing messages..."
|
||||
let st = dbStore toStore
|
||||
(_, inserted) <-
|
||||
withTransaction st $ \db -> do
|
||||
DB.copy_
|
||||
db
|
||||
[sql|
|
||||
COPY messages (recipient_id, msg_id, msg_ts, msg_quota, msg_ntf_flag, msg_body)
|
||||
FROM STDIN WITH (FORMAT CSV)
|
||||
|]
|
||||
foldLogLines tty f (putMessage db) (0 :: Int, 0) >>= (DB.putCopyEnd db $>)
|
||||
Only mCnt : _ <- withTransaction st (`DB.query_` "SELECT count(*) FROM messages")
|
||||
unless (inserted == mCnt) $ putStrLn $ "WARNING: inserted " <> show inserted <> " rows, table has " <> show mCnt <> " messages."
|
||||
pure inserted
|
||||
where
|
||||
putMessage db (!i, !cnt) _eof s = do
|
||||
let i' = i + 1
|
||||
cnt' <- case strDecode s of
|
||||
Right (MLRv3 rId msg) -> (cnt + 1) <$ DB.putCopyData db (messageRecToText rId msg)
|
||||
Left e -> cnt <$ putStrLn ("Error parsing line " <> show i' <> ": " <> e)
|
||||
pure (i', cnt')
|
||||
|
||||
messageRecToText :: RecipientId -> Message -> B.ByteString
|
||||
messageRecToText rId msg =
|
||||
LB.toStrict $ BB.toLazyByteString $ mconcat tabFields <> BB.char7 '\n'
|
||||
where
|
||||
tabFields = BB.char7 ',' `intersperse` fields
|
||||
fields =
|
||||
[ renderField (toField rId),
|
||||
renderField (toField $ Binary (messageId msg)),
|
||||
renderField (toField $ systemSeconds (messageTs msg)),
|
||||
renderField (toField msgQuota),
|
||||
renderField (toField ntf),
|
||||
renderField (toField $ Binary body)
|
||||
]
|
||||
(msgQuota, ntf, body) = case msg of
|
||||
Message {msgFlags = MsgFlags ntf', msgBody = C.MaxLenBS body'} -> (False, ntf', body')
|
||||
MessageQuota {} -> (True, False, B.empty)
|
||||
Reference in New Issue
Block a user