--max-sessions option

This commit is contained in:
ssrlive 2024-10-09 16:57:14 +08:00
parent b03032b8cd
commit 4ef71a5b4c
2 changed files with 12 additions and 6 deletions

View file

@ -107,6 +107,10 @@ pub struct Args {
/// Exit immediately when fatal error occurs, useful for running as a service /// Exit immediately when fatal error occurs, useful for running as a service
#[arg(long)] #[arg(long)]
pub exit_on_fatal_error: bool, 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> { 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(), virtual_dns_pool: IpCidr::from_str("198.18.0.0/15").unwrap(),
daemonize: false, daemonize: false,
exit_on_fatal_error: false, exit_on_fatal_error: false,
max_sessions: 200,
} }
} }
} }

View file

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