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 <[email protected]>

Author: spaced4ndy

simplexmq.cabal

@@ -152,6 +152,7 @@ library
           Simplex.Messaging.Agent.Store.Postgres.DB
           Simplex.Messaging.Agent.Store.Postgres.Migrations
           Simplex.Messaging.Agent.Store.Postgres.Migrations.M20241210_initial
+          Simplex.Messaging.Agent.Store.Postgres.Migrations.M20250203_msg_bodies
       if !flag(client_library)
           exposed-modules:
               Simplex.Messaging.Agent.Store.Postgres.Util

src/Simplex/Messaging/Agent.hs

@@ -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

src/Simplex/Messaging/Agent/Client.hs

@@ -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)

src/Simplex/Messaging/Agent/Protocol.hs

@@ -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)
 

src/Simplex/Messaging/Agent/Store.hs

@@ -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,26 +845,33 @@ 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
         getMsgData_ =
           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 =

src/Simplex/Messaging/Agent/Store/AgentStore.hs

@@ -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

src/Simplex/Messaging/Agent/Store/Postgres/Migrations.hs

@@ -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;
+|]

src/Simplex/Messaging/Agent/Store/Postgres/Migrations/M20250203_msg_bodies.hs

@@ -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|
-ALTER TABLE snd_messages DROP COLUMN msg_encrypt_key;
-ALTER TABLE snd_messages DROP COLUMN padded_msg_len;
+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_body_id;
-
--- DROP TABLE msg_bodies;
+ALTER TABLE snd_messages DROP COLUMN msg_encrypt_key;
+ALTER TABLE snd_messages DROP COLUMN padded_msg_len;
 |]

src/Simplex/Messaging/Agent/Store/SQLite/Migrations/M20250203_msg_bodies.hs

@@ -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
+);