hotfix: correctly increment ws connection counter (#5620)

Author: jstuczyn

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,

gateway/src/node/client_handling/websocket/connection_handler/fresh.rs

@@ -131,9 +131,6 @@ impl<R, S> FreshHandler<R, S> {
 
     // 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,

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<OsRng, TcpStream> {
+        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,47 +115,14 @@ impl Listener {
             }
         };
 
-        let open_connections = Arc::new(AtomicUsize::new(0));
-
         while !self.shutdown.is_shutdown() {
             tokio::select! {
                 biased;
                 _ = self.shutdown.recv() => {
                     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)
                 }
 
             }

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<AtomicUsize> {