From e490bc7829f6bb0c049a8668ea7d8f34c7d9ae04 Mon Sep 17 00:00:00 2001 From: ssrlive <30760636+ssrlive@users.noreply.github.com> Date: Tue, 9 Sep 2025 13:47:11 +0800 Subject: [PATCH] Fix process stuck issue while exiting --- src/bin/main.rs | 19 +++++++------ src/general_api.rs | 68 ++++++++++++++++++++++++++++------------------ src/lib.rs | 2 ++ src/win_svc.rs | 17 ++++++------ 4 files changed, 63 insertions(+), 43 deletions(-) diff --git a/src/bin/main.rs b/src/bin/main.rs index 0c6e89e..192abf9 100644 --- a/src/bin/main.rs +++ b/src/bin/main.rs @@ -24,7 +24,16 @@ fn main() -> Result<(), BoxError> { } let rt = tokio::runtime::Builder::new_multi_thread().enable_all().build()?; - rt.block_on(main_async(args)) + rt.block_on(async move { + let res = main_async(args).await; + let _h = tokio::spawn(async move { + // Delay some seconds then try to exit current process if not exited yet, normally this case should not happen + tokio::time::sleep(std::time::Duration::from_secs(tun2proxy::FORCE_EXIT_TIMEOUT)).await; + log::info!("Forcing exit now."); + std::process::exit(-1); + }); + res + }) } fn setup_logging(args: &Args) { @@ -78,7 +87,7 @@ async fn main_async(args: Args) -> Result<(), BoxError> { true })?; - let tasks = main_loop_handle.await??; + let _tasks = main_loop_handle.await??; if ctrlc_fired.load(std::sync::atomic::Ordering::SeqCst) { log::info!("Ctrl-C fired, waiting the handler to finish..."); @@ -89,12 +98,6 @@ async fn main_async(args: Args) -> Result<(), BoxError> { } } - if args.exit_on_fatal_error && tasks >= args.max_sessions { - // Because `main_async` function perhaps stuck in `await` state, so we need to exit the process forcefully - log::info!("Internal fatal error, max sessions reached ({tasks}/{})", args.max_sessions); - std::process::exit(-1); - } - Ok(()) } diff --git a/src/general_api.rs b/src/general_api.rs index e713409..bdf0371 100644 --- a/src/general_api.rs +++ b/src/general_api.rs @@ -4,8 +4,6 @@ use crate::{ }; use std::os::raw::{c_char, c_int, c_ushort}; -static TUN_QUIT: std::sync::Mutex> = std::sync::Mutex::new(None); - /// # Safety /// /// Run the tun2proxy component with some arguments. @@ -97,6 +95,18 @@ pub unsafe extern "C" fn tun2proxy_run_with_cli_args(cli_args: *const c_char, tu general_run_for_api(args, tun_mtu, packet_information) } +static TUN_QUIT: std::sync::Mutex> = std::sync::Mutex::new(None); + +pub(crate) fn tun2proxy_stop_internal() -> c_int { + if let Ok(mut lock) = TUN_QUIT.lock() { + if let Some(shutdown_token) = lock.take() { + shutdown_token.cancel(); + return 0; + } + } + -1 +} + pub fn general_run_for_api(args: Args, tun_mtu: u16, packet_information: bool) -> c_int { log::set_max_level(args.verbosity.into()); if let Err(err) = log::set_boxed_logger(Box::::default()) { @@ -119,26 +129,34 @@ pub fn general_run_for_api(args: Args, tun_mtu: u16, packet_information: bool) - log::error!("failed to create tokio runtime with"); return -3; }; - match rt.block_on(async move { - let ret = general_run_async(args.clone(), tun_mtu, packet_information, shutdown_token).await; - match &ret { - Ok(sessions) => { - if args.exit_on_fatal_error && *sessions >= args.max_sessions { - log::error!("Forced exit due to max sessions reached ({sessions}/{})", args.max_sessions); - std::process::exit(-1); - } - log::debug!("tun2proxy exited normally, current sessions: {sessions}"); - } - Err(err) => log::error!("main loop error: {err}"), - } + let args_clone = args.clone(); + let res = rt.block_on(async move { + let ret = general_run_async(args_clone, tun_mtu, packet_information, shutdown_token).await; + let _h = tokio::spawn(async move { + // Delay some seconds then try to exit current process if not exited yet, normally this case should not happen + tokio::time::sleep(std::time::Duration::from_secs(crate::FORCE_EXIT_TIMEOUT)).await; + log::info!("Forcing exit now."); + std::process::exit(-1); + }); ret - }) { - Ok(_) => 0, + }); + + let res = match res { + Ok(sessions) => { + log::debug!("tun2proxy exited normally, current session count: {sessions}"); + 0 + } Err(e) => { log::error!("failed to run tun2proxy with error: {e:?}"); -4 } + }; + + if let Ok(mut lock) = TUN_QUIT.lock() { + lock.take(); } + + res } /// Run the tun2proxy component with some arguments. @@ -238,12 +256,18 @@ pub async fn general_run_async( } } - let join_handle = tokio::spawn(crate::run(device, tun_mtu, args, shutdown_token.clone())); + let join_handle = tokio::spawn(crate::run(device, tun_mtu, args.clone(), shutdown_token.clone())); match join_handle.await? { Ok(sessions) => { #[cfg(any(target_os = "linux", target_os = "windows", target_os = "macos"))] tproxy_config::tproxy_remove(restore).await?; + + let max_sessions = args.max_sessions; + if args.exit_on_fatal_error && sessions >= max_sessions { + let info = format!("Forced exit due to max sessions reached ({sessions}/{max_sessions})"); + return Err(std::io::Error::other(info)); + } Ok(sessions) } Err(err) => Err(std::io::Error::from(err)), @@ -257,13 +281,3 @@ pub async fn general_run_async( pub unsafe extern "C" fn tun2proxy_stop() -> c_int { tun2proxy_stop_internal() } - -pub(crate) fn tun2proxy_stop_internal() -> c_int { - if let Ok(mut lock) = TUN_QUIT.lock() { - if let Some(shutdown_token) = lock.take() { - shutdown_token.cancel(); - return 0; - } - } - -1 -} diff --git a/src/lib.rs b/src/lib.rs index 1f9e68f..eb4f628 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -44,6 +44,8 @@ static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc; pub use general_api::general_run_async; +pub const FORCE_EXIT_TIMEOUT: u64 = 2; // seconds + mod android; mod args; mod directions; diff --git a/src/win_svc.rs b/src/win_svc.rs index 2fed698..a63d64b 100644 --- a/src/win_svc.rs +++ b/src/win_svc.rs @@ -80,15 +80,16 @@ fn run_service(_arguments: Vec) -> Result<(), crate::BoxErro let ret = crate::general_run_async(args.clone(), tun::DEFAULT_MTU, false, shutdown_token).await; match &ret { - Ok(sessions) => { - if args.exit_on_fatal_error && *sessions >= args.max_sessions { - log::error!("Forced exit due to max sessions reached ({sessions}/{})", args.max_sessions); - std::process::exit(-1); - } - log::debug!("tun2proxy exited normally, current sessions: {sessions}"); - } - Err(err) => log::error!("main loop error: {err}"), + Ok(sessions) => log::debug!("tun2proxy exited normally, current session count: {sessions}"), + Err(e) => log::error!("failed to run tun2proxy with error: {e:?}"), } + let _h = tokio::spawn(async move { + // Delay some seconds then try to exit current process if not exited yet, normally this case should not happen + tokio::time::sleep(std::time::Duration::from_secs(crate::FORCE_EXIT_TIMEOUT)).await; + log::info!("Forcing exit now."); + std::process::exit(-1); + }); + Ok::<(), crate::Error>(()) })?; }