agent: store shared message body only once (when it is the same across messages when batching) (#1453)

* agent: store shared message body only once (when it is the same across messages when batching) * rename * refactor * refactor * save bodies and messages in single transaction * comment * comment * comment * box * mapME * box * ValueOrRef * remove instances * refactor * comments * test * refactor * mapAccumLM compatibility with ghc 8.10.7 --------- Co-authored-by: Evgeny Poberezkin <evgeny@poberezkin.com>
2026-08-28 05:04:40 +00:00 · 2025-02-14 18:01:40 +00:00
parent 0d8a1a2879
commit 7ac80bffcb
13 changed files with 294 additions and 95 deletions
@@ -35,6 +35,8 @@ module Simplex.Messaging.Agent
    AE,
    SubscriptionsInfo (..),
    MsgReq,
+    ValueOrRef (..),
+    vrValue,
    getSMPAgentClient,
    getSMPAgentClient_,
    disconnectAgentClient,
@@ -140,6 +142,9 @@ import Data.Either (isRight, partitionEithers, rights)
 import Data.Foldable (foldl', toList)
 import Data.Functor (($>))
 import Data.Functor.Identity
+import Data.Int (Int64)
+import Data.IntMap.Strict (IntMap)
+import qualified Data.IntMap.Strict as IM
 import Data.List (find)
 import Data.List.NonEmpty (NonEmpty (..))
 import qualified Data.List.NonEmpty as L
@@ -409,11 +414,25 @@ sendMessage :: AgentClient -> ConnId -> PQEncryption -> MsgFlags -> MsgBody -> A
 sendMessage c = withAgentEnv c .:: sendMessage' c
 {-# INLINE sendMessage #-}

+data ValueOrRef a = VRValue (Maybe Int) a | VRRef Int
+
+instance Functor ValueOrRef where
+  fmap f = \case
+    VRValue i_ a -> VRValue i_ (f a)
+    VRRef i -> VRRef i
+
+vrValue :: a -> ValueOrRef a
+vrValue = VRValue Nothing
+
 -- When sending multiple messages to the same connection,
 -- only the first MsgReq for this connection should have non-empty ConnId.
 -- All subsequent MsgReq in traversable for this connection must be empty.
 -- This is done to optimize processing by grouping all messages to one connection together.
-type MsgReq = (ConnId, PQEncryption, MsgFlags, MsgBody)
+-- Also, repeated msg bodies should us MBRef constructor to reference previously used body.
+-- It is an error:
+-- - to use MBBody with the same Int
+-- - to use MBRef with Int that wasn't previously used in MBBody
+type MsgReq = (ConnId, PQEncryption, MsgFlags, ValueOrRef MsgBody)

 -- | Send multiple messages to different connections (SEND command)
 sendMessages :: AgentClient -> [MsgReq] -> AE [Either AgentErrorType (AgentMsgId, PQEncryption)]
@@ -1125,7 +1144,7 @@ getNotificationConns' c nonce encNtfInfo =

 -- | Send message to the connection (SEND command) in Reader monad
 sendMessage' :: AgentClient -> ConnId -> PQEncryption -> MsgFlags -> MsgBody -> AM (AgentMsgId, PQEncryption)
-sendMessage' c connId pqEnc msgFlags msg = ExceptT $ runIdentity <$> sendMessagesB_ c (Identity (Right (connId, pqEnc, msgFlags, msg))) (S.singleton connId)
+sendMessage' c connId pqEnc msgFlags msg = ExceptT $ runIdentity <$> sendMessagesB_ c (Identity (Right (connId, pqEnc, msgFlags, vrValue msg))) (S.singleton connId)
 {-# INLINE sendMessage' #-}

 -- | Send multiple messages to different connections (SEND command) in Reader monad
@@ -1160,14 +1179,14 @@ sendMessagesB_ c reqs connIds = withConnLocks c connIds "sendMessages" $ do
          else do
            conn <- first storeError <$> getConn db connId
            conn <$ atomically (writeTVar prev $ Just conn)
-    prepareConn :: Set ConnId -> Either AgentErrorType (MsgReq, SomeConn) -> (Set ConnId, Either AgentErrorType (ConnData, NonEmpty SndQueue, Maybe PQEncryption, MsgFlags, AMessage))
+    prepareConn :: Set ConnId -> Either AgentErrorType (MsgReq, SomeConn) -> (Set ConnId, Either AgentErrorType (ConnData, NonEmpty SndQueue, Maybe PQEncryption, MsgFlags, ValueOrRef AMessage))
    prepareConn s (Left e) = (s, Left e)
-    prepareConn s (Right ((_, pqEnc, msgFlags, msg), SomeConn _ conn)) = case conn of
+    prepareConn s (Right ((_, pqEnc, msgFlags, msgOrRef), SomeConn _ conn)) = case conn of
      DuplexConnection cData _ sqs -> prepareMsg cData sqs
      SndConnection cData sq -> prepareMsg cData [sq]
      _ -> (s, Left $ CONN SIMPLEX)
      where
-        prepareMsg :: ConnData -> NonEmpty SndQueue -> (Set ConnId, Either AgentErrorType (ConnData, NonEmpty SndQueue, Maybe PQEncryption, MsgFlags, AMessage))
+        prepareMsg :: ConnData -> NonEmpty SndQueue -> (Set ConnId, Either AgentErrorType (ConnData, NonEmpty SndQueue, Maybe PQEncryption, MsgFlags, ValueOrRef AMessage))
        prepareMsg cData@ConnData {connId, pqSupport} sqs
          | ratchetSyncSendProhibited cData = (s, Left $ CMD PROHIBITED "sendMessagesB: send prohibited")
          -- connection is only updated if PQ encryption was disabled, and now it has to be enabled.
@@ -1177,7 +1196,7 @@ sendMessagesB_ c reqs connIds = withConnLocks c connIds "sendMessages" $ do
               in (S.insert connId s, mkReq cData')
          | otherwise = (s, mkReq cData)
          where
-            mkReq cData' = Right (cData', sqs, Just pqEnc, msgFlags, A_MSG msg)
+            mkReq cData' = Right (cData', sqs, Just pqEnc, msgFlags, A_MSG <$> msgOrRef)

 -- / async command processing v v v

@@ -1361,10 +1380,10 @@ enqueueMessages c cData sqs msgFlags aMessage = do

 enqueueMessages' :: AgentClient -> ConnData -> NonEmpty SndQueue -> MsgFlags -> AMessage -> AM (AgentMsgId, CR.PQEncryption)
 enqueueMessages' c cData sqs msgFlags aMessage =
-  ExceptT $ runIdentity <$> enqueueMessagesB c (Identity (Right (cData, sqs, Nothing, msgFlags, aMessage)))
+  ExceptT $ runIdentity <$> enqueueMessagesB c (Identity (Right (cData, sqs, Nothing, msgFlags, vrValue aMessage)))
 {-# INLINE enqueueMessages' #-}

-enqueueMessagesB :: Traversable t => AgentClient -> t (Either AgentErrorType (ConnData, NonEmpty SndQueue, Maybe PQEncryption, MsgFlags, AMessage)) -> AM' (t (Either AgentErrorType (AgentMsgId, PQEncryption)))
+enqueueMessagesB :: Traversable t => AgentClient -> t (Either AgentErrorType (ConnData, NonEmpty SndQueue, Maybe PQEncryption, MsgFlags, ValueOrRef AMessage)) -> AM' (t (Either AgentErrorType (AgentMsgId, PQEncryption)))
 enqueueMessagesB c reqs = do
  reqs' <- enqueueMessageB c reqs
  enqueueSavedMessageB c $ mapMaybe snd $ rights $ toList reqs'
@@ -1376,40 +1395,62 @@ isActiveSndQ SndQueue {status} = status == Secured || status == Active

 enqueueMessage :: AgentClient -> ConnData -> SndQueue -> MsgFlags -> AMessage -> AM (AgentMsgId, PQEncryption)
 enqueueMessage c cData sq msgFlags aMessage =
-  ExceptT $ fmap fst . runIdentity <$> enqueueMessageB c (Identity (Right (cData, [sq], Nothing, msgFlags, aMessage)))
+  ExceptT $ fmap fst . runIdentity <$> enqueueMessageB c (Identity (Right (cData, [sq], Nothing, msgFlags, vrValue aMessage)))
 {-# INLINE enqueueMessage #-}

-- TODO [save once] IntMap of msg bodies.
 -- this function is used only for sending messages in batch, it returns the list of successes to enqueue additional deliveries
-enqueueMessageB :: forall t. Traversable t => AgentClient -> t (Either AgentErrorType (ConnData, NonEmpty SndQueue, Maybe PQEncryption, MsgFlags, AMessage)) -> AM' (t (Either AgentErrorType ((AgentMsgId, PQEncryption), Maybe (ConnData, [SndQueue], AgentMsgId))))
+enqueueMessageB :: forall t. Traversable t => AgentClient -> t (Either AgentErrorType (ConnData, NonEmpty SndQueue, Maybe PQEncryption, MsgFlags, ValueOrRef AMessage)) -> AM' (t (Either AgentErrorType ((AgentMsgId, PQEncryption), Maybe (ConnData, [SndQueue], AgentMsgId))))
 enqueueMessageB c reqs = do
  cfg <- asks config
-  reqMids <- withStoreBatch c $ \db -> fmap (bindRight $ storeSentMsg db cfg) reqs
+  (_, reqMids) <- unsafeWithStore c $ \db -> do
+    mapAccumLM (\ids r -> storeSentMsg db cfg ids r `E.catchAny` \e -> (ids,) <$> handleInternal e) IM.empty reqs
  forME reqMids $ \((cData, sq :| sqs, _, _, _), InternalId msgId, pqSecr) -> do
    submitPendingMsg c cData sq
    let sqs' = filter isActiveSndQ sqs
    pure $ Right ((msgId, pqSecr), if null sqs' then Nothing else Just (cData, sqs', msgId))
  where
-    storeSentMsg :: DB.Connection -> AgentConfig -> (ConnData, NonEmpty SndQueue, Maybe PQEncryption, MsgFlags, AMessage) -> IO (Either AgentErrorType ((ConnData, NonEmpty SndQueue, Maybe PQEncryption, MsgFlags, AMessage), InternalId, PQEncryption))
-    storeSentMsg db cfg req@(cData@ConnData {connId}, sq :| _, pqEnc_, msgFlags, aMessage) = fmap (first storeError) $ runExceptT $ do
-      let AgentConfig {e2eEncryptVRange} = cfg
-      internalTs <- liftIO getCurrentTime
-      (internalId, internalSndId, prevMsgHash) <- ExceptT $ updateSndIds db connId
-      let privHeader = APrivHeader (unSndId internalSndId) prevMsgHash
-          agentMsg = AgentMessage privHeader aMessage
-          agentMsgStr = smpEncode agentMsg
-          internalHash = C.sha256Hash agentMsgStr
-          currentE2EVersion = maxVersion e2eEncryptVRange
-      -- TODO [save once] Save single MsgBody / enveloped body agentMsgStr (outside of withStoreBatch ... storeSentMsg).
-      -- TODO             Link messages to it, save encryption data per message.
-      -- TODO             'msg_body' field is not nullable - use default empty strings?
-      (mek, paddedLen, pqEnc) <- agentRatchetEncryptHeader db cData e2eEncAgentMsgLength pqEnc_ currentE2EVersion
-      withExceptT (SEAgentError . cryptoError) $ CR.rcCheckCanPad paddedLen agentMsgStr
-      let msgType = agentMessageType agentMsg
-          msgData = SndMsgData {internalId, internalSndId, internalTs, msgType, msgFlags, msgBody = agentMsgStr, pqEncryption = pqEnc, internalHash, prevMsgHash, encryptKey_ = Just mek, paddedLen_ = Just paddedLen}
-      liftIO $ createSndMsg db connId msgData
-      liftIO $ createSndMsgDelivery db connId sq internalId
-      pure (req, internalId, pqEnc)
+    storeSentMsg :: DB.Connection -> AgentConfig -> IntMap (Int64, AMessage) -> Either AgentErrorType (ConnData, NonEmpty SndQueue, Maybe PQEncryption, MsgFlags, ValueOrRef AMessage) -> IO (IntMap (Int64, AMessage), Either AgentErrorType ((ConnData, NonEmpty SndQueue, Maybe PQEncryption, MsgFlags, ValueOrRef AMessage), InternalId, PQEncryption))
+    storeSentMsg db cfg aMessageIds = \case
+      Left e -> pure (aMessageIds, Left e)
+      Right req@(cData@ConnData {connId}, sq :| _, pqEnc_, msgFlags, mbr) -> case mbr of
+        VRValue i_ aMessage -> case  i_ >>= (`IM.lookup` aMessageIds) of
+          Just _ -> pure (aMessageIds, Left $ INTERNAL "enqueueMessageB: storeSentMsg duplicate saved message body")
+          Nothing -> do
+            mbId <- createSndMsgBody db aMessage
+            let aMessageIds' = maybe id (`IM.insert` (mbId, aMessage)) i_ aMessageIds
+            (aMessageIds',) <$> storeSentMsg_ mbId aMessage
+        VRRef i -> (aMessageIds,) <$> case IM.lookup i aMessageIds of
+          Just (mbId, aMessage) -> storeSentMsg_ mbId aMessage
+          Nothing -> pure $ Left $ INTERNAL "enqueueMessageB: storeSentMsg missing saved message body id"
+        where
+          storeSentMsg_ sndMsgBodyId aMessage = fmap (first storeError) $ runExceptT $ do
+            let AgentConfig {e2eEncryptVRange} = cfg
+            internalTs <- liftIO getCurrentTime
+            (internalId, internalSndId, prevMsgHash) <- ExceptT $ updateSndIds db connId
+            -- We need to do pre-flight encoding that is not stored in database
+            -- to calculate its hash and remember it on connection (createSndMsg -> updateSndMsgHash)
+            -- to enable next enqueue.
+            -- (As encoding is different per connection, we can't store shared body, so it's repeated on delivery)
+            let agentMsgStr = encodeAgentMsgStr aMessage internalSndId prevMsgHash
+                internalHash = C.sha256Hash agentMsgStr
+                currentE2EVersion = maxVersion e2eEncryptVRange
+            (mek, paddedLen, pqEnc) <- agentRatchetEncryptHeader db cData e2eEncAgentMsgLength pqEnc_ currentE2EVersion
+            withExceptT (SEAgentError . cryptoError) $ CR.rcCheckCanPad paddedLen agentMsgStr
+            let msgType = aMessageType aMessage
+                -- msgBody is empty, because snd_messages record is linked to snd_message_bodies
+                msgData = SndMsgData {internalId, internalSndId, internalTs, msgType, msgFlags, msgBody = "", pqEncryption = pqEnc, internalHash, prevMsgHash, sndMsgPrepData_ = Just SndMsgPrepData {encryptKey = mek, paddedLen, sndMsgBodyId}}
+            liftIO $ createSndMsg db connId msgData
+            liftIO $ createSndMsgDelivery db connId sq internalId
+            pure (req, internalId, pqEnc)
+    handleInternal :: E.SomeException -> IO (Either AgentErrorType b)
+    handleInternal = pure . Left . INTERNAL . show
+
+
+encodeAgentMsgStr :: AMessage -> InternalSndId -> PrevSndMsgHash -> ByteString
+encodeAgentMsgStr aMessage internalSndId prevMsgHash = do
+  let privHeader = APrivHeader (unSndId internalSndId) prevMsgHash
+      agentMsg = AgentMessage privHeader aMessage
+   in smpEncode agentMsg

 enqueueSavedMessage :: AgentClient -> ConnData -> AgentMsgId -> SndQueue -> AM' ()
 enqueueSavedMessage c cData msgId sq = enqueueSavedMessageB c $ Identity (cData, [sq], msgId)
@@ -1454,7 +1495,7 @@ runSmpQueueMsgDelivery c@AgentClient {subQ} ConnData {connId} sq@SndQueue {userI
    liftIO $ throwWhenNoDelivery c sq
    atomically $ beginAgentOperation c AOSndNetwork
    withWork c doWork (\db -> getPendingQueueMsg db connId sq) $
-      \(rq_, PendingMsgData {msgId, msgType, msgBody, pqEncryption, msgFlags, msgRetryState, internalTs, encryptKey_, paddedLen_}) -> do
+      \(rq_, PendingMsgData {msgId, msgType, msgBody, pqEncryption, msgFlags, msgRetryState, internalTs, internalSndId, prevMsgHash, pendingMsgPrepData_}) -> do
        atomically $ endAgentOperation c AOMsgDelivery -- this operation begins in submitPendingMsg
        let mId = unId msgId
            ri' = maybe id updateRetryInterval2 msgRetryState ri
@@ -1464,15 +1505,15 @@ runSmpQueueMsgDelivery c@AgentClient {subQ} ConnData {connId} sq@SndQueue {userI
          resp <- tryError $ case msgType of
            AM_CONN_INFO -> sendConfirmation c sq msgBody
            AM_CONN_INFO_REPLY -> sendConfirmation c sq msgBody
-            _ -> case (encryptKey_, paddedLen_) of
-              (Nothing, Nothing) -> sendAgentMessage c sq msgFlags msgBody
-              (Just mek, Just paddedLen) -> do
+            _ -> case pendingMsgPrepData_ of
+              Nothing -> sendAgentMessage c sq msgFlags msgBody
+              Just PendingMsgPrepData {encryptKey, paddedLen, sndMsgBody} -> do
+                let agentMsgStr = encodeAgentMsgStr sndMsgBody internalSndId prevMsgHash
                AgentConfig {smpAgentVRange} <- asks config
-                encAgentMessage <- liftError cryptoError $ CR.rcEncryptMsg mek paddedLen msgBody
+                encAgentMessage <- liftError cryptoError $ CR.rcEncryptMsg encryptKey paddedLen agentMsgStr
                let agentVersion = maxVersion smpAgentVRange
                    msgBody' = smpEncode $ AgentMsgEnvelope {agentVersion, encAgentMessage}
                sendAgentMessage c sq msgFlags msgBody'
-              _ -> throwE $ INTERNAL "runSmpQueueMsgDelivery: missing encryption data"
          case resp of
            Left e -> do
              let err = if msgType == AM_A_MSG_ then MERR mId e else ERR e
@@ -2963,7 +3004,7 @@ storeConfirmation c cData@ConnData {connId, pqSupport, connAgentVersion = v} sq
    (encConnInfo, pqEncryption) <- agentRatchetEncrypt db cData agentMsgStr e2eEncConnInfoLength (Just pqEnc) currentE2EVersion
    let msgBody = smpEncode $ AgentConfirmation {agentVersion = v, e2eEncryption_, encConnInfo}
        msgType = agentMessageType agentMsg
-        msgData = SndMsgData {internalId, internalSndId, internalTs, msgType, msgBody, pqEncryption, msgFlags = SMP.MsgFlags {notification = True}, internalHash, prevMsgHash, encryptKey_ = Nothing, paddedLen_ = Nothing}
+        msgData = SndMsgData {internalId, internalSndId, internalTs, msgType, msgBody, pqEncryption, msgFlags = SMP.MsgFlags {notification = True}, internalHash, prevMsgHash, sndMsgPrepData_ = Nothing}
    liftIO $ createSndMsg db connId msgData
    liftIO $ createSndMsgDelivery db connId sq internalId

@@ -2989,7 +3030,7 @@ enqueueRatchetKey c cData@ConnData {connId} sq e2eEncryption = do
      let msgBody = smpEncode $ AgentRatchetKey {agentVersion, e2eEncryption, info = agentMsgStr}
          msgType = agentMessageType agentMsg
          -- this message is e2e encrypted with queue key, not with double ratchet
-          msgData = SndMsgData {internalId, internalSndId, internalTs, msgType, msgBody, pqEncryption = PQEncOff, msgFlags = SMP.MsgFlags {notification = True}, internalHash, prevMsgHash, encryptKey_ = Nothing, paddedLen_ = Nothing}
+          msgData = SndMsgData {internalId, internalSndId, internalTs, msgType, msgBody, pqEncryption = PQEncOff, msgFlags = SMP.MsgFlags {notification = True}, internalHash, prevMsgHash, sndMsgPrepData_ = Nothing}
      liftIO $ createSndMsg db connId msgData
      liftIO $ createSndMsgDelivery db connId sq internalId
      pure internalId
@@ -146,6 +146,7 @@ module Simplex.Messaging.Agent.Client
    withStore',
    withStoreBatch,
    withStoreBatch',
+    unsafeWithStore,
    storeError,
    userServers,
    pickServer,
@@ -2009,6 +2010,11 @@ withStore c action = do
      ]
 #endif

+unsafeWithStore :: AgentClient -> (DB.Connection -> IO a) -> AM' a
+unsafeWithStore c action = do
+  st <- asks store
+  liftIO $ agentOperationBracket c AODatabase (\_ -> pure ()) $ withTransaction st action
+
 withStoreBatch :: Traversable t => AgentClient -> (DB.Connection -> t (IO (Either AgentErrorType a))) -> AM' (t (Either AgentErrorType a))
 withStoreBatch c actions = do
  st <- asks store
@@ -140,6 +140,7 @@ module Simplex.Messaging.Agent.Protocol
    serializeQueueStatus,
    queueStatusT,
    agentMessageType,
+    aMessageType,
    extraSMPServerHosts,
    updateSMPServerHosts,
  )
@@ -167,7 +168,7 @@ import Data.Time.Clock.System (SystemTime)
 import Data.Type.Equality
 import Data.Typeable ()
 import Data.Word (Word16, Word32)
-import Simplex.Messaging.Agent.Store.DB (FromField (..), ToField (..))
+import Simplex.Messaging.Agent.Store.DB (Binary (..), FromField (..), ToField (..))
 import Simplex.FileTransfer.Description
 import Simplex.FileTransfer.Protocol (FileParty (..))
 import Simplex.FileTransfer.Transport (XFTPErrorType)
@@ -855,20 +856,7 @@ agentMessageType = \case
  AgentConnInfo _ -> AM_CONN_INFO
  AgentConnInfoReply {} -> AM_CONN_INFO_REPLY
  AgentRatchetInfo _ -> AM_RATCHET_INFO
-  AgentMessage _ aMsg -> case aMsg of
-    -- HELLO is used both in v1 and in v2, but differently.
-    -- - in v1 (and, possibly, in v2 for simplex connections) can be sent multiple times,
-    --   until the queue is secured - the OK response from the server instead of initial AUTH errors confirms it.
-    -- - in v2 duplexHandshake it is sent only once, when it is known that the queue was secured.
-    HELLO -> AM_HELLO_
-    A_MSG _ -> AM_A_MSG_
-    A_RCVD {} -> AM_A_RCVD_
-    A_QCONT _ -> AM_QCONT_
-    QADD _ -> AM_QADD_
-    QKEY _ -> AM_QKEY_
-    QUSE _ -> AM_QUSE_
-    QTEST _ -> AM_QTEST_
-    EREADY _ -> AM_EREADY_
+  AgentMessage _ aMsg -> aMessageType aMsg

 data APrivHeader = APrivHeader
  { -- | sequential ID assigned by the sending agent
@@ -946,6 +934,22 @@ data AMessage
    EREADY AgentMsgId
  deriving (Show)

+aMessageType :: AMessage -> AgentMessageType
+aMessageType = \case
+  -- HELLO is used both in v1 and in v2, but differently.
+  -- - in v1 (and, possibly, in v2 for simplex connections) can be sent multiple times,
+  --   until the queue is secured - the OK response from the server instead of initial AUTH errors confirms it.
+  -- - in v2 duplexHandshake it is sent only once, when it is known that the queue was secured.
+  HELLO -> AM_HELLO_
+  A_MSG _ -> AM_A_MSG_
+  A_RCVD {} -> AM_A_RCVD_
+  A_QCONT _ -> AM_QCONT_
+  QADD _ -> AM_QADD_
+  QKEY _ -> AM_QKEY_
+  QUSE _ -> AM_QUSE_
+  QTEST _ -> AM_QTEST_
+  EREADY _ -> AM_EREADY_
+
 -- | this type is used to send as part of the protocol between different clients
 -- TODO possibly, rename fields and types referring to external and internal IDs to make them different
 data AMessageReceipt = AMessageReceipt
@@ -1010,6 +1014,10 @@ instance Encoding AMessage where
        QTEST_ -> QTEST <$> smpP
        EREADY_ -> EREADY <$> smpP

+instance ToField AMessage where toField = toField . Binary . smpEncode
+
+instance FromField AMessage where fromField = blobFieldParser smpP
+
 instance Encoding AMessageReceipt where
  smpEncode AMessageReceipt {agentMsgId, msgHash, rcptInfo} =
    smpEncode (agentMsgId, msgHash, Large rcptInfo)
@@ -543,10 +543,16 @@ data SndMsgData = SndMsgData
    pqEncryption :: PQEncryption,
    internalHash :: MsgHash,
    prevMsgHash :: MsgHash,
-    encryptKey_ :: Maybe MsgEncryptKeyX448,
-    paddedLen_ :: Maybe Int
+    sndMsgPrepData_ :: Maybe SndMsgPrepData
  }

+data SndMsgPrepData = SndMsgPrepData
+  { encryptKey :: MsgEncryptKeyX448,
+    paddedLen :: Int,
+    sndMsgBodyId :: Int64
+  }
+  deriving (Show)
+
 data SndMsg = SndMsg
  { internalId :: InternalId,
    internalSndId :: InternalSndId,
@@ -563,8 +569,16 @@ data PendingMsgData = PendingMsgData
    pqEncryption :: PQEncryption,
    msgRetryState :: Maybe RI2State,
    internalTs :: InternalTs,
-    encryptKey_ :: Maybe MsgEncryptKeyX448,
-    paddedLen_ :: Maybe Int
+    internalSndId :: InternalSndId,
+    prevMsgHash :: PrevSndMsgHash,
+    pendingMsgPrepData_ :: Maybe PendingMsgPrepData
+  }
+  deriving (Show)
+
+data PendingMsgPrepData = PendingMsgPrepData
+  { encryptKey :: MsgEncryptKeyX448,
+    paddedLen :: Int,
+    sndMsgBody :: AMessage
  }
  deriving (Show)

@@ -92,6 +92,7 @@ module Simplex.Messaging.Agent.Store.AgentStore
    updateRcvIds,
    createRcvMsg,
    updateRcvMsgHash,
+    createSndMsgBody,
    updateSndIds,
    createSndMsg,
    updateSndMsgHash,
@@ -770,6 +771,14 @@ createRcvMsg db connId rq@RcvQueue {dbQueueId} rcvMsgData@RcvMsgData {msgMeta =
  updateRcvMsgHash db connId sndMsgId internalRcvId internalHash
  DB.execute db "UPDATE rcv_queues SET last_broker_ts = ? WHERE conn_id = ? AND rcv_queue_id = ?" (brokerTs, connId, dbQueueId)

+createSndMsgBody :: DB.Connection -> AMessage -> IO Int64
+createSndMsgBody db aMessage =
+  fromOnly . head <$>
+    DB.query
+      db
+      "INSERT INTO snd_message_bodies (agent_msg) VALUES (?) RETURNING snd_message_body_id"
+      (Only aMessage)
+
 updateSndIds :: DB.Connection -> ConnId -> IO (Either StoreError (InternalId, InternalSndId, PrevSndMsgHash))
 updateSndIds db connId = runExceptT $ do
  (lastInternalId, lastInternalSndId, prevSndHash) <- ExceptT $ retrieveLastIdsAndHashSnd_ db connId
@@ -836,7 +845,7 @@ getPendingQueueMsg db connId SndQueue {dbQueueId} =
          (connId, dbQueueId)
    getMsgData :: InternalId -> IO (Either StoreError (Maybe RcvQueue, PendingMsgData))
    getMsgData msgId = runExceptT $ do
-      msg <- ExceptT $ firstRow pendingMsgData err getMsgData_
+      msg <- ExceptT $ firstRow' pendingMsgData err getMsgData_
      rq_ <- liftIO $ L.head <$$> getRcvQueuesByConnId_ db connId
      pure (rq_, msg)
      where
@@ -844,18 +853,25 @@ getPendingQueueMsg db connId SndQueue {dbQueueId} =
          DB.query
            db
            [sql|
-              SELECT m.msg_type, m.msg_flags, m.msg_body, m.pq_encryption, m.internal_ts, s.retry_int_slow, s.retry_int_fast, s.msg_encrypt_key, s.padded_msg_len
+              SELECT
+                m.msg_type, m.msg_flags, m.msg_body, m.pq_encryption, m.internal_ts, m.internal_snd_id, s.previous_msg_hash,
+                s.retry_int_slow, s.retry_int_fast, s.msg_encrypt_key, s.padded_msg_len, sb.agent_msg
              FROM messages m
              JOIN snd_messages s ON s.conn_id = m.conn_id AND s.internal_id = m.internal_id
+              LEFT JOIN snd_message_bodies sb ON sb.snd_message_body_id = s.snd_message_body_id
              WHERE m.conn_id = ? AND m.internal_id = ?
            |]
            (connId, msgId)
        err = SEInternal $ "msg delivery " <> bshow msgId <> " returned []"
-        pendingMsgData :: (AgentMessageType, Maybe MsgFlags, MsgBody, PQEncryption, InternalTs, Maybe Int64, Maybe Int64, Maybe CR.MsgEncryptKeyX448, Maybe Int) -> PendingMsgData
-        pendingMsgData (msgType, msgFlags_, msgBody, pqEncryption, internalTs, riSlow_, riFast_, encryptKey_, paddedLen_) =
+        pendingMsgData :: (AgentMessageType, Maybe MsgFlags, MsgBody, PQEncryption, InternalTs, InternalSndId, PrevSndMsgHash, Maybe Int64, Maybe Int64, Maybe CR.MsgEncryptKeyX448, Maybe Int, Maybe AMessage) -> Either StoreError PendingMsgData
+        pendingMsgData (msgType, msgFlags_, msgBody, pqEncryption, internalTs, internalSndId, prevMsgHash, riSlow_, riFast_, encryptKey_, paddedLen_, sndMsgBody_) = do
          let msgFlags = fromMaybe SMP.noMsgFlags msgFlags_
              msgRetryState = RI2State <$> riSlow_ <*> riFast_
-           in PendingMsgData {msgId, msgType, msgFlags, msgBody, pqEncryption, msgRetryState, internalTs, encryptKey_, paddedLen_}
+              result pendingMsgPrepData_ = PendingMsgData {msgId, msgType, msgFlags, msgBody, pqEncryption, msgRetryState, internalTs, internalSndId, prevMsgHash, pendingMsgPrepData_}
+           in result <$> case (encryptKey_, paddedLen_, sndMsgBody_) of
+                (Nothing, Nothing, Nothing) -> Right Nothing
+                (Just encryptKey, Just paddedLen, Just sndMsgBody) -> Right $ Just PendingMsgPrepData {encryptKey, paddedLen, sndMsgBody}
+                _ -> Left $ SEInternal "unexpected snd msg data"
    markMsgFailed msgId = DB.execute db "UPDATE snd_message_deliveries SET failed = 1 WHERE conn_id = ? AND internal_id = ?" (connId, msgId)

 getWorkItem :: Show i => ByteString -> IO (Maybe i) -> (i -> IO (Either StoreError a)) -> (i -> IO ()) -> IO (Either StoreError (Maybe a))
@@ -997,7 +1013,6 @@ deleteDeliveredSndMsg db connId msgId = do
  cnt <- countPendingSndDeliveries_ db connId msgId
  when (cnt == 0) $ deleteMsg db connId msgId

-- TODO [save once] Delete from shared message bodies if no deliveries reference it. (`when (cnt == 0)`)
 deleteSndMsgDelivery :: DB.Connection -> ConnId -> SndQueue -> InternalId -> Bool -> IO ()
 deleteSndMsgDelivery db connId SndQueue {dbQueueId} msgId keepForReceipt = do
  DB.execute
@@ -1006,11 +1021,19 @@ deleteSndMsgDelivery db connId SndQueue {dbQueueId} msgId keepForReceipt = do
    (connId, dbQueueId, msgId)
  cnt <- countPendingSndDeliveries_ db connId msgId
  when (cnt == 0) $ do
-    del <-
-      maybeFirstRow id (DB.query db "SELECT rcpt_internal_id, rcpt_status FROM snd_messages WHERE conn_id = ? AND internal_id = ?" (connId, msgId)) >>= \case
-        Just (Just (_ :: Int64), Just MROk) -> pure deleteMsg
-        _ -> pure $ if keepForReceipt then deleteMsgContent else deleteMsg
-    del db connId msgId
+    maybeFirstRow id (DB.query db "SELECT rcpt_internal_id, rcpt_status, snd_message_body_id FROM snd_messages WHERE conn_id = ? AND internal_id = ?" (connId, msgId)) >>= \case
+      Just (Just (_ :: Int64), Just MROk, sndMsgBodyId_) -> do
+        forM_ sndMsgBodyId_ deleteSndMsgBody
+        deleteMsg db connId msgId
+      Just (_, _, Just (sndMsgBodyId :: Int64)) -> do
+        deleteSndMsgBody sndMsgBodyId
+        delKeepForReceipt
+      _ ->
+        delKeepForReceipt
+  where
+    delKeepForReceipt = if keepForReceipt then deleteMsgContent db connId msgId else deleteMsg db connId msgId
+    deleteSndMsgBody sndMsgBodyId =
+      DB.execute db "DELETE FROM snd_message_bodies WHERE snd_message_body_id = ?" (Only sndMsgBodyId)

 countPendingSndDeliveries_ :: DB.Connection -> ConnId -> InternalId -> IO Int
 countPendingSndDeliveries_ db connId msgId = do
@@ -2207,11 +2230,15 @@ insertSndMsgDetails_ dbConn connId SndMsgData {..} =
    dbConn
    [sql|
      INSERT INTO snd_messages
-        ( conn_id, internal_snd_id, internal_id, internal_hash, previous_msg_hash, msg_encrypt_key, padded_msg_len)
+        ( conn_id, internal_snd_id, internal_id, internal_hash, previous_msg_hash, msg_encrypt_key, padded_msg_len, snd_message_body_id)
      VALUES
-        (?,?,?,?,?,?,?)
+        (?,?,?,?,?,?,?,?)
    |]
-    (connId, internalSndId, internalId, Binary internalHash, Binary prevMsgHash, encryptKey_, paddedLen_)
+    (connId, internalSndId, internalId, Binary internalHash, Binary prevMsgHash, encryptKey_, paddedLen_, sndMsgBodyId_)
+  where
+    (encryptKey_, paddedLen_, sndMsgBodyId_) = case sndMsgPrepData_ of
+      Nothing -> (Nothing, Nothing, Nothing)
+      Just SndMsgPrepData {encryptKey, paddedLen, sndMsgBodyId} -> (Just encryptKey, Just paddedLen, Just sndMsgBodyId)

 updateSndMsgHash :: DB.Connection -> ConnId -> InternalSndId -> MsgHash -> IO ()
 updateSndMsgHash db connId internalSndId internalHash =
@@ -25,12 +25,14 @@ import Database.PostgreSQL.Simple.Internal (Connection (..))
 import Database.PostgreSQL.Simple.SqlQQ (sql)
 import Simplex.Messaging.Agent.Store.Postgres.Common
 import Simplex.Messaging.Agent.Store.Postgres.Migrations.M20241210_initial
+import Simplex.Messaging.Agent.Store.Postgres.Migrations.M20250203_msg_bodies
 import Simplex.Messaging.Agent.Store.Shared
 import UnliftIO.MVar

 schemaMigrations :: [(String, Text, Maybe Text)]
 schemaMigrations =
-  [ ("20241210_initial", m20241210_initial, Nothing)
+  [ ("20241210_initial", m20241210_initial, Nothing),
+    ("20250203_msg_bodies", m20250203_msg_bodies, Just down_m20250203_msg_bodies)
  ]

 -- | The list of migrations in ascending order by date
@@ -0,0 +1,37 @@
+{-# LANGUAGE QuasiQuotes #-}
+
+module Simplex.Messaging.Agent.Store.Postgres.Migrations.M20250203_msg_bodies where
+
+import Data.Text (Text)
+import qualified Data.Text as T
+import Text.RawString.QQ (r)
+
+m20250203_msg_bodies :: Text
+m20250203_msg_bodies =
+  T.pack
+    [r|
+ALTER TABLE snd_messages ADD COLUMN msg_encrypt_key BYTEA;
+ALTER TABLE snd_messages ADD COLUMN padded_msg_len BIGINT;
+
+
+CREATE TABLE snd_message_bodies (
+  snd_message_body_id BIGINT PRIMARY KEY GENERATED ALWAYS AS IDENTITY,
+  agent_msg BYTEA NOT NULL DEFAULT ''::BYTEA
+);
+ALTER TABLE snd_messages ADD COLUMN snd_message_body_id BIGINT REFERENCES snd_message_bodies ON DELETE SET NULL;
+CREATE INDEX idx_snd_messages_snd_message_body_id ON snd_messages(snd_message_body_id);
+|]
+
+
+down_m20250203_msg_bodies :: Text
+down_m20250203_msg_bodies =
+  T.pack
+    [r|
+DROP INDEX idx_snd_messages_snd_message_body_id;
+ALTER TABLE snd_messages DROP COLUMN snd_message_body_id;
+DROP TABLE snd_message_bodies;
+
+
+ALTER TABLE snd_messages DROP COLUMN msg_encrypt_key;
+ALTER TABLE snd_messages DROP COLUMN padded_msg_len;
+|]
@@ -12,25 +12,22 @@ ALTER TABLE snd_messages ADD COLUMN msg_encrypt_key BLOB;
 ALTER TABLE snd_messages ADD COLUMN padded_msg_len INTEGER;


-- CREATE TABLE msg_bodies (
--   msg_body_id INTEGER PRIMARY KEY,
--   msg_body BLOB NOT NULL DEFAULT x''
-- )
-
-- ALTER TABLE snd_messages ADD COLUMN msg_body_id INTEGER REFERENCES msg_bodies ON DELETE CASCADE;
-
-- fkey to msg_bodies
-- on each delivery check if other deliveries reference the same msg_body_id, if not delete it
+CREATE TABLE snd_message_bodies (
+  snd_message_body_id INTEGER PRIMARY KEY,
+  agent_msg BLOB NOT NULL DEFAULT x''
+);
+ALTER TABLE snd_messages ADD COLUMN snd_message_body_id INTEGER REFERENCES snd_message_bodies ON DELETE SET NULL;
+CREATE INDEX idx_snd_messages_snd_message_body_id ON snd_messages(snd_message_body_id);
 |]

 down_m20250203_msg_bodies :: Query
 down_m20250203_msg_bodies =
    [sql|
+DROP INDEX idx_snd_messages_snd_message_body_id;
+ALTER TABLE snd_messages DROP COLUMN snd_message_body_id;
+DROP TABLE snd_message_bodies;
+
+
 ALTER TABLE snd_messages DROP COLUMN msg_encrypt_key;
 ALTER TABLE snd_messages DROP COLUMN padded_msg_len;
-
-
-- ALTER TABLE snd_messages DROP COLUMN msg_body_id;
-
-- DROP TABLE msg_bodies;
 |]
@@ -129,6 +129,7 @@ CREATE TABLE snd_messages(
  rcpt_status TEXT,
  msg_encrypt_key BLOB,
  padded_msg_len INTEGER,
+  snd_message_body_id INTEGER REFERENCES snd_message_bodies ON DELETE SET NULL,
  PRIMARY KEY(conn_id, internal_snd_id),
  FOREIGN KEY(conn_id, internal_id) REFERENCES messages
  ON DELETE CASCADE
@@ -417,6 +418,10 @@ CREATE TABLE ntf_tokens_to_delete(
 del_failed INTEGER DEFAULT 0,
 created_at TEXT NOT NULL DEFAULT(datetime('now'))
 );
+CREATE TABLE snd_message_bodies(
+  snd_message_body_id INTEGER PRIMARY KEY,
+  agent_msg BLOB NOT NULL DEFAULT x''
+);
 CREATE UNIQUE INDEX idx_rcv_queues_ntf ON rcv_queues(host, port, ntf_id);
 CREATE UNIQUE INDEX idx_rcv_queue_id ON rcv_queues(conn_id, rcv_queue_id);
 CREATE UNIQUE INDEX idx_snd_queue_id ON snd_queues(conn_id, snd_queue_id);
@@ -543,3 +548,6 @@ CREATE INDEX idx_snd_message_deliveries_expired ON snd_message_deliveries(
  internal_id
 );
 CREATE INDEX idx_rcv_files_redirect_id on rcv_files(redirect_id);
+CREATE INDEX idx_snd_messages_snd_message_body_id ON snd_messages(
+  snd_message_body_id
+);
@@ -1,3 +1,4 @@
+{-# LANGUAGE MonadComprehensions #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE ScopedTypeVariables #-}

@@ -8,6 +9,7 @@ import Control.Monad
 import Control.Monad.Except
 import Control.Monad.IO.Unlift
 import Control.Monad.Trans.Except
+import Control.Monad.Trans.State.Strict (StateT (..))
 import Data.Aeson (FromJSON, ToJSON)
 import qualified Data.Aeson as J
 import Data.Bifunctor (first)
@@ -17,7 +19,7 @@ import qualified Data.ByteString.Lazy.Char8 as LB
 import Data.IORef
 import Data.Int (Int64)
 import Data.List (groupBy, sortOn)
-import Data.List.NonEmpty (NonEmpty)
+import Data.List.NonEmpty (NonEmpty (..))
 import qualified Data.List.NonEmpty as L
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as M
@@ -25,6 +27,7 @@ import Data.Text (Text)
 import qualified Data.Text as T
 import Data.Text.Encoding (decodeUtf8With, encodeUtf8)
 import Data.Time (NominalDiffTime)
+import Data.Tuple (swap)
 import GHC.Conc (labelThread, myThreadId, threadDelay)
 import UnliftIO hiding (atomicModifyIORef')
 import qualified UnliftIO.Exception as UE
@@ -100,6 +103,47 @@ forME :: (Monad m, Traversable t) => t (Either e a) -> (a -> m (Either e b)) ->
 forME = flip mapME
 {-# INLINE forME #-}

+
+-- | Monadic version of mapAccumL
+-- Copied from ghc-9.6.3 package: https://hackage.haskell.org/package/ghc-9.12.1/docs/GHC-Utils-Monad.html#v:mapAccumLM
+-- for backward compatibility with 8.10.7.
+mapAccumLM :: (Monad m, Traversable t)
+            => (acc -> x -> m (acc, y)) -- ^ combining function
+            -> acc                      -- ^ initial state
+            -> t x                      -- ^ inputs
+            -> m (acc, t y)             -- ^ final state, outputs
+{-# INLINE [1] mapAccumLM #-}
+-- INLINE pragma.  mapAccumLM is called in inner loops.  Like 'map',
+-- we inline it so that we can take advantage of knowing 'f'.
+-- This makes a few percent difference (in compiler allocations)
+-- when compiling perf/compiler/T9675
+mapAccumLM f s = fmap swap . flip runStateT s . traverse f'
+  where
+    f' = StateT . (fmap . fmap) swap . flip f
+{-# RULES "mapAccumLM/List" mapAccumLM = mapAccumLM_List #-}
+{-# RULES "mapAccumLM/NonEmpty" mapAccumLM = mapAccumLM_NonEmpty #-}
+
+mapAccumLM_List
+ :: Monad m
+ => (acc -> x -> m (acc, y))
+ -> acc -> [x] -> m (acc, [y])
+{-# INLINE mapAccumLM_List #-}
+mapAccumLM_List f = go
+  where
+    go s (x : xs) = do
+      (s1, x')  <- f s x
+      (s2, xs') <- go s1 xs
+      return    (s2, x' : xs')
+    go s [] = return (s, [])
+
+mapAccumLM_NonEmpty
+ :: Monad m
+ => (acc -> x -> m (acc, y))
+ -> acc -> NonEmpty x -> m (acc, NonEmpty y)
+{-# INLINE mapAccumLM_NonEmpty #-}
+mapAccumLM_NonEmpty f s (x :| xs) =
+  [(s2, x' :| xs') | (s1, x') <- f s x, (s2, xs') <- mapAccumLM_List f s1 xs]
+
 catchAll :: IO a -> (E.SomeException -> IO a) -> IO a
 catchAll = E.catch
 {-# INLINE catchAll #-}