From 2cb3817b2ccd63d38496fee8ecc7e14b7919e4f7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C4=99drzej=20Stuczy=C5=84ski?= Date: Wed, 26 Feb 2025 09:48:13 +0000 Subject: [PATCH] feat: add config option for maximum number of client connections (#5513) --- gateway/src/config.rs | 3 +++ .../node/client_handling/websocket/listener.rs | 16 ++++++++++++++++ gateway/src/node/mod.rs | 1 + nym-node/src/config/gateway_tasks.rs | 5 +++++ nym-node/src/config/helpers.rs | 1 + 5 files changed, 26 insertions(+) diff --git a/gateway/src/config.rs b/gateway/src/config.rs index 66d6180f31..02d58ad143 100644 --- a/gateway/src/config.rs +++ b/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, } diff --git a/gateway/src/node/client_handling/websocket/listener.rs b/gateway/src/node/client_handling/websocket/listener.rs index b5eebc3cab..2d358840e7 100644 --- a/gateway/src/node/client_handling/websocket/listener.rs +++ b/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}"), diff --git a/gateway/src/node/mod.rs b/gateway/src/node/mod.rs index 1dd6f2c217..9323bb572b 100644 --- a/gateway/src/node/mod.rs +++ b/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"), )) diff --git a/nym-node/src/config/gateway_tasks.rs b/nym-node/src/config/gateway_tasks.rs index 9f814f2bf5..1a4c7454ac 100644 --- a/nym-node/src/config/gateway_tasks.rs +++ b/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(), diff --git a/nym-node/src/config/helpers.rs b/nym-node/src/config/helpers.rs index 2afd58301c..f3380ef223 100644 --- a/nym-node/src/config/helpers.rs +++ b/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