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

Author: epoberezkin

src/Simplex/Messaging/Protocol.hs

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

src/Simplex/Messaging/Server.hs

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

src/Simplex/Messaging/Server/Env/STM.hs

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

src/Simplex/Messaging/Server/QueueStore/Postgres.hs

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

tests/CoreTests/BatchingTests.hs

@@ -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,16 +1141,15 @@ 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 ()
             Just _ -> error "nothing else should be delivered"
 
 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 ()