From 8f1853e7737b2d5dbd731976cd235e3394d2f0f8 Mon Sep 17 00:00:00 2001
From: eternal-flame-AD <yume@yumechi.jp>
Date: Fri, 22 Nov 2024 18:21:59 -0600
Subject: [PATCH] restart process when too many runaway process happened

Signed-off-by: eternal-flame-AD <yume@yumechi.jp>
---
 mac/apparmor/yumechi-no-kuni-proxy-worker |  3 ++-
 src/lib.rs                                | 27 ++++++++++++++++++--
 src/main.rs                               | 30 +++++++++++------------
 3 files changed, 42 insertions(+), 18 deletions(-)

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<C: UpstreamClient, S: Sandboxing> {
 /// 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();