diff --git a/mac/apparmor/yumechi-no-kuni-proxy-worker b/mac/apparmor/yumechi-no-kuni-proxy-worker index 0fb5240..f4e7516 100644 --- a/mac/apparmor/yumechi-no-kuni-proxy-worker +++ b/mac/apparmor/yumechi-no-kuni-proxy-worker @@ -74,7 +74,8 @@ profile yumechi-no-kuni-proxy-worker @{prog_path} { /{,usr/}{,local/}{,s}bin/@{prog} ixr, owner /var/lib/@{prog}/{,bin}/@{prog} ixr, - signal (send) peer=yume-proxy-workers//serve//image, + signal (send, receive) set=int,term,kill peer=yume-proxy-workers//serve, + signal (send) set=int,term,kill,usr1 peer=yume-proxy-workers//serve//image, ^image { diff --git a/src/lib.rs b/src/lib.rs index b52b041..ad6c5c5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -828,19 +828,42 @@ pub struct App { /// without triggering the resource limits. #[allow(unsafe_code)] pub fn register_cancel_handler() { + static STRIKES: AtomicU64 = AtomicU64::new(0); + #[cfg(target_family = "unix")] unsafe { + fn estop() { + unsafe { + let pid = libc::getpid(); + libc::kill(pid, libc::SIGTERM); + libc::sleep(5); + libc::kill(pid, libc::SIGTERM); + libc::sleep(2); + libc::kill(pid, libc::SIGKILL); + } + } unsafe extern "C" fn cancel_handler(_: libc::c_int) { - log::error!("Received cancel signal, stopping thread"); + let strikes = STRIKES.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + log::error!( + "Thread blocked for too long, force pthread_exit() (strikes: {}/{})", + strikes, + 5 + ); + if strikes == 5 { + log::error!("Too many strikes, exiting"); + estop(); + } libc::pthread_exit(&sandbox::EXIT_TIMEOUT as *const _ as *mut _); } if libc::signal(libc::SIGUSR1, cancel_handler as usize) == libc::SIG_ERR { log::error!( - "Failed to register cancel handler: {}", + "Failed to register worker cancel handler: {}", std::io::Error::last_os_error() ); } + + log::info!("Registered worker cancel handler."); } } diff --git a/src/main.rs b/src/main.rs index 967ffc3..4c00e6c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -227,6 +227,8 @@ fn main() { } } + yumechi_no_kuni_proxy_worker::register_cancel_handler(); + let runtime = tokio::runtime::Builder::new_multi_thread() .max_blocking_threads(512) .enable_all() @@ -236,23 +238,8 @@ fn main() { log::info!("Spawned Tokio reactor."); - yumechi_no_kuni_proxy_worker::register_cancel_handler(); - runtime .block_on(async move { - #[cfg(not(feature = "axum-server"))] - { - let listener = tokio::net::TcpListener::from_std(listener) - .expect("Failed to pass socket to Tokio"); - - log::warn!( - "Built without axum-server feature, using hyper without graceful shutdown" - ); - log::info!("Ready for connections."); - axum::serve(listener, ms).await.expect("Failed to serve"); - return; - } - #[cfg(feature = "metrics")] { let reg = prometheus::default_registry(); @@ -298,6 +285,19 @@ fn main() { }); } + #[cfg(not(feature = "axum-server"))] + { + let listener = tokio::net::TcpListener::from_std(listener) + .expect("Failed to pass socket to Tokio"); + + log::warn!( + "Built without axum-server feature, using hyper without graceful shutdown" + ); + log::info!("Ready for connections."); + axum::serve(listener, ms).await.expect("Failed to serve"); + return; + } + #[cfg(feature = "axum-server")] { let handle = axum_server::Handle::new();