From eacaf8e5aebcb73545231b602d5fd98a8f89bba9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C4=99drzej=20Stuczy=C5=84ski?= Date: Fri, 14 Mar 2025 15:47:17 +0000 Subject: [PATCH] hotfix: correctly increment ws connection counter (#5620) --- .../client_handling/websocket/common_state.rs | 1 - .../websocket/connection_handler/fresh.rs | 3 - .../client_handling/websocket/listener.rs | 110 ++++++++++++------ nym-node/nym-node-metrics/src/network.rs | 2 +- 4 files changed, 74 insertions(+), 42 deletions(-) diff --git a/gateway/src/node/client_handling/websocket/common_state.rs b/gateway/src/node/client_handling/websocket/common_state.rs index a571117895..d7a92616ea 100644 --- a/gateway/src/node/client_handling/websocket/common_state.rs +++ b/gateway/src/node/client_handling/websocket/common_state.rs @@ -19,7 +19,6 @@ pub(crate) struct Config { pub(crate) bandwidth: BandwidthFlushingBehaviourConfig, } -// I can see this being possible expanded with say storage or client store #[derive(Clone)] pub(crate) struct CommonHandlerState { pub(crate) cfg: Config, diff --git a/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs b/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs index bd1f9fd601..c1f1bb0812 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs @@ -131,9 +131,6 @@ impl FreshHandler { // for time being we assume handle is always constructed from raw socket. // if we decide we want to change it, that's not too difficult - // also at this point I'm not entirely sure how to deal with this warning without - // some considerable refactoring - #[allow(clippy::too_many_arguments)] pub(crate) fn new( rng: R, conn: S, diff --git a/gateway/src/node/client_handling/websocket/listener.rs b/gateway/src/node/client_handling/websocket/listener.rs index 2d358840e7..94e11122ef 100644 --- a/gateway/src/node/client_handling/websocket/listener.rs +++ b/gateway/src/node/client_handling/websocket/listener.rs @@ -6,9 +6,8 @@ use crate::node::client_handling::websocket::connection_handler::FreshHandler; 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 std::{io, process}; +use tokio::net::TcpStream; use tokio::task::JoinHandle; use tracing::*; @@ -34,6 +33,76 @@ impl Listener { } } + fn active_connections(&self) -> usize { + self.shared_state + .metrics + .network + .active_ingress_websocket_connections_count() + } + + fn prepare_connection_handler( + &self, + socket: TcpStream, + remote_address: SocketAddr, + ) -> FreshHandler { + let shutdown = self + .shutdown + .fork(format!("websocket_handler_{remote_address}")); + FreshHandler::new( + OsRng, + socket, + self.shared_state.clone(), + remote_address, + shutdown, + ) + } + + fn try_handle_accepted_connection(&self, accepted: io::Result<(TcpStream, SocketAddr)>) { + match accepted { + Ok((socket, remote_address)) => { + trace!("received a socket connection from {remote_address}"); + + let active = self.active_connections(); + + // 1. check if we're within the connection limit + if active >= self.maximum_open_connections { + warn!( + "connection limit exceeded ({}). can't accept request from {remote_address}", + self.maximum_open_connections + ); + return; + } + + debug!("there are currently {active} connected clients on the gateway websocket"); + + // 2. prepare shared data for the new connection handler + let handle = self.prepare_connection_handler(socket, remote_address); + + // 3. increment the connection counter. + // make sure to do it before spawning the task, + // as another connection might get accepted before the task is scheduled + // for execution + self.shared_state + .metrics + .network + .new_ingress_websocket_client(); + + // 4. spawn the task handling the client connection + tokio::spawn(async move { + // TODO: refactor it similarly to the mixnet listener on the nym-node + let metrics_ref = handle.shared_state.metrics.clone(); + + // 4.1. handle all client requests until connection gets terminated + handle.start_handling().await; + + // 4.2. decrement the connection counter + metrics_ref.network.disconnected_ingress_websocket_client(); + }); + } + Err(err) => warn!("failed to accept client connection: {err}"), + } + } + // TODO: change the signature to pub(crate) async fn run(&self, handler: Handler) pub(crate) async fn run(&mut self) { @@ -46,8 +115,6 @@ impl Listener { } }; - let open_connections = Arc::new(AtomicUsize::new(0)); - while !self.shutdown.is_shutdown() { tokio::select! { biased; @@ -55,38 +122,7 @@ impl Listener { trace!("client_handling::Listener: received shutdown"); } connection = tcp_listener.accept() => { - match connection { - 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( - OsRng, - socket, - self.shared_state.clone(), - 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}"), - } + self.try_handle_accepted_connection(connection) } } diff --git a/nym-node/nym-node-metrics/src/network.rs b/nym-node/nym-node-metrics/src/network.rs index 7de912fc3c..ea2b41ff52 100644 --- a/nym-node/nym-node-metrics/src/network.rs +++ b/nym-node/nym-node-metrics/src/network.rs @@ -45,7 +45,7 @@ impl NetworkStats { pub fn active_ingress_websocket_connections_count(&self) -> usize { self.active_ingress_websocket_connections - .load(Ordering::Relaxed) + .load(Ordering::SeqCst) } pub fn active_egress_mixnet_connections_counter(&self) -> Arc {