smp-server: limit forwards in flight per relay

src/Simplex/Messaging/Server.hs

@@ -1454,25 +1454,35 @@ client
                         Nothing -> ERR $ transportErr TENoServerAuth
                       _ -> ERR $ transportErr TEVersion
       PFWD fwdV pubKey encBlock -> do
-        ProxyAgent {smpAgent = a} <- asks proxyAgent
+        ProxyAgent {smpAgent = a, relayForwards} <- asks proxyAgent
         ServerStats {pMsgFwds, pMsgFwdsOwn} <- asks serverStats
         let inc = mkIncProxyStats pMsgFwds pMsgFwdsOwn
         liftIO (lookupSMPServerClient a sessId) >>= \case
           Just (own, smp) -> do
             inc own pRequests
-            forkProxiedCmd $ do
-              liftIO (runExceptT (forwardSMPTransmission smp corrId fwdV pubKey encBlock) `E.catches` clientHandlers)  >>= \case
-                Right r -> PRES r <$ inc own pSuccesses
-                Left e -> ERR (smpProxyError e) <$ case e of
-                  PCEProtocolError {} -> inc own pSuccesses
-                  PCEResponseTimeout -> do
-                    inc own pErrorsOther
-                    liftIO $ closeTimedOutClient (smpPingCount $ networkConfig $ smpCfg $ agentCfg a) smp
-                  _ -> inc own pErrorsOther
+            maxFwds <- asks $ proxyRelayConcurrency . config
+            -- reserved on the forked thread, so that the slot is released when the forward completes
+            forkProxiedCmd $
+              bracket (atomically $ reserve maxFwds relayForwards) (\r -> when r $ atomically $ release relayForwards) $ \case
+                -- the relay is not keeping up, the client retries TIMEOUT later
+                False -> ERR (smpProxyError PCEResponseTimeout) <$ inc own pErrorsOther
+                True ->
+                  liftIO (runExceptT (forwardSMPTransmission smp corrId fwdV pubKey encBlock) `E.catches` clientHandlers) >>= \case
+                    Right r -> PRES r <$ inc own pSuccesses
+                    Left e -> ERR (smpProxyError e) <$ case e of
+                      PCEProtocolError {} -> inc own pSuccesses
+                      PCEResponseTimeout -> do
+                        inc own pErrorsOther
+                        liftIO $ closeTimedOutClient (smpPingCount $ networkConfig $ smpCfg $ agentCfg a) smp
+                      _ -> inc own pErrorsOther
           Nothing -> inc False pRequests >> inc False pErrorsConnect $> Just (ERR $ PROXY NO_SESSION)
       where
         forkProxiedCmd :: M s BrokerMsg -> M s (Maybe BrokerMsg)
         forkProxiedCmd = forkCmd serverClientConcurrency corrId (EntityId sessId)
+        reserve maxFwds v = stateTVar v $ \fwds -> case M.lookup sessId fwds of
+          Just n | n >= maxFwds -> (False, fwds)
+          n_ -> (True, M.insert sessId (maybe 1 (+ 1) n_) fwds)
+        release v = modifyTVar' v $ M.update (\n -> if n > 1 then Just (n - 1) else Nothing) sessId
     -- Run a slow command on a thread
     forkCmd :: (ServerConfig s -> Int) -> CorrId -> EntityId -> M s BrokerMsg -> M s (Maybe a)
     forkCmd concurrency corrId entId cmdAction = do

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

@@ -67,6 +67,7 @@ module Simplex.Messaging.Server.Env.STM
     defaultNtfExpiration,
     defaultInactiveClientExpiration,
     defaultProxyClientConcurrency,
+    defaultProxyRelayConcurrency,
     defaultNameResolverConcurrency,
     defaultMaxJournalMsgCount,
     defaultMaxJournalStateLines,
@@ -128,7 +129,7 @@ import Simplex.Messaging.Server.StoreLog.ReadWrite
 import Simplex.Messaging.SystemTime
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
-import Simplex.Messaging.Transport (ASrvTransport, SMPVersion, THandleParams, TransportPeer (..), VersionRangeSMP)
+import Simplex.Messaging.Transport (ASrvTransport, SMPVersion, SessionId, THandleParams, TransportPeer (..), VersionRangeSMP)
 import Simplex.Messaging.Transport.Server
 import Simplex.Messaging.Util (ifM, tshow, whenM, ($>>=))
 import System.Directory (doesFileExist)
@@ -199,6 +200,8 @@ data ServerConfig s = ServerConfig
     smpAgentCfg :: SMPClientAgentConfig,
     allowSMPProxy :: Bool, -- auth is the same with `newQueueBasicAuth`
     serverClientConcurrency :: Int,
+    -- | max forwarded commands awaiting a response from one destination relay
+    proxyRelayConcurrency :: Int,
     -- | max concurrent name resolutions per connection, enforced in forkCmd.
     -- Much higher than serverClientConcurrency: forwarded RSLVs from many clients
     -- aggregate over a single proxy->relay connection (only servers send proxied
@@ -252,6 +255,9 @@ defaultInactiveClientExpiration =
 defaultProxyClientConcurrency :: Int
 defaultProxyClientConcurrency = 32
 
+defaultProxyRelayConcurrency :: Int
+defaultProxyRelayConcurrency = 512
+
 defaultNameResolverConcurrency :: Int
 defaultNameResolverConcurrency = 1000
 
@@ -441,8 +447,10 @@ data ClientSub
   | CSDeleted QueueId (Maybe ServiceId) -- includes previously associated service IDs
   | CSService ServiceId (Int64, IdsHash) -- only send END to idividual client subs on message delivery, not of SSUB/NSSUB
 
-newtype ProxyAgent = ProxyAgent
-  { smpAgent :: SMPClientAgent 'Sender
+data ProxyAgent = ProxyAgent
+  { smpAgent :: SMPClientAgent 'Sender,
+    -- forwarded commands awaiting a response, per relay session
+    relayForwards :: TVar (Map SessionId Int)
   }
 
 type ClientId = Int
@@ -731,7 +739,8 @@ mkJournalStoreConfig queueStoreCfg storePath msgQueueQuota maxJournalMsgCount ma
 newSMPProxyAgent :: SMPClientAgentConfig -> TVar ChaChaDRG -> IO ProxyAgent
 newSMPProxyAgent smpAgentCfg random = do
   smpAgent <- newSMPClientAgent SSender smpAgentCfg Nothing random
-  pure ProxyAgent {smpAgent}
+  relayForwards <- newTVarIO mempty
+  pure ProxyAgent {smpAgent, relayForwards}
 
 readWriteQueueStore :: forall q. StoreQueueClass q => Bool -> (RecipientId -> QueueRec -> IO q) -> FilePath -> STMQueueStore q -> IO (StoreLog 'WriteMode)
 readWriteQueueStore tty mkQ = readWriteStoreLog (readQueueStore tty mkQ) (writeQueueStore @q)

src/Simplex/Messaging/Server/Main.hs

@@ -609,6 +609,7 @@ smpServerCLI_ generateSite serveStaticFiles attachStaticFiles cfgPath logPath =
                   },
               allowSMPProxy = True,
               serverClientConcurrency = readIniDefault defaultProxyClientConcurrency "PROXY" "client_concurrency" ini,
+              proxyRelayConcurrency = readIniDefault defaultProxyRelayConcurrency "PROXY" "relay_concurrency" ini,
               serverResolverConcurrency = readIniDefault defaultNameResolverConcurrency "NAMES" "resolver_concurrency" ini,
               namesConfig = readNamesConfig ini,
               information = serverPublicInfo ini,

src/Simplex/Messaging/Server/Main/Init.hs

@@ -154,6 +154,9 @@ iniFileContent cfgPath logPath opts host basicAuth controlPortPwds =
         \# socks_mode = onion\n\n\
         \# Limit number of threads a client can spawn to process proxy commands in parrallel.\n"
     <> ("# client_concurrency = " <> tshow defaultProxyClientConcurrency)
+    <> "\n\n\
+        \# Limit number of messages forwarded to one destination server and awaiting its response.\n"
+    <> ("# relay_concurrency = " <> tshow defaultProxyRelayConcurrency)
     <> "\n\n\
         \[NAMES]\n\
         \# Public-namespace resolution via the snrc-resolve.py REST resolver.\n\