--max-sessions option

src/args.rs

@@ -107,6 +107,10 @@ pub struct Args {
     /// Exit immediately when fatal error occurs, useful for running as a service
     #[arg(long)]
     pub exit_on_fatal_error: bool,
+
+    /// Maximum number of sessions to be handled concurrently
+    #[arg(long, value_name = "number", default_value = "200")]
+    pub max_sessions: usize,
 }
 
 fn validate_tun(p: &str) -> Result<String> {
@@ -149,6 +153,7 @@ impl Default for Args {
             virtual_dns_pool: IpCidr::from_str("198.18.0.0/15").unwrap(),
             daemonize: false,
             exit_on_fatal_error: false,
+            max_sessions: 200,
         }
     }
 }

src/lib.rs

@@ -65,8 +65,6 @@ pub mod win_svc;
 
 const DNS_PORT: u16 = 53;
 
-const MAX_SESSIONS: u64 = 200;
-
 static TASK_COUNT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
 use std::sync::atomic::Ordering::Relaxed;
 
@@ -244,13 +242,15 @@ where
                 ip_stack_stream?
             }
         };
+        let max_sessions = args.max_sessions as u64;
         match ip_stack_stream {
             IpStackStream::Tcp(tcp) => {
-                if TASK_COUNT.load(Relaxed) > MAX_SESSIONS {
-                    log::warn!("Too many sessions that over {MAX_SESSIONS}, dropping new session");
+                if TASK_COUNT.load(Relaxed) > max_sessions {
                     if args.exit_on_fatal_error {
+                        log::info!("Too many sessions that over {max_sessions}, exiting...");
                         break;
                     }
+                    log::warn!("Too many sessions that over {max_sessions}, dropping new session");
                     continue;
                 }
                 log::trace!("Session count {}", TASK_COUNT.fetch_add(1, Relaxed) + 1);
@@ -272,11 +272,12 @@ where
                 });
             }
             IpStackStream::Udp(udp) => {
-                if TASK_COUNT.load(Relaxed) > MAX_SESSIONS {
-                    log::warn!("Too many sessions that over {MAX_SESSIONS}, dropping new session");
+                if TASK_COUNT.load(Relaxed) > max_sessions {
                     if args.exit_on_fatal_error {
+                        log::info!("Too many sessions that over {max_sessions}, exiting...");
                         break;
                     }
+                    log::warn!("Too many sessions that over {max_sessions}, dropping new session");
                     continue;
                 }
                 log::trace!("Session count {}", TASK_COUNT.fetch_add(1, Relaxed) + 1);