From 266ade40d9eb84f98d605fdfe5df0f95917bc8d2 Mon Sep 17 00:00:00 2001
From: ssrlive <30760636+ssrlive@users.noreply.github.com>
Date: Wed, 15 May 2024 14:45:18 +0800
Subject: [PATCH] session count on server and client

---
 .gitignore    |  1 +
 src/client.rs | 10 ++++++++++
 src/server.rs | 10 ++++++++++
 3 files changed, 21 insertions(+)

diff --git a/.gitignore b/.gitignore
index 6288446..0aa4f14 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,6 +5,7 @@ nginx_signing.key
 overtls-daemon.sh
 project.xcworkspace/
 xcuserdata/
+.vs/
 .vscode/
 .VSCodeCounter/
 .env
diff --git a/src/client.rs b/src/client.rs
index 67efe8a..cc3a053 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -73,6 +73,9 @@ where
     let (udp_tx, _, incomings) = udprelay::create_udp_tunnel();
     udprelay::udp_handler_watchdog(config, &incomings, &udp_tx, quit.clone()).await?;
 
+    let session_id = Arc::new(std::sync::atomic::AtomicUsize::new(0));
+    let session_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
+
     loop {
         tokio::select! {
             _ = quit.cancelled() => {
@@ -84,10 +87,17 @@ where
                 let config = config.clone();
                 let udp_tx = udp_tx.clone();
                 let incomings = incomings.clone();
+                let session_id = session_id.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+                let session_count = session_count.clone();
+                let peer_addr = conn.peer_addr()?;
                 tokio::spawn(async move {
+                    let count = session_count.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1;
+                    log::debug!("session #{} from {} started, session count {}", session_id, peer_addr, count);
                     if let Err(e) = handle_incoming(conn, config, Some(udp_tx), incomings).await {
                         log::debug!("{}", e);
                     }
+                    let count = session_count.fetch_sub(1, std::sync::atomic::Ordering::SeqCst) - 1;
+                    log::debug!("session #{} from {} ended, session count {}", session_id, peer_addr, count);
                 });
             }
         }
diff --git a/src/server.rs b/src/server.rs
index 6113a69..d7a502b 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -83,6 +83,9 @@ pub async fn run_server(config: &Config, exiting_flag: crate::CancellationToken)
         }
     };
 
+    let session_id = Arc::new(std::sync::atomic::AtomicUsize::new(0));
+    let session_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
+
     loop {
         tokio::select! {
             _ = exiting_flag.cancelled() => {
@@ -105,10 +108,17 @@ pub async fn run_server(config: &Config, exiting_flag: crate::CancellationToken)
                     Ok::<_, Error>(())
                 };
 
+                let session_id = session_id.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+                let session_count = session_count.clone();
+
                 tokio::spawn(async move {
+                    let count = session_count.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1;
+                    log::debug!("session #{} from {} started, session count {}", session_id, peer_addr, count);
                     if let Err(e) = incoming_task.await {
                         log::debug!("{peer_addr}: {e}");
                     }
+                    let count = session_count.fetch_sub(1, std::sync::atomic::Ordering::SeqCst) - 1;
+                    log::debug!("session #{} from {} ended, session count {}", session_id, peer_addr, count);
                 });
             }
         }