--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
#[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,
}
}
}

View file

@ -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);