smp server: refactor subscriptions and delivery in order to always response SOK on subscription with an optional message to follow. (#1573)

* smp server: refactor subscriptions and delivery

* metric for time between MSG and ACK

* cleanup

* refactor pattern match for ghc 8.10.7

* time buckets

* split max time metric

* histogram

* fix

Author: epoberezkin

src/Simplex/Messaging/Client.hs

@@ -1224,8 +1224,8 @@ okSMPCommands cmd c nm qs = L.map process <$> sendProtocolCommands c nm cs
       Left e -> Left e
 
 -- | Send SMP command
-sendSMPCommand :: PartyI p => SMPClient -> NetworkRequestMode -> Maybe C.APrivateAuthKey -> QueueId -> Command p -> ExceptT SMPClientError IO BrokerMsg
-sendSMPCommand c nm pKey qId cmd = sendProtocolCommand c nm pKey qId (Cmd sParty cmd)
+sendSMPCommand :: PartyI p => SMPClient -> NetworkRequestMode -> Maybe C.APrivateAuthKey -> EntityId -> Command p -> ExceptT SMPClientError IO BrokerMsg
+sendSMPCommand c nm pKey entId cmd = sendProtocolCommand c nm pKey entId (Cmd sParty cmd)
 {-# INLINE sendSMPCommand #-}
 
 type PCTransmission err msg = (Either TransportError SentRawTransmission, Request err msg)

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,
@@ -377,7 +378,7 @@ sameClient c cv = maybe False (sameClientId c) <$> readTVar cv
 data ClientSub
   = CSClient QueueId (Maybe ServiceId) (Maybe ServiceId) -- includes previous and new associated service IDs
   | CSDeleted QueueId (Maybe ServiceId) -- includes previously associated service IDs
-  | CSService ServiceId -- only send END to idividual client subs on message delivery, not of SSUB/NSSUB
+  | CSService ServiceId Int64 -- only send END to idividual client subs on message delivery, not of SSUB/NSSUB
 
 newtype ProxyAgent = ProxyAgent
   { smpAgent :: SMPClientAgent 'Sender
@@ -394,7 +395,7 @@ 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)),
+    sndQ :: TBQueue (NonEmpty (Transmission BrokerMsg), [Transmission BrokerMsg]),
     msgQ :: TBQueue (NonEmpty (Transmission BrokerMsg)),
     procThreads :: TVar Int,
     endThreads :: TVar (IntMap (Weak ThreadId)),
@@ -408,13 +409,15 @@ data Client s = Client
 
 type VerifiedTransmission s = (Maybe (StoreQueue s, QueueRec), Transmission Cmd)
 
+type ResponseAndMessage = (Transmission BrokerMsg, Maybe (Transmission BrokerMsg))
+
 data ServerSub = ServerSub (TVar SubscriptionThread) | ProhibitSub
 
 data SubscriptionThread = NoSub | SubPending | SubThread (Weak ThreadId)
 
 data Sub = Sub
   { subThread :: ServerSub, -- Nothing value indicates that sub
-    delivered :: TMVar (MsgId, RoundedSystemTime)
+    delivered :: TVar (Maybe (MsgId, RoundedSystemTime))
   }
 
 newServer :: IO (Server s)
@@ -497,13 +500,13 @@ newClient clientId qSize clientTHParams createdAt = do
 
 newSubscription :: SubscriptionThread -> STM Sub
 newSubscription st = do
-  delivered <- newEmptyTMVar
+  delivered <- newTVar Nothing
   subThread <- ServerSub <$> newTVar st
   return Sub {subThread, delivered}
 
 newProhibitedSub :: STM Sub
 newProhibitedSub = do
-  delivered <- newEmptyTMVar
+  delivered <- newTVar Nothing
   return Sub {subThread = ProhibitSub, delivered}
 
 newEnv :: ServerConfig s -> IO (Env s)

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

@@ -356,8 +356,8 @@ instance QueueStoreClass (JournalQueue s) (QStore s) where
   {-# INLINE setQueueService #-}
   getQueueNtfServices = withQS (getQueueNtfServices @(JournalQueue s))
   {-# INLINE getQueueNtfServices #-}
-  getNtfServiceQueueCount = withQS (getNtfServiceQueueCount @(JournalQueue s))
-  {-# INLINE getNtfServiceQueueCount #-}
+  getServiceQueueCount = withQS (getServiceQueueCount @(JournalQueue s))
+  {-# INLINE getServiceQueueCount #-}
 
 makeQueue_ :: JournalMsgStore s -> RecipientId -> QueueRec -> Lock -> IO (JournalQueue s)
 makeQueue_ JournalMsgStore {sharedLock} rId qr queueLock = do

src/Simplex/Messaging/Server/MsgStore/Journal.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.Text.Encoding (decodeLatin1, encodeUtf8)
@@ -65,7 +64,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)
@@ -84,7 +83,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
@@ -485,11 +484,15 @@ instance StoreQueueClass q => QueueStoreClass q (PostgresQueueStore q) where
         let (sNtfs, restNtfs) = partition (\(nId, _) -> S.member nId snIds) ntfs'
          in ((serviceId, sNtfs) : ssNtfs, restNtfs)
 
-  getNtfServiceQueueCount :: PostgresQueueStore q -> ServiceId -> IO (Either ErrorType Int64)
-  getNtfServiceQueueCount st serviceId =
-    E.uninterruptibleMask_ $ runExceptT $ withDB' "getNtfServiceQueueCount" st $ \db ->
+  getServiceQueueCount :: (PartyI p, ServiceParty p) => PostgresQueueStore q -> SParty p -> ServiceId -> IO (Either ErrorType Int64)
+  getServiceQueueCount st party serviceId =
+    E.uninterruptibleMask_ $ runExceptT $ withDB' "getServiceQueueCount" st $ \db ->
       fmap (fromMaybe 0) $ maybeFirstRow fromOnly $
-        DB.query db "SELECT count(1) FROM msg_queues WHERE ntf_service_id = ? AND deleted_at IS NULL" (Only serviceId)
+        DB.query db query (Only serviceId)
+    where
+      query = case party of
+        SRecipientService -> "SELECT count(1) FROM msg_queues WHERE rcv_service_id = ? AND deleted_at IS NULL"
+        SNotifierService -> "SELECT count(1) FROM msg_queues WHERE ntf_service_id = ? AND deleted_at IS NULL"
 
 batchInsertServices :: [STMService] -> PostgresQueueStore q -> IO Int64
 batchInsertServices services' toStore =

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

@@ -346,10 +346,15 @@ instance StoreQueueClass q => QueueStoreClass q (STMQueueStore q) where
         let (sNtfs, restNtfs) = partition (\(nId, _) -> S.member nId snIds) ntfs'
         pure ((Just serviceId, sNtfs) : ssNtfs, restNtfs)
 
-  getNtfServiceQueueCount :: STMQueueStore q -> ServiceId -> IO (Either ErrorType Int64)
-  getNtfServiceQueueCount st serviceId =
+  getServiceQueueCount :: (PartyI p, ServiceParty p) => STMQueueStore q -> SParty p -> ServiceId -> IO (Either ErrorType Int64)
+  getServiceQueueCount st party serviceId =
     TM.lookupIO serviceId (services st) >>=
-      maybe (pure $ Left AUTH) (fmap (Right . fromIntegral . S.size) . readTVarIO . serviceNtfQueues)
+      maybe (pure $ Left AUTH) (fmap (Right . fromIntegral . S.size) . readTVarIO . serviceSel)
+    where
+      serviceSel :: STMService -> TVar (Set QueueId)
+      serviceSel = case party of
+        SRecipientService -> serviceRcvQueues
+        SNotifierService -> serviceNtfQueues
 
 withQueueRec :: TVar (Maybe QueueRec) -> (QueueRec -> STM a) -> IO (Either ErrorType a)
 withQueueRec qr a = atomically $ readQueueRec qr >>= mapM a

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

@@ -48,7 +48,7 @@ class StoreQueueClass q => QueueStoreClass q s where
   getCreateService :: s -> ServiceRec -> IO (Either ErrorType ServiceId)
   setQueueService :: (PartyI p, ServiceParty p) => s -> q -> SParty p -> Maybe ServiceId -> IO (Either ErrorType ())
   getQueueNtfServices :: s -> [(NotifierId, a)] -> IO (Either ErrorType ([(Maybe ServiceId, [(NotifierId, a)])], [(NotifierId, a)]))
-  getNtfServiceQueueCount :: s -> ServiceId -> IO (Either ErrorType Int64)
+  getServiceQueueCount :: (PartyI p, ServiceParty p) => s -> SParty p -> ServiceId -> IO (Either ErrorType Int64)
 
 data EntityCounts = EntityCounts
   { queueCount :: Int,

src/Simplex/Messaging/Server/QueueStore/Types.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