feat: add config option for maximum number of client connections (#5513)

gateway/src/config.rs

@@ -97,6 +97,9 @@ pub struct Debug {
     /// Specifies maximum age of stored messages before they are removed from the storage
     pub stale_messages_max_age: Duration,
 
+    /// The maximum number of client connections the gateway will keep open at once.
+    pub maximum_open_connections: usize,
+
     pub zk_nym_tickets: ZkNymTicketHandlerDebug,
 }
 

gateway/src/node/client_handling/websocket/listener.rs

@@ -7,11 +7,14 @@ use nym_task::TaskClient;
 use rand::rngs::OsRng;
 use std::net::SocketAddr;
 use std::process;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
 use tokio::task::JoinHandle;
 use tracing::*;
 
 pub struct Listener {
     address: SocketAddr,
+    maximum_open_connections: usize,
     shared_state: CommonHandlerState,
     shutdown: TaskClient,
 }
@@ -19,11 +22,13 @@ pub struct Listener {
 impl Listener {
     pub(crate) fn new(
         address: SocketAddr,
+        maximum_open_connections: usize,
         shared_state: CommonHandlerState,
         shutdown: TaskClient,
     ) -> Self {
         Listener {
             address,
+            maximum_open_connections,
             shared_state,
             shutdown,
         }
@@ -41,6 +46,8 @@ impl Listener {
             }
         };
 
+        let open_connections = Arc::new(AtomicUsize::new(0));
+
         while !self.shutdown.is_shutdown() {
             tokio::select! {
                 biased;
@@ -52,6 +59,12 @@ impl Listener {
                         Ok((socket, remote_addr)) => {
                             let shutdown = self.shutdown.fork(format!("websocket_handler_{remote_addr}"));
                             trace!("received a socket connection from {remote_addr}");
+
+                            if open_connections.fetch_add(1, Ordering::SeqCst) >= self.maximum_open_connections {
+                                warn!("connection limit exceeded ({}). can't accept request from {remote_addr}", self.maximum_open_connections);
+                                continue;
+                            }
+
                             // TODO: I think we *REALLY* need a mechanism for having a maximum number of connected
                             // clients or spawned tokio tasks -> perhaps a worker system?
                             let handle = FreshHandler::new(
@@ -61,12 +74,15 @@ impl Listener {
                                 remote_addr,
                                 shutdown,
                             );
+                            let open_connections = open_connections.clone();
                             tokio::spawn(async move {
                                 // TODO: refactor it similarly to the mixnet listener on the nym-node
                                 let metrics_ref = handle.shared_state.metrics.clone();
                                 metrics_ref.network.new_ingress_websocket_client();
+                                open_connections.fetch_add(1, Ordering::SeqCst);
                                 handle.start_handling().await;
                                 metrics_ref.network.disconnected_ingress_websocket_client();
+                                open_connections.fetch_sub(1, Ordering::SeqCst);
                             });
                         }
                         Err(err) => warn!("failed to get client: {err}"),

gateway/src/node/mod.rs

@@ -262,6 +262,7 @@ impl GatewayTasksBuilder {
 
         Ok(websocket::Listener::new(
             self.config.gateway.websocket_bind_address,
+            self.config.debug.maximum_open_connections,
             shared_state,
             self.shutdown.fork("websocket"),
         ))

nym-node/src/config/gateway_tasks.rs

@@ -47,6 +47,9 @@ pub struct Debug {
     /// Number of messages from offline client that can be pulled at once (i.e. with a single SQL query) from the storage.
     pub message_retrieval_limit: i64,
 
+    /// The maximum number of client connections the gateway will keep open at once.
+    pub maximum_open_connections: usize,
+
     pub stale_messages: StaleMessageDebug,
 
     pub client_bandwidth: ClientBandwidthDebug,
@@ -56,12 +59,14 @@ pub struct Debug {
 
 impl Debug {
     const DEFAULT_MESSAGE_RETRIEVAL_LIMIT: i64 = 100;
+    const DEFAULT_MAXIMUM_OPEN_CONNECTIONS: usize = 8192;
 }
 
 impl Default for Debug {
     fn default() -> Self {
         Debug {
             message_retrieval_limit: Self::DEFAULT_MESSAGE_RETRIEVAL_LIMIT,
+            maximum_open_connections: Self::DEFAULT_MAXIMUM_OPEN_CONNECTIONS,
             stale_messages: Default::default(),
             client_bandwidth: Default::default(),
             zk_nym_tickets: Default::default(),

nym-node/src/config/helpers.rs

@@ -40,6 +40,7 @@ fn ephemeral_gateway_config(config: &Config) -> nym_gateway::config::Config {
                 .stale_messages
                 .cleaner_run_interval,
             stale_messages_max_age: config.gateway_tasks.debug.stale_messages.max_age,
+            maximum_open_connections: config.gateway_tasks.debug.maximum_open_connections,
             zk_nym_tickets: nym_gateway::config::ZkNymTicketHandlerDebug {
                 revocation_bandwidth_penalty: config
                     .gateway_tasks