From c37723b0adbe5cbca9382533ff559f20fca8d819 Mon Sep 17 00:00:00 2001
From: Johnathan Sharratt <johnathan.sharratt@gmail.com>
Date: Sat, 25 May 2024 17:02:07 +1000
Subject: [PATCH] Fix for interrupts

---
 lib/wasix/src/os/task/thread.rs |  3 ++
 lib/wasix/src/state/env.rs      |  6 +++-
 lib/wasix/src/syscalls/mod.rs   | 50 ++++++++++++++++++++-------------
 3 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/lib/wasix/src/os/task/thread.rs b/lib/wasix/src/os/task/thread.rs
index 69837851f39..293addc0146 100644
--- a/lib/wasix/src/os/task/thread.rs
+++ b/lib/wasix/src/os/task/thread.rs
@@ -341,6 +341,9 @@ impl WasiThread {
 
     /// Adds a signal for this thread to process
     pub fn signal(&self, signal: Signal) {
+        let tid = self.tid();
+        tracing::trace!(%tid, "signal-thread({:?})", signal);
+
         let mut guard = self.state.signals.lock().unwrap();
         if !guard.0.contains(&signal) {
             guard.0.push(signal);
diff --git a/lib/wasix/src/state/env.rs b/lib/wasix/src/state/env.rs
index a43c9d29155..a626aa09236 100644
--- a/lib/wasix/src/state/env.rs
+++ b/lib/wasix/src/state/env.rs
@@ -806,7 +806,7 @@ impl WasiEnv {
                 tracing::trace!(
                     pid=%ctx.data().pid(),
                     ?signal,
-                    "Processing signal",
+                    "processing signal via handler",
                 );
                 if let Err(err) = handler.call(ctx, signal as i32) {
                     match err.downcast::<WasiError>() {
@@ -832,6 +832,10 @@ impl WasiEnv {
                         }
                     }
                 }
+                tracing::trace!(
+                    pid=%ctx.data().pid(),
+                    "signal processed",
+                );
             }
             Ok(true)
         } else {
diff --git a/lib/wasix/src/syscalls/mod.rs b/lib/wasix/src/syscalls/mod.rs
index f4b1ee10911..90e75752fa0 100644
--- a/lib/wasix/src/syscalls/mod.rs
+++ b/lib/wasix/src/syscalls/mod.rs
@@ -358,7 +358,6 @@ struct AsyncifyPoller<'a, 'b, 'c, T, Fut>
 where
     Fut: Future<Output = T> + Send + Sync + 'static,
 {
-    process_signals: bool,
     ctx: &'b mut FunctionEnvMut<'c, WasiEnv>,
     work: &'a mut Pin<Box<Fut>>,
 }
@@ -380,18 +379,37 @@ where
                 Errno::Child.into()
             }))));
         }
-        if self.process_signals && env.thread.has_signals_or_subscribe(cx.waker()) {
-            let signals = env.thread.signals().lock().unwrap();
-            for sig in signals.0.iter() {
-                if *sig == Signal::Sigint
-                    || *sig == Signal::Sigquit
-                    || *sig == Signal::Sigkill
-                    || *sig == Signal::Sigabrt
-                {
-                    let exit_code = env.thread.set_or_get_exit_code_for_signal(*sig);
-                    return Poll::Ready(Err(WasiError::Exit(exit_code)));
+        if env.thread.has_signals_or_subscribe(cx.waker()) {
+            let has_exit = {
+                let signals = env.thread.signals().lock().unwrap();
+                signals
+                    .0
+                    .iter()
+                    .filter_map(|sig| {
+                        if *sig == Signal::Sigint
+                            || *sig == Signal::Sigquit
+                            || *sig == Signal::Sigkill
+                            || *sig == Signal::Sigabrt
+                        {
+                            Some(env.thread.set_or_get_exit_code_for_signal(*sig))
+                        } else {
+                            None
+                        }
+                    })
+                    .next()
+            };
+
+            return match WasiEnv::process_signals_and_exit(self.ctx) {
+                Ok(Ok(_)) => {
+                    if let Some(exit_code) = has_exit {
+                        Poll::Ready(Err(WasiError::Exit(exit_code)))
+                    } else {
+                        Poll::Pending
+                    }
                 }
-            }
+                Ok(Err(err)) => Poll::Ready(Err(WasiError::Exit(ExitCode::Errno(err)))),
+                Err(err) => Poll::Ready(Err(err)),
+            };
         }
         Poll::Pending
     }
@@ -465,13 +483,6 @@ where
         false => Duration::from_millis(50),
     };
 
-    // Determine if we should process signals or now
-    let process_signals = ctx
-        .data()
-        .try_inner()
-        .map(|i| !i.signal_set)
-        .unwrap_or(true);
-
     // Box up the trigger
     let mut trigger = Box::pin(work);
 
@@ -498,7 +509,6 @@ where
         Ok(tokio::select! {
             // Inner wait with finializer
             res = AsyncifyPoller {
-                process_signals,
                 ctx: &mut ctx,
                 work: &mut trigger,
             } => {