smp server: refactor message delivery to always respond SOK to subscriptions

2026-10-16 10:19:57 +00:00 · 2025-06-12 23:05:52 +01:00
parent da37384335
commit 4ac4e42a86
6 changed files with 156 additions and 127 deletions
@@ -114,6 +114,7 @@ module Simplex.Messaging.Protocol
    BasicAuth (..),
    SrvLoc (..),
    CorrId (..),
+    pattern NoCorrId,
    EntityId (..),
    pattern NoEntity,
    QueueId,
@@ -1370,6 +1371,9 @@ newtype CorrId = CorrId {bs :: ByteString}
  deriving (Eq, Ord, Show)
  deriving newtype (Encoding)

+pattern NoCorrId :: CorrId
+pattern NoCorrId = CorrId ""
+
 instance IsString CorrId where
  fromString = CorrId . fromString
  {-# INLINE fromString #-}
@@ -66,12 +66,13 @@ import Data.IORef
 import Data.Int (Int64)
 import qualified Data.IntMap.Strict as IM
 import qualified Data.IntSet as IS
-import Data.List (foldl', intercalate, mapAccumR)
+import Data.Foldable (foldrM)
+import Data.List (foldl', intercalate)
 import Data.List.NonEmpty (NonEmpty (..), (<|))
 import qualified Data.List.NonEmpty as L
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as M
-import Data.Maybe (catMaybes, fromMaybe, isJust, isNothing)
+import Data.Maybe (fromMaybe, isJust, isNothing)
 import Data.Semigroup (Sum (..))
 import qualified Data.Set as S
 import Data.Text (Text)
@@ -382,7 +383,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt
                let cancelNtfs s = throwSTM $ userError $ s <> ", " <> show (length ts_) <> " ntfs kept"
                unlessM (currentClient readTVar c) $ cancelNtfs "not current client"
                whenM (isFullTBQueue sndQ) $ cancelNtfs "sending queue full"
-                writeTBQueue sndQ ts
+                writeTBQueue sndQ (ts, [])
              pure $ length ts_
            currentClient :: Monad m => (forall a. TVar a -> m a) -> Client s' -> m Bool
            currentClient rd Client {clientId, connected} = (&&) <$> rd connected <*> (IS.member clientId <$> rd subClients)
@@ -393,7 +394,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt
                ntfs -> do
                  writeTVar v []
                  pure $ foldl' (\acc' ntf -> nmsg nId ntf : acc') acc ntfs -- reverses, to order by time
-            nmsg nId MsgNtf {ntfNonce, ntfEncMeta} = (CorrId "", nId, NMSG ntfNonce ntfEncMeta)
+            nmsg nId MsgNtf {ntfNonce, ntfEncMeta} = (NoCorrId, nId, NMSG ntfNonce ntfEncMeta)
            updateNtfStats :: Client s' -> Either SomeException Int -> IO ()
            updateNtfStats Client {clientId} = \case
              Right 0 -> pure ()
@@ -419,14 +420,14 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt
          where
            enqueueEvts evts c@Client {connected, sndQ} =
              whenM (readTVarIO connected) $ do
-                sent <- atomically $ tryWriteTBQueue sndQ ts
+                sent <- atomically $ tryWriteTBQueue sndQ (ts, [])
                if sent
                  then updateEndStats
                  else -- if queue is full it can block
                    forkClient c ("sendPendingEvtsThread.queueEvts") $
-                      atomically (writeTBQueue sndQ ts) >> updateEndStats
+                      atomically (writeTBQueue sndQ (ts, [])) >> updateEndStats
              where
-                ts = L.map (\(entId, evt) -> (CorrId "", entId, evt)) evts
+                ts = L.map (\(entId, evt) -> (NoCorrId, entId, evt)) evts
                -- this accounts for both END and DELD events
                updateEndStats = do
                  let len = L.length evts
@@ -1082,7 +1083,7 @@ receive h@THandle {params = THandleParams {thAuth, sessionId}} ms Client {rcvQ,
    atomically . (writeTVar rcvActiveAt $!) =<< getSystemTime
    let (es, ts') = partitionEithers $ L.toList ts
        errs = map (second ERR) es
-    case ts' of
+    errs' <- case ts' of
      (_, _, (_, _, Cmd p cmd)) : rest -> do
        let service = peerClientService =<< thAuth
        (errs', cmds) <- partitionEithers <$> case batchParty p of
@@ -1092,9 +1093,10 @@ receive h@THandle {params = THandleParams {thAuth, sessionId}} ms Client {rcvQ,
            qs <- getQueueRecs ms p $ map queueId ts'
            zipWithM (\t -> verified stats t . verifyLoadedQueue service thAuth t) ts' qs
          _ -> mapM (\t -> verified stats t =<< verifyTransmission ms service thAuth t) ts'
-        write rcvQ cmds
-        write sndQ $ errs ++ errs'
-      [] -> write sndQ errs
+        mapM_ (atomically . writeTBQueue rcvQ) $ L.nonEmpty cmds
+        pure $ errs ++ errs'
+      [] -> pure errs
+    mapM_ (atomically . writeTBQueue sndQ . (,[])) $ L.nonEmpty errs'
  where
    sameParty :: SParty p -> SignedTransmission Cmd -> Bool
    sameParty p (_, _, (_, _, Cmd p' _)) = isJust $ testEquality p p'
@@ -1116,37 +1118,32 @@ receive h@THandle {params = THandleParams {thAuth, sessionId}} ms Client {rcvQ,
            NSUB -> incStat $ ntfSubAuth stats
            GET -> incStat $ msgGetAuth stats
            _ -> pure ()
-    write q = mapM_ (atomically . writeTBQueue q) . L.nonEmpty

 send :: Transport c => MVar (THandleSMP c 'TServer) -> Client s -> IO ()
 send th c@Client {sndQ, msgQ, clientTHParams = THandleParams {sessionId}} = do
  labelMyThread . B.unpack $ "client $" <> encode sessionId <> " send"
  forever $ atomically (readTBQueue sndQ) >>= sendTransmissions
  where
-    sendTransmissions :: NonEmpty (Transmission BrokerMsg) -> IO ()
-    sendTransmissions ts
-      | L.length ts <= 2 = tSend th c ts
+    -- If the request had batched subscriptions
+    -- this will reply SOKs to all SUBs in the first batched transmission,
+    -- to reduce client timeouts.
+    -- After that all messages will be sent in separate transmissions,
+    -- without any client response timeouts, and allowing them to interleave
+    -- with other requests responses.
+    sendTransmissions :: (NonEmpty (Transmission BrokerMsg), [(RecipientId, RcvMessage)]) -> IO ()
+    sendTransmissions (ts, []) = tSend th c ts
+    sendTransmissions (ts, msg@(rId, encMsg) : msgs)
+      | length ts <= 4 = do -- up to 4 SOKs can be in one block with MSG (see testBatchSubResponses test)
+          tSend th c $ ts <> [(NoCorrId, rId, MSG encMsg)]
+          mapM_ (atomically . writeTBQueue msgQ) $ L.nonEmpty msgs
      | otherwise = do
-          let (msgs_, ts') = mapAccumR splitMessages [] ts
-          -- If the request had batched subscriptions and L.length ts > 2
-          -- this will reply OK to all SUBs in the first batched transmission,
-          -- to reduce client timeouts.
-          tSend th c ts'
-          -- After that all messages will be sent in separate transmissions,
-          -- without any client response timeouts, and allowing them to interleave
-          -- with other requests responses.
-          mapM_ (atomically . writeTBQueue msgQ) $ L.nonEmpty msgs_
-      where
-        splitMessages :: [Transmission BrokerMsg] -> Transmission BrokerMsg -> ([Transmission BrokerMsg], Transmission BrokerMsg)
-        splitMessages msgs t@(corrId, entId, cmd) = case cmd of
-          -- replace MSG response with OK, accumulating MSG in a separate list.
-          MSG {} -> ((CorrId "", entId, cmd) : msgs, (corrId, entId, OK))
-          _ -> (msgs, t)
+          tSend th c ts
+          atomically $ writeTBQueue msgQ (msg :| msgs)

 sendMsg :: Transport c => MVar (THandleSMP c 'TServer) -> Client s -> IO ()
 sendMsg th c@Client {msgQ, clientTHParams = THandleParams {sessionId}} = do
  labelMyThread . B.unpack $ "client $" <> encode sessionId <> " sendMsg"
-  forever $ atomically (readTBQueue msgQ) >>= mapM_ (\t -> tSend th c [t])
+  forever $ atomically (readTBQueue msgQ) >>= mapM_ (\(rId, msg) -> tSend th c [(NoCorrId, rId, MSG msg)])

 tSend :: Transport c => MVar (THandleSMP c 'TServer) -> Client s -> NonEmpty (Transmission BrokerMsg) -> IO ()
 tSend th Client {sndActiveAt} ts = do
@@ -1312,15 +1309,16 @@ client
    labelMyThread . B.unpack $ "client $" <> encode sessionId <> " commands"
    let THandleParams {thVersion} = thParams'
        service = peerClientService =<< thAuth thParams'
+        process t acc@(rs, msgs) =
+          (maybe acc (\(!r, !msg_) -> (r : rs, maybe msgs (: msgs) msg_)))
+            <$> processCommand service thVersion t
    forever $
      atomically (readTBQueue rcvQ)
-        >>= mapM (processCommand service thVersion)
-        >>= mapM_ reply . L.nonEmpty . catMaybes . L.toList
+        >>= foldrM process ([], [])
+        >>= \(rs_, msgs) -> mapM_ (atomically . writeTBQueue sndQ . (,msgs)) (L.nonEmpty rs_)
  where
-    reply :: MonadIO m => NonEmpty (Transmission BrokerMsg) -> m ()
-    reply = atomically . writeTBQueue sndQ
-    processProxiedCmd :: Transmission (Command 'ProxiedClient) -> M s (Maybe (Transmission BrokerMsg))
-    processProxiedCmd (corrId, EntityId sessId, command) = (corrId,EntityId sessId,) <$$> case command of
+    processProxiedCmd :: Transmission (Command 'ProxiedClient) -> M s (Maybe ResponseAndMessage)
+    processProxiedCmd (corrId, EntityId sessId, command) = (\t -> ((corrId, EntityId sessId, t), Nothing)) <$$> case command of
      PRXY srv auth -> ifM allowProxy getRelay (pure $ Just $ ERR $ PROXY BASIC_AUTH)
        where
          allowProxy = do
@@ -1383,7 +1381,7 @@ client
        forkProxiedCmd cmdAction = do
          bracket_ wait signal . forkClient clnt (B.unpack $ "client $" <> encode sessionId <> " proxy") $ do
            -- commands MUST be processed under a reasonable timeout or the client would halt
-            cmdAction >>= \t -> reply [(corrId, EntityId sessId, t)]
+            cmdAction >>= \t -> atomically $ writeTBQueue sndQ ([(corrId, EntityId sessId, t)], [])
          pure Nothing
          where
            wait = do
@@ -1399,32 +1397,32 @@ client
    mkIncProxyStats ps psOwn own sel = do
      incStat $ sel ps
      when own $ incStat $ sel psOwn
-    processCommand :: Maybe THPeerClientService -> VersionSMP -> VerifiedTransmission s -> M s (Maybe (Transmission BrokerMsg))
+    processCommand :: Maybe THPeerClientService -> VersionSMP -> VerifiedTransmission s -> M s (Maybe ResponseAndMessage)
    processCommand service clntVersion (q_, (corrId, entId, cmd)) = case cmd of
      Cmd SProxiedClient command -> processProxiedCmd (corrId, entId, command)
-      Cmd SSender command -> Just <$> case command of
+      Cmd SSender command -> response <$> case command of
        SKEY k -> withQueue $ \q qr -> checkMode QMMessaging qr $ secureQueue_ q k
-        SEND flags msgBody -> withQueue_ False $ sendMessage flags msgBody
-      Cmd SIdleClient PING -> pure $ Just (corrId, NoEntity, PONG)
-      Cmd SProxyService (RFWD encBlock) -> Just . (corrId,NoEntity,) <$> processForwardedCommand encBlock
-      Cmd SSenderLink command -> Just <$> case command of
+        SEND flags msgBody -> withQueue_ False err $ sendMessage flags msgBody
+      Cmd SIdleClient PING -> pure $ response (corrId, NoEntity, PONG)
+      Cmd SProxyService (RFWD encBlock) -> (response . (corrId,NoEntity,)) <$> processForwardedCommand encBlock
+      Cmd SSenderLink command -> response <$> case command of
        LKEY k -> withQueue $ \q qr -> checkMode QMMessaging qr $ secureQueue_ q k $>> getQueueLink_ q qr
        LGET -> withQueue $ \q qr -> checkContact qr $ getQueueLink_ q qr
-      Cmd SNotifier NSUB -> Just . (corrId,entId,) <$> case q_ of
+      Cmd SNotifier NSUB -> response . (corrId,entId,) <$> case q_ of
        Just (q, QueueRec {notifier = Just ntfCreds}) -> subscribeNotifications q ntfCreds
        _ -> pure $ ERR INTERNAL
-      Cmd SNotifierService NSUBS -> Just . (corrId,entId,) <$> case service of
+      Cmd SNotifierService NSUBS -> response . (corrId,entId,) <$> case service of
        Just s -> subscribeServiceNotifications s
        Nothing -> pure $ ERR INTERNAL
      Cmd SCreator (NEW nqr@NewQueueReq {auth_}) ->
-        Just <$> ifM allowNew (createQueue nqr) (pure (corrId, entId, ERR AUTH))
+        response <$> ifM allowNew (createQueue nqr) (pure (corrId, entId, ERR AUTH))
          where
            allowNew = do
              ServerConfig {allowNewQueues, newQueueBasicAuth} <- asks config
              pure $ allowNewQueues && maybe True ((== auth_) . Just) newQueueBasicAuth
+      Cmd SRecipient SUB -> Just <$> withQueue' subscribeQueueAndDeliver
      Cmd SRecipient command ->
-        Just <$> case command of
-          SUB -> withQueue subscribeQueue
+        response <$> case command of
          GET -> withQueue getMessage
          ACK msgId -> withQueue $ acknowledgeMsg msgId
          KEY sKey -> withQueue $ \q _ -> either err (corrId,entId,) <$> secureQueue_ q sKey
@@ -1443,12 +1441,12 @@ client
          OFF -> maybe (pure $ err INTERNAL) suspendQueue_ q_
          DEL -> maybe (pure $ err INTERNAL) delQueueAndMsgs q_
          QUE -> withQueue $ \q qr -> (corrId,entId,) <$> getQueueInfo q qr
-      Cmd SRecipientService SUBS -> pure $ Just $ err (CMD PROHIBITED) -- "TODO [certs rcv]"
+      Cmd SRecipientService SUBS -> pure $ response $ err (CMD PROHIBITED) -- "TODO [certs rcv]"
      where
        createQueue :: NewQueueReq -> M s (Transmission BrokerMsg)
        createQueue NewQueueReq {rcvAuthKey, rcvDhKey, subMode, queueReqData}
          | isJust service && subMode == SMOnlyCreate = pure (corrId, entId, ERR $ CMD PROHIBITED)
-          | otherwise = time "NEW" $ do
+          | otherwise = do
              g <- asks random
              idSize <- asks $ queueIdBytes . config
              updatedAt <- Just <$> liftIO getSystemDate
@@ -1506,7 +1504,7 @@ client
                            -- when (isJust ntf) $ incStat $ ntfCreated stats
                            case subMode of
                              SMOnlyCreate -> pure ()
-                              SMSubscribe -> void $ subscribeNewQueue rcvId qr -- no need to check if message is available, it's a new queue
+                              SMSubscribe -> void $ subscribeQueue rcvId qr -- no need to check if message is available, it's a new queue
                            pure $ IDS QIK {rcvId, sndId, rcvPublicDhKey, queueMode, linkId = fst <$> queueData, serviceId = rcvServiceId} -- , serverNtfCreds = snd <$> ntf
              (corrId,entId,) <$> tryCreate (3 :: Int)

@@ -1531,7 +1529,7 @@ client
        getQueueLink_ q qr = liftIO $ LNK (senderId qr) <$$> getQueueLinkData (queueStore ms) q entId

        addQueueNotifier_ :: StoreQueue s -> NtfPublicAuthKey -> RcvNtfPublicDhKey -> M s (Transmission BrokerMsg)
-        addQueueNotifier_ q notifierKey dhKey = time "NKEY" $ do
+        addQueueNotifier_ q notifierKey dhKey = do
          (rcvPublicDhKey, privDhKey) <- atomically . C.generateKeyPair =<< asks random
          let rcvNtfDhSecret = C.dh' dhKey privDhKey
          (corrId,entId,) <$> addNotifierRetry 3 rcvPublicDhKey rcvNtfDhSecret
@@ -1568,32 +1566,32 @@ client
        suspendQueue_ (q, _) = liftIO $ either err (const ok) <$> suspendQueue (queueStore ms) q

        -- TODO [certs rcv] if serviceId is passed, associate with the service and respond with SOK
-        subscribeQueue :: StoreQueue s -> QueueRec -> M s (Transmission BrokerMsg)
-        subscribeQueue q qr =
+        subscribeQueueAndDeliver :: StoreQueue s -> QueueRec -> M s ResponseAndMessage
+        subscribeQueueAndDeliver q qr =
          liftIO (TM.lookupIO rId subscriptions) >>= \case
-            Nothing -> subscribeNewQueue rId qr >>= deliver True
+            Nothing -> subscribeQueue rId qr >>= deliver True
            Just s@Sub {subThread} -> do
              stats <- asks serverStats
              case subThread of
                ProhibitSub -> do
                  -- cannot use SUB in the same connection where GET was used
                  incStat $ qSubProhibited stats
-                  pure (corrId, rId, ERR $ CMD PROHIBITED)
+                  pure ((corrId, rId, ERR $ CMD PROHIBITED), Nothing)
                _ -> do
                  incStat $ qSubDuplicate stats
                  atomically (tryTakeTMVar $ delivered s) >> deliver False s
          where
            rId = recipientId q
-            deliver :: Bool -> Sub -> M s (Transmission BrokerMsg)
+            deliver :: Bool -> Sub -> M s ResponseAndMessage
            deliver inc sub = do
              stats <- asks serverStats
-              fmap (either (\e -> (corrId, rId, ERR e)) id) $ liftIO $ runExceptT $ do
+              fmap (either (\e -> ((corrId, rId, ERR e), Nothing)) id) $ liftIO $ runExceptT $ do
                msg_ <- tryPeekMsg ms q
-                liftIO $ when (inc && isJust msg_) $ incStat (qSub stats)
-                liftIO $ deliverMessage "SUB" qr rId sub msg_
+                mapM_ (\msg -> atomically (setDelivered sub msg) >> when inc (incStat $ qSub stats)) msg_
+                pure ((corrId, rId, SOK Nothing), (rId,) . encryptMsg qr <$> msg_)

-        subscribeNewQueue :: RecipientId -> QueueRec -> M s Sub
-        subscribeNewQueue rId QueueRec {rcvServiceId} = time "SUB newSub" . atomically $ do
+        subscribeQueue :: RecipientId -> QueueRec -> M s Sub
+        subscribeQueue rId QueueRec {rcvServiceId} = atomically $ do
          writeTQueue (subQ subscribers) (CSClient rId rcvServiceId Nothing, clientId)
          sub <- newSubscription NoSub
          TM.insert rId sub subscriptions
@@ -1601,7 +1599,7 @@ client

        -- clients that use GET are not added to server subscribers
        getMessage :: StoreQueue s -> QueueRec -> M s (Transmission BrokerMsg)
-        getMessage q qr = time "GET" $ do
+        getMessage q qr = do
          atomically (TM.lookup entId subscriptions) >>= \case
            Nothing ->
              atomically newSub >>= (`getMessage_` Nothing)
@@ -1637,19 +1635,22 @@ client
                  Nothing -> incStat (msgGetNoMsg stats) $> ok

        withQueue :: (StoreQueue s -> QueueRec -> M s (Transmission BrokerMsg)) -> M s (Transmission BrokerMsg)
-        withQueue = withQueue_ True
+        withQueue = withQueue_ True err
+
+        withQueue' :: (StoreQueue s -> QueueRec -> M s ResponseAndMessage) -> M s ResponseAndMessage
+        withQueue' = withQueue_ True ((,Nothing) . err)

        -- SEND passes queueNotBlocked False here to update time, but it fails anyway on blocked queues (see code for SEND).
-        withQueue_ :: Bool -> (StoreQueue s -> QueueRec -> M s (Transmission BrokerMsg)) -> M s (Transmission BrokerMsg)
-        withQueue_ queueNotBlocked action = case q_ of
-          Nothing -> pure $ err INTERNAL
+        withQueue_ :: Bool -> (ErrorType -> r) -> (StoreQueue s -> QueueRec -> M s r) -> M s r
+        withQueue_ queueNotBlocked err' action = case q_ of
+          Nothing -> pure $ err' INTERNAL
          Just (q, qr@QueueRec {status, updatedAt}) -> case status of
-            EntityBlocked info | queueNotBlocked -> pure $ err $ BLOCKED info
+            EntityBlocked info | queueNotBlocked -> pure $ err' $ BLOCKED info
            _ -> do
              t <- liftIO getSystemDate
              if updatedAt == Just t
                then action q qr
-                else liftIO (updateQueueTime (queueStore ms) q t) >>= either (pure . err) (action q)
+                else liftIO (updateQueueTime (queueStore ms) q t) >>= either (pure . err') (action q)

        subscribeNotifications :: StoreQueue s -> NtfCreds -> M s BrokerMsg
        subscribeNotifications q NtfCreds {ntfServiceId} = do
@@ -1721,7 +1722,7 @@ client
            else pure $ SOKS srvSubs

        acknowledgeMsg :: MsgId -> StoreQueue s -> QueueRec -> M s (Transmission BrokerMsg)
-        acknowledgeMsg msgId q qr = time "ACK" $ do
+        acknowledgeMsg msgId q qr = do
          liftIO (TM.lookupIO entId subscriptions) >>= \case
            Nothing -> pure $ err NO_MSG
            Just sub ->
@@ -1736,8 +1737,10 @@ client
                        pure ok
                      _ -> do
                        (deletedMsg_, msg_) <- tryDelPeekMsg ms q msgId
-                        liftIO $ mapM_ (updateStats stats False) deletedMsg_
-                        liftIO $ deliverMessage "ACK" qr entId sub msg_
+                        liftIO $ do
+                          mapM_ (updateStats stats False) deletedMsg_
+                          mapM_ (atomically . setDelivered sub) msg_
+                          pure (corrId, entId, maybe OK (MSG . encryptMsg qr) msg_)
                _ -> pure $ err NO_MSG
          where
            getDelivered :: Sub -> STM (Maybe ServerSub)
@@ -1788,7 +1791,7 @@ client
                        deleteQueueLinkData (queueStore ms) q
                      ServerConfig {messageExpiration, msgIdBytes} <- asks config
                      msgId <- randomId' msgIdBytes
-                      msg_ <- liftIO $ time "SEND" $ runExceptT $ do
+                      msg_ <- liftIO $ runExceptT $ do
                        expireMessages messageExpiration stats
                        msg <- liftIO $ mkMessage msgId body
                        writeMsg ms q True msg
@@ -1797,7 +1800,7 @@ client
                        Right Nothing -> do
                          incStat $ msgSentQuota stats
                          pure $ err QUOTA
-                        Right (Just (msg, wasEmpty)) -> time "SEND ok" $ do
+                        Right (Just (msg, wasEmpty)) -> do
                          when wasEmpty $ liftIO $ tryDeliverMessage msg
                          when (notification msgFlags) $ do
                            mapM_ (`enqueueNotification` msg) (notifier qr)
@@ -1856,7 +1859,7 @@ client
                        _ -> pure Nothing
                deliver sndQ' s = do
                  let encMsg = encryptMsg qr msg
-                  writeTBQueue sndQ' [(CorrId "", rId, MSG encMsg)]
+                  writeTBQueue sndQ' ([(NoCorrId, rId, MSG encMsg)], [])
                  void $ setDelivered s msg
                forkDeliver (rc@Client {sndQ = sndQ'}, s@Sub {delivered}, st) = do
                  t <- mkWeakThreadId =<< forkIO deliverThread
@@ -1870,7 +1873,7 @@ client
                      -- lookup can be outside of STM transaction,
                      -- as long as the check that it is the same client is inside.
                      getSubscribedClient rId (queueSubscribers subscribers) >>= mapM_ deliverIfSame
-                    deliverIfSame rcv = time "deliver" . atomically $
+                    deliverIfSame rcv = atomically $
                      whenM (sameClient rc rcv) $
                        tryTakeTMVar delivered >>= \case
                          Just _ -> pure () -- if a message was already delivered, should not deliver more
@@ -1916,7 +1919,8 @@ client
              Left r -> pure r
              -- rejectOrVerify filters allowed commands, no need to repeat it here.
              -- INTERNAL is used because processCommand never returns Nothing for sender commands (could be extracted for better types).
-              Right t''@(_, (corrId', entId', _)) -> fromMaybe (corrId', entId', ERR INTERNAL) <$> lift (processCommand Nothing fwdVersion t'')
+              -- `fst` removes empty message that is only returned for `SUB` command
+              Right t''@(_, (corrId', entId', _)) -> maybe (corrId', entId', ERR INTERNAL) fst <$> lift (processCommand Nothing fwdVersion t'')
          -- encode response
          r' <- case batchTransmissions clntTHParams [Right (Nothing, encodeTransmission clntTHParams r)] of
            [] -> throwE INTERNAL -- at least 1 item is guaranteed from NonEmpty/Right
@@ -1949,21 +1953,6 @@ client
                    VRVerified q -> Right (q, t'')
                    VRFailed e -> Left (corrId', entId', ERR e)

-        deliverMessage :: T.Text -> QueueRec -> RecipientId -> Sub -> Maybe Message -> IO (Transmission BrokerMsg)
-        deliverMessage name qr rId s@Sub {subThread} msg_ = time (name <> " deliver") . atomically $
-          case subThread of
-            ProhibitSub -> pure resp
-            _ -> case msg_ of
-              Just msg ->
-                let encMsg = encryptMsg qr msg
-                 in setDelivered s msg $> (corrId, rId, MSG encMsg)
-              _ -> pure resp
-          where
-            resp = (corrId, rId, OK)
-
-        time :: MonadIO m => T.Text -> m a -> m a
-        time name = timed name entId
-
        encryptMsg :: QueueRec -> Message -> RcvMessage
        encryptMsg qr msg = encrypt . encodeRcvMsgBody $ case msg of
          Message {msgFlags, msgBody} -> RcvMsgBody {msgTs = msgTs', msgFlags, msgBody}
@@ -2021,9 +2010,15 @@ client

        ok :: Transmission BrokerMsg
        ok = (corrId, entId, OK)
+        {-# INLINE ok #-}

        err :: ErrorType -> Transmission BrokerMsg
        err e = (corrId, entId, ERR e)
+        {-# INLINE err #-}
+
+        response :: Transmission BrokerMsg -> Maybe ResponseAndMessage
+        response = Just . (,Nothing)
+        {-# INLINE response #-}

 updateDeletedStats :: QueueRec -> M s ()
 updateDeletedStats q = do
@@ -2037,18 +2032,6 @@ incStat :: MonadIO m => IORef Int -> m ()
 incStat r = liftIO $ atomicModifyIORef'_ r (+ 1)
 {-# INLINE incStat #-}

-timed :: MonadIO m => T.Text -> RecipientId -> m a -> m a
-timed name (EntityId qId) a = do
-  t <- liftIO getSystemTime
-  r <- a
-  t' <- liftIO getSystemTime
-  let int = diff t t'
-  when (int > sec) . logDebug $ T.unwords [name, tshow $ encode qId, tshow int]
-  pure r
-  where
-    diff t t' = (systemSeconds t' - systemSeconds t) * sec + fromIntegral (systemNanoseconds t' - systemNanoseconds t)
-    sec = 1000_000000
-
 randomId' :: Int -> M s ByteString
 randomId' n = atomically . C.randomBytes n =<< asks random

@@ -40,6 +40,7 @@ module Simplex.Messaging.Server.Env.STM
    MsgStore (..),
    AStoreType (..),
    VerifiedTransmission,
+    ResponseAndMessage,
    newEnv,
    mkJournalStoreConfig,
    msgStore,
@@ -392,8 +393,8 @@ data Client s = Client
    serviceSubsCount :: TVar Int64, -- only one service can be subscribed, based on its certificate, this is subscription count
    ntfServiceSubsCount :: TVar Int64, -- only one service can be subscribed, based on its certificate, this is subscription count
    rcvQ :: TBQueue (NonEmpty (VerifiedTransmission s)),
-    sndQ :: TBQueue (NonEmpty (Transmission BrokerMsg)),
-    msgQ :: TBQueue (NonEmpty (Transmission BrokerMsg)),
+    sndQ :: TBQueue (NonEmpty (Transmission BrokerMsg), [(RecipientId, RcvMessage)]),
+    msgQ :: TBQueue (NonEmpty (RecipientId, RcvMessage)),
    procThreads :: TVar Int,
    endThreads :: TVar (IntMap (Weak ThreadId)),
    endThreadSeq :: TVar Int,
@@ -406,6 +407,8 @@ data Client s = Client

 type VerifiedTransmission s = (Maybe (StoreQueue s, QueueRec), Transmission Cmd)

+type ResponseAndMessage = (Transmission BrokerMsg, Maybe (RecipientId, RcvMessage))
+
 data ServerSub = ServerSub (TVar SubscriptionThread) | ProhibitSub

 data SubscriptionThread = NoSub | SubPending | SubThread (Weak ThreadId)
@@ -43,14 +43,13 @@ import qualified Data.ByteString.Builder as BB
 import Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Lazy as LB
 import Data.Bitraversable (bimapM)
-import Data.Either (fromRight, lefts, rights)
+import Data.Either (fromRight, lefts)
 import Data.Functor (($>))
 import Data.Int (Int64)
 import Data.List (foldl', intersperse, partition)
 import Data.List.NonEmpty (NonEmpty)
-import qualified Data.List.NonEmpty as L
 import qualified Data.Map.Strict as M
-import Data.Maybe (catMaybes, fromMaybe, mapMaybe)
+import Data.Maybe (catMaybes, fromMaybe)
 import qualified Data.Set as S
 import Data.Text (Text)
 import Data.Time.Clock.System (SystemTime (..), getSystemTime)
@@ -64,7 +63,7 @@ import Database.PostgreSQL.Simple.ToField (Action (..), ToField (..))
 import Database.PostgreSQL.Simple.Errors (ConstraintViolation (..), constraintViolation)
 import Database.PostgreSQL.Simple.SqlQQ (sql)
 import GHC.IO (catchAny)
-import Simplex.Messaging.Agent.Client (withLockMap, withLocksMap)
+import Simplex.Messaging.Agent.Client (withLockMap)
 import Simplex.Messaging.Agent.Lock (Lock)
 import Simplex.Messaging.Agent.Store.AgentStore ()
 import Simplex.Messaging.Agent.Store.Postgres (createDBStore, closeDBStore)
@@ -83,7 +82,7 @@ import Simplex.Messaging.Server.StoreLog
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
 import Simplex.Messaging.Transport (SMPServiceRole (..))
-import Simplex.Messaging.Util (eitherToMaybe, firstRow, ifM, maybeFirstRow, tshow, (<$$>), ($>>=))
+import Simplex.Messaging.Util (eitherToMaybe, firstRow, ifM, maybeFirstRow, tshow, (<$$>))
 import System.Exit (exitFailure)
 import System.IO (IOMode (..), hFlush, stdout)
 import UnliftIO.STM
@@ -47,6 +47,7 @@ batchingTests = do
      it "should batch with 135 subscriptions per batch" testClientBatchSubscriptions
      it "should batch with 255 ENDs per batch" testClientBatchENDs
      it "should batch with 80 NMSGs per batch" testClientBatchNMSGs
+      it "should batch subscription responses with message" testBatchSubResponses
      it "should break on message that does not fit" testClientBatchWithMessage
      it "should break on large message" testClientBatchWithLargeMessage

@@ -207,6 +208,20 @@ testClientBatchNMSGs = do
  (length rs1, length rs2, length rs3) `shouldBe` (40, 80, 80)
  all lenOk [s1, s2, s3] `shouldBe` True

+-- 4 responses are used in Simplex.Messaging.Server / `send`
+testBatchSubResponses :: IO ()
+testBatchSubResponses = do
+  client <- testClientStub
+  soks <- replicateM 4 $ randomSOK
+  msg <- randomMSG
+  let msgs = map (\t -> Right (Nothing, encodeTransmission (thParams client) t)) (soks <> [msg])
+      batches = batchTransmissions (thParams client) $ L.fromList msgs
+  length batches `shouldBe` 1
+  soks' <- replicateM 5 $ randomSOK
+  let msgs' = map (\t -> Right (Nothing, encodeTransmission (thParams client) t)) (soks' <> [msg])
+      batches' = batchTransmissions (thParams client) $ L.fromList msgs'
+  length batches' `shouldBe` 2
+
 testClientBatchWithMessageV6 :: IO ()
 testClientBatchWithMessageV6 = do
  client <- testClientStubV6
@@ -361,6 +376,22 @@ randomNMSGCmd ts = do
  Right encNMsgMeta <- pure $ C.cbEncrypt (C.dh' k pk) nonce (smpEncode msgMeta) 128
  pure (CorrId "", EntityId nId, NMSG nonce encNMsgMeta)

+randomSOK :: IO (Transmission BrokerMsg)
+randomSOK = do
+  g <- C.newRandom
+  corrId <- atomically $ C.randomBytes 24 g
+  rId <- atomically $ C.randomBytes 24 g
+  pure (CorrId corrId, EntityId rId, SOK Nothing)
+
+randomMSG :: IO (Transmission BrokerMsg)
+randomMSG = do
+  g <- C.newRandom
+  corrId <- atomically $ C.randomBytes 24 g
+  rId <- atomically $ C.randomBytes 24 g
+  msgId <- atomically $ C.randomBytes 24 g
+  msg <- atomically $ C.randomBytes (maxMessageLength currentClientSMPRelayVersion) g
+  pure (CorrId corrId, EntityId rId, MSG RcvMessage {msgId, msgBody = EncRcvMsgBody msg})
+
 randomSENDv6 :: ByteString -> Int -> IO (Either TransportError (Maybe TAuthorizations, ByteString))
 randomSENDv6 = randomSEND_ C.SEd25519 minServerSMPRelayVersion

@@ -29,6 +29,7 @@ import Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Char8 as B
 import Data.Hashable (hash)
 import qualified Data.IntSet as IS
+import Data.List.NonEmpty (NonEmpty)
 import Data.String (IsString (..))
 import Data.Type.Equality
 import qualified Data.X509.Validation as XV
@@ -111,16 +112,25 @@ sendRecv h@THandle {params} (sgn, corrId, qId, cmd) = do
  tGet1 h

 signSendRecv :: forall c p. (Transport c, PartyI p) => THandleSMP c 'TClient -> C.APrivateAuthKey -> (ByteString, EntityId, Command p) -> IO (Transmission (Either ErrorType BrokerMsg))
-signSendRecv h pk = signSendRecv_ h pk Nothing
+signSendRecv h pk t = do
+  [r] <- signSendRecv_ h pk Nothing t
+  pure r
+
+signSendRecv2 :: forall c p. (Transport c, PartyI p) => THandleSMP c 'TClient -> C.APrivateAuthKey -> (ByteString, EntityId, Command p) -> IO (Transmission (Either ErrorType BrokerMsg), Transmission (Either ErrorType BrokerMsg))
+signSendRecv2 h pk t = do
+  [r1, r2] <- signSendRecv_ h pk Nothing t
+  pure (r1, r2)

 serviceSignSendRecv :: forall c p. (Transport c, PartyI p) => THandleSMP c 'TClient -> C.APrivateAuthKey -> C.PrivateKeyEd25519 -> (ByteString, EntityId, Command p) -> IO (Transmission (Either ErrorType BrokerMsg))
-serviceSignSendRecv h pk = signSendRecv_ h pk . Just
+serviceSignSendRecv h pk serviceKey t = do
+  [r] <- signSendRecv_ h pk (Just serviceKey) t
+  pure r

-signSendRecv_ :: forall c p. (Transport c, PartyI p) => THandleSMP c 'TClient -> C.APrivateAuthKey -> Maybe C.PrivateKeyEd25519 -> (ByteString, EntityId, Command p) -> IO (Transmission (Either ErrorType BrokerMsg))
+signSendRecv_ :: forall c p. (Transport c, PartyI p) => THandleSMP c 'TClient -> C.APrivateAuthKey -> Maybe C.PrivateKeyEd25519 -> (ByteString, EntityId, Command p) -> IO (NonEmpty (Transmission (Either ErrorType BrokerMsg)))
 signSendRecv_ h@THandle {params} (C.APrivateAuthKey a pk) serviceKey_ (corrId, qId, cmd) = do
  let TransmissionForAuth {tForAuth, tToSend} = encodeTransmissionForAuth params (CorrId corrId, qId, cmd)
  Right () <- tPut1 h (authorize tForAuth, tToSend)
-  tGet1 h
+  liftIO $ tGetClient h
  where
    authorize t = (,(`C.sign'` t) <$> serviceKey_) <$> case a of
      C.SEd25519 -> Just . TASignature . C.ASignature C.SEd25519 $ C.sign' pk t'
@@ -365,7 +375,7 @@ testCreateDelete =
      Resp "bcda" _ ok4 <- signSendRecv rh rKey ("bcda", rId, OFF)
      (ok4, OK) #== "accepts OFF when suspended"

-      Resp "cdab" _ (Msg mId2 msg2) <- signSendRecv rh rKey ("cdab", rId, SUB)
+      (Resp "cdab" _ (SOK Nothing), Resp "" _ (Msg mId2 msg2)) <- signSendRecv2 rh rKey ("cdab", rId, SUB)
      (dec mId2 msg2, Right "hello") #== "accepts SUB when suspended and delivers the message again (because was not ACKed)"

      Resp "dabc" _ err5 <- sendRecv rh (sampleSig, "dabc", rId, DEL)
@@ -404,7 +414,7 @@ stressTest =
        Resp "" NoEntity (Ids rId _ _) <- signSendRecv h1 rKey ("", NoEntity, New rPub dhPub)
        pure rId
      let subscribeQueues h = forM_ rIds $ \rId -> do
-            Resp "" rId' OK <- signSendRecv h rKey ("", rId, SUB)
+            Resp "" rId' (SOK Nothing) <- signSendRecv h rKey ("", rId, SUB)
            rId' `shouldBe` rId
      closeConnection $ connection h1
      subscribeQueues h2
@@ -497,7 +507,7 @@ testSwitchSub =
      Resp "abcd" _ (Msg mId2 msg2) <- signSendRecv rh1 rKey ("abcd", rId, ACK mId1)
      (dec mId2 msg2, Right "test2, no ACK") #== "test message 2 delivered, no ACK"

-      Resp "bcda" _ (Msg mId2' msg2') <- signSendRecv rh2 rKey ("bcda", rId, SUB)
+      (Resp "bcda" _ (SOK Nothing), Resp "" _ (Msg mId2' msg2')) <- signSendRecv2 rh2 rKey ("bcda", rId, SUB)
      (dec mId2' msg2', Right "test2, no ACK") #== "same simplex queue via another TCP connection, tes2 delivered again (no ACK in 1st queue)"
      Resp "cdab" _ OK <- signSendRecv rh2 rKey ("cdab", rId, ACK mId2')

@@ -684,7 +694,7 @@ testWithStoreLog =
      nId <- readTVarIO notifierId
      Resp "dabc" _ (SOK Nothing) <- signSendRecv h1 nKey ("dabc", nId, NSUB)
      Resp "bcda" _ OK <- signSendRecv h sKey1 ("bcda", sId1, _SEND' "hello")
-      Resp "cdab" _ (Msg mId3 msg3) <- signSendRecv h rKey1 ("cdab", rId1, SUB)
+      (Resp "cdab" _ (SOK Nothing), Resp "" _ (Msg mId3 msg3)) <- signSendRecv2 h rKey1 ("cdab", rId1, SUB)
      (decryptMsgV3 dh1 mId3 msg3, Right "hello") #== "delivered from restored queue"
      Resp "" _ (NMSG _ _) <- tGet1 h1
      -- this queue is removed - not restored
@@ -769,7 +779,7 @@ testRestoreMessages =
      Just rKey <- readTVarIO recipientKey
      Just dh <- readTVarIO dhShared
      let dec = decryptMsgV3 dh
-      Resp "2" _ (Msg mId2 msg2) <- signSendRecv h rKey ("2", rId, SUB)
+      (Resp "2" _ (SOK Nothing), Resp "" _ (Msg mId2 msg2)) <- signSendRecv2 h rKey ("2", rId, SUB)
      (dec mId2 msg2, Right "hello 2") #== "restored message delivered"
      Resp "3" _ (Msg mId3 msg3) <- signSendRecv h rKey ("3", rId, ACK mId2)
      (dec mId3 msg3, Right "hello 3") #== "restored message delivered"
@@ -786,7 +796,7 @@ testRestoreMessages =
      Just rKey <- readTVarIO recipientKey
      Just dh <- readTVarIO dhShared
      let dec = decryptMsgV3 dh
-      Resp "4" _ (Msg mId4 msg4) <- signSendRecv h rKey ("4", rId, SUB)
+      (Resp "4" _ (SOK Nothing), Resp "" _ (Msg mId4 msg4)) <- signSendRecv2 h rKey ("4", rId, SUB)
      (dec mId4 msg4, Right "hello 4") #== "restored message delivered"
      Resp "5" _ (Msg mId5 msg5) <- signSendRecv h rKey ("5", rId, ACK mId4)
      (dec mId5 msg5, Right "hello 5") #== "restored message delivered"
@@ -1131,7 +1141,7 @@ testMsgExpireOnSend =
        threadDelay 2500000
        Resp "2" _ OK <- signSendRecv sh sKey ("2", sId, _SEND "hello (should NOT expire)")
        testSMPClient @c $ \rh -> do
-          Resp "3" _ (Msg mId msg) <- signSendRecv rh rKey ("3", rId, SUB)
+          (Resp "3" _ (SOK Nothing), Resp "" _ (Msg mId msg)) <- signSendRecv2 rh rKey ("3", rId, SUB)
          (dec mId msg, Right "hello (should NOT expire)") #== "delivered"
          1000 `timeout` tGetClient @SMPVersion @ErrorType @BrokerMsg rh >>= \case
            Nothing -> return ()
@@ -1139,8 +1149,7 @@ testMsgExpireOnSend =

 testMsgExpireOnInterval :: SpecWith (ASrvTransport, AStoreType)
 testMsgExpireOnInterval =
-  -- fails on ubuntu
-  xit' "should expire messages that are not received before messageTTL after expiry interval" $ \(ATransport (t :: TProxy c 'TServer), msType) -> do
+  it "should expire messages that are not received before messageTTL after expiry interval" $ \(ATransport (t :: TProxy c 'TServer), msType) -> do
    g <- C.newRandom
    (sPub, sKey) <- atomically $ C.generateAuthKeyPair C.SEd25519 g
    let cfg' = updateCfg (cfgMS msType) $ \cfg_ -> cfg_ {messageExpiration = Just ExpirationConfig {ttl = 1, checkInterval = 1}, idleQueueInterval = 1}
@@ -1151,7 +1160,7 @@ testMsgExpireOnInterval =
        threadDelay 3000000
        testSMPClient @c $ \rh -> do
          signSendRecv rh rKey ("2", rId, SUB) >>= \case
-            Resp "2" _ OK -> pure ()
+            Resp "2" _ (SOK Nothing) -> pure ()
            r -> unexpected r
          1000 `timeout` tGetClient @SMPVersion @ErrorType @BrokerMsg rh >>= \case
            Nothing -> return ()
@@ -1170,7 +1179,7 @@ testMsgNOTExpireOnInterval =
        Resp "1" _ OK <- signSendRecv sh sKey ("1", sId, _SEND "hello (should NOT expire)")
        threadDelay 2500000
        testSMPClient @c $ \rh -> do
-          Resp "2" _ (Msg mId msg) <- signSendRecv rh rKey ("2", rId, SUB)
+          (Resp "2" _ (SOK Nothing), Resp "" _ (Msg mId msg)) <- signSendRecv2 rh rKey ("2", rId, SUB)
          (dec mId msg, Right "hello (should NOT expire)") #== "delivered"
          1000 `timeout` tGetClient @SMPVersion @ErrorType @BrokerMsg rh >>= \case
            Nothing -> return ()