Terminate entire wasix process when a worker thread fails or calls exit

lib/wasix/src/bin_factory/exec.rs

@@ -279,6 +279,7 @@ fn call_module(
         if let Err(err) = call_ret {
             match err.downcast::<WasiError>() {
                 Ok(WasiError::Exit(code)) if code.is_success() => Ok(Errno::Success),
+                Ok(WasiError::ThreadExit) => Ok(Errno::Success),
                 Ok(WasiError::Exit(code)) => {
                     runtime.on_taint(TaintReason::NonZeroExitCode(code));
                     Err(WasiError::Exit(code).into())

lib/wasix/src/lib.rs

@@ -114,6 +114,8 @@ pub use crate::{
 pub enum WasiError {
     #[error("WASI exited with code: {0}")]
     Exit(ExitCode),
+    #[error("WASI thread exited")]
+    ThreadExit,
     #[error("WASI deep sleep: {0:?}")]
     DeepSleep(DeepSleepWork),
     #[error("The WASI version could not be determined")]

lib/wasix/src/runners/wasi.rs

@@ -404,6 +404,9 @@ impl crate::runners::Runner for WasiRunner {
                                     WasiRuntimeError::Wasi(WasiError::Exit(a)) => {
                                         WasiRuntimeError::Wasi(WasiError::Exit(*a))
                                     }
+                                    WasiRuntimeError::Wasi(WasiError::ThreadExit) => {
+                                        WasiRuntimeError::Wasi(WasiError::ThreadExit)
+                                    }
                                     WasiRuntimeError::Wasi(WasiError::UnknownWasiVersion) => {
                                         WasiRuntimeError::Wasi(WasiError::UnknownWasiVersion)
                                     }

lib/wasix/src/state/env.rs

@@ -1252,32 +1252,33 @@ impl WasiEnv {
 
     /// Cleans up all the open files (if this is the main thread)
     #[allow(clippy::await_holding_lock)]
-    pub fn blocking_on_exit(&self, exit_code: Option<ExitCode>) {
-        let cleanup = self.on_exit(exit_code);
+    pub fn blocking_on_exit(&self, process_exit_code: Option<ExitCode>) {
+        let cleanup = self.on_exit(process_exit_code);
         InlineWaker::block_on(cleanup);
     }
 
     /// Cleans up all the open files (if this is the main thread)
     #[allow(clippy::await_holding_lock)]
-    pub fn on_exit(&self, exit_code: Option<ExitCode>) -> BoxFuture<'static, ()> {
+    pub fn on_exit(&self, process_exit_code: Option<ExitCode>) -> BoxFuture<'static, ()> {
         const CLEANUP_TIMEOUT: Duration = Duration::from_secs(10);
 
         // If snap-shooting is enabled then we should record an event that the thread has exited.
         #[cfg(feature = "journal")]
         if self.should_journal() && self.has_active_journal() {
-            if let Err(err) = JournalEffector::save_thread_exit(self, self.tid(), exit_code) {
+            if let Err(err) = JournalEffector::save_thread_exit(self, self.tid(), process_exit_code)
+            {
                 tracing::warn!("failed to save snapshot event for thread exit - {}", err);
             }
 
             if self.thread.is_main() {
-                if let Err(err) = JournalEffector::save_process_exit(self, exit_code) {
+                if let Err(err) = JournalEffector::save_process_exit(self, process_exit_code) {
                     tracing::warn!("failed to save snapshot event for process exit - {}", err);
                 }
             }
         }
 
-        // If this is the main thread then also close all the files
-        if self.thread.is_main() {
+        // If the process wants to exit, also close all files and terminate it
+        if let Some(process_exit_code) = process_exit_code {
             let process = self.process.clone();
             let disable_fs_cleanup = self.disable_fs_cleanup;
             let pid = self.pid();
@@ -1303,8 +1304,7 @@ impl WasiEnv {
                 }
 
                 // Terminate the process
-                let exit_code = exit_code.unwrap_or_else(|| Errno::Canceled.into());
-                process.terminate(exit_code);
+                process.terminate(process_exit_code);
             })
         } else {
             Box::pin(async {})

lib/wasix/src/state/func_env.rs

@@ -284,15 +284,15 @@ impl WasiFunctionEnv {
     /// This function should only be called from within a syscall
     /// as it can potentially execute local thread variable cleanup
     /// code
-    pub fn on_exit(&self, store: &mut impl AsStoreMut, exit_code: Option<ExitCode>) {
+    pub fn on_exit(&self, store: &mut impl AsStoreMut, process_exit_code: Option<ExitCode>) {
         trace!(
             "wasi[{}:{}]::on_exit",
             self.data(store).pid(),
             self.data(store).tid()
         );
 
         // Cleans up all the open files (if this is the main thread)
-        self.data(store).blocking_on_exit(exit_code);
+        self.data(store).blocking_on_exit(process_exit_code);
     }
 
     /// Bootstraps this main thread and context with any journals that

lib/wasix/src/syscalls/wasix/thread_exit.rs

@@ -3,15 +3,14 @@ use crate::syscalls::*;
 
 /// ### `thread_exit()`
 /// Terminates the current running thread, if this is the last thread then
-/// the process will also exit with the specified exit code. An exit code
-/// of 0 indicates successful termination of the thread. The meanings of
-/// other values is dependent on the environment.
+/// the process will also exit with code 0.
+/// The exit code parameter is a left over from a previous version of this
+/// syscall, maintained here to keep the syscall backwards-compatible, but
+/// is otherwise unused.
 ///
-/// ## Parameters
-///
-/// * `rval` - The exit code returned by the process.
-#[instrument(level = "trace", skip_all, fields(%exitcode), ret)]
-pub fn thread_exit(ctx: FunctionEnvMut<'_, WasiEnv>, exitcode: ExitCode) -> Result<(), WasiError> {
-    tracing::debug!(tid=%ctx.data().thread.id(), %exitcode);
-    Err(WasiError::Exit(exitcode))
+/// This syscall does not return.
+#[instrument(level = "trace", skip_all, fields(%_exitcode), ret)]
+pub fn thread_exit(ctx: FunctionEnvMut<'_, WasiEnv>, _exitcode: ExitCode) -> Result<(), WasiError> {
+    tracing::debug!(tid=%ctx.data().thread.id(), "thread exit");
+    Err(WasiError::ThreadExit)
 }

lib/wasix/src/syscalls/wasix/thread_spawn.rs

@@ -203,10 +203,19 @@ fn call_module<M: MemorySize>(
                 .map_err(|_| Errno::Overflow)
                 .unwrap(),
         );
+        trace!("callback finished (ret={:?})", call_ret);
+
         let mut ret = Errno::Success;
+        let mut exit_code = None;
         if let Err(err) = call_ret {
             match err.downcast::<WasiError>() {
+                Ok(WasiError::ThreadExit) => {
+                    trace!("thread exited cleanly");
+                    ret = Errno::Success;
+                }
                 Ok(WasiError::Exit(code)) => {
+                    trace!(exit_code = ?code, "thread requested exit");
+                    exit_code = Some(code);
                     ret = if code.is_success() {
                         Errno::Success
                     } else {
@@ -226,20 +235,23 @@ fn call_module<M: MemorySize>(
                         .runtime
                         .on_taint(TaintReason::UnknownWasiVersion);
                     ret = Errno::Noexec;
+                    exit_code = Some(ExitCode::Other(128 + ret as i32));
                 }
                 Err(err) => {
                     debug!("failed with runtime error: {}", err);
                     env.data(&store)
                         .runtime
                         .on_taint(TaintReason::RuntimeError(err));
                     ret = Errno::Noexec;
+                    exit_code = Some(ExitCode::Other(128 + ret as i32));
                 }
             }
+        } else {
+            debug!("thread exited cleanly without calling thread_exit");
         }
-        trace!("callback finished (ret={})", ret);
 
         // Clean up the environment
-        env.on_exit(store, Some(ret.into()));
+        env.on_exit(store, exit_code);
 
         // Return the result
         Ok(ret as u32)

tests/integration/cli/tests/snapshot.rs

@@ -1281,3 +1281,81 @@ fn test_snapshot_fs_rename() {
         .run_wasm(include_bytes!("./wasm/fs-rename.wasm"));
     assert_json_snapshot!(snapshot);
 }
+
+#[cfg_attr(any(target_env = "musl", target_os = "windows"), ignore)]
+#[test]
+fn test_snapshot_exit_0_from_main() {
+    let snapshot = TestBuilder::new()
+        .with_name(function!())
+        .run_wasm(include_bytes!("./wasm/exit-0-from-main.wasm"));
+    assert!(matches!(
+        snapshot.result,
+        TestResult::Success(TestOutput { exit_code: 0, .. })
+    ));
+    assert_json_snapshot!(snapshot);
+}
+
+#[cfg_attr(any(target_env = "musl", target_os = "windows"), ignore)]
+#[test]
+fn test_snapshot_exit_1_from_main() {
+    let snapshot = TestBuilder::new()
+        .with_name(function!())
+        .run_wasm(include_bytes!("./wasm/exit-1-from-main.wasm"));
+    assert!(matches!(
+        snapshot.result,
+        TestResult::Success(TestOutput { exit_code: 1, .. })
+    ));
+    assert_json_snapshot!(snapshot);
+}
+
+#[cfg_attr(any(target_env = "musl", target_os = "windows"), ignore)]
+#[test]
+fn test_snapshot_exit_0_from_worker() {
+    let snapshot = TestBuilder::new()
+        .with_name(function!())
+        .run_wasm(include_bytes!("./wasm/exit-0-from-worker.wasm"));
+    assert!(matches!(
+        snapshot.result,
+        TestResult::Success(TestOutput { exit_code: 0, .. })
+    ));
+    assert_json_snapshot!(snapshot);
+}
+
+#[cfg_attr(any(target_env = "musl", target_os = "windows"), ignore)]
+#[test]
+fn test_snapshot_exit_1_from_worker() {
+    let snapshot = TestBuilder::new()
+        .with_name(function!())
+        .run_wasm(include_bytes!("./wasm/exit-1-from-worker.wasm"));
+    assert!(matches!(
+        snapshot.result,
+        TestResult::Success(TestOutput { exit_code: 1, .. })
+    ));
+    assert_json_snapshot!(snapshot);
+}
+
+#[cfg_attr(any(target_env = "musl", target_os = "windows"), ignore)]
+#[test]
+fn test_snapshot_worker_terminating_normally() {
+    let snapshot = TestBuilder::new()
+        .with_name(function!())
+        .run_wasm(include_bytes!("./wasm/worker-terminating-normally.wasm"));
+    assert!(matches!(
+        snapshot.result,
+        TestResult::Success(TestOutput { exit_code: 0, .. })
+    ));
+    assert_json_snapshot!(snapshot);
+}
+
+#[cfg_attr(any(target_env = "musl", target_os = "windows"), ignore)]
+#[test]
+fn test_snapshot_worker_panicking() {
+    let snapshot = TestBuilder::new()
+        .with_name(function!())
+        .run_wasm(include_bytes!("./wasm/worker-panicking.wasm"));
+    assert!(matches!(
+        snapshot.result,
+        TestResult::Success(TestOutput { exit_code: 173, .. })
+    ));
+    assert_json_snapshot!(snapshot);
+}

tests/integration/cli/tests/snapshots/snapshot__snapshot_exit_0_from_main.snap

@@ -0,0 +1,21 @@
+---
+source: tests/integration/cli/tests/snapshot.rs
+expression: snapshot
+---
+{
+  "spec": {
+    "name": "snapshot::test_snapshot_exit_0_from_main",
+    "use_packages": [],
+    "include_webcs": [],
+    "cli_args": [],
+    "enable_threads": true,
+    "enable_network": false
+  },
+  "result": {
+    "Success": {
+      "stdout": "",
+      "stderr": "",
+      "exit_code": 0
+    }
+  }
+}

tests/integration/cli/tests/snapshots/snapshot__snapshot_exit_0_from_worker.snap

@@ -0,0 +1,22 @@
+---
+source: tests/integration/cli/tests/snapshot.rs
+assertion_line: 1309
+expression: snapshot
+---
+{
+  "spec": {
+    "name": "snapshot::test_snapshot_exit_0_from_worker",
+    "use_packages": [],
+    "include_webcs": [],
+    "cli_args": [],
+    "enable_threads": true,
+    "enable_network": false
+  },
+  "result": {
+    "Success": {
+      "stdout": "",
+      "stderr": "",
+      "exit_code": 0
+    }
+  }
+}

tests/integration/cli/tests/snapshots/snapshot__snapshot_exit_1_from_main.snap

@@ -0,0 +1,22 @@
+---
+source: tests/integration/cli/tests/snapshot.rs
+assertion_line: 1300
+expression: snapshot
+---
+{
+  "spec": {
+    "name": "snapshot::test_snapshot_exit_1_from_main",
+    "use_packages": [],
+    "include_webcs": [],
+    "cli_args": [],
+    "enable_threads": true,
+    "enable_network": false
+  },
+  "result": {
+    "Success": {
+      "stdout": "",
+      "stderr": "",
+      "exit_code": 1
+    }
+  }
+}

tests/integration/cli/tests/snapshots/snapshot__snapshot_exit_1_from_worker.snap

@@ -0,0 +1,22 @@
+---
+source: tests/integration/cli/tests/snapshot.rs
+assertion_line: 1318
+expression: snapshot
+---
+{
+  "spec": {
+    "name": "snapshot::test_snapshot_exit_1_from_worker",
+    "use_packages": [],
+    "include_webcs": [],
+    "cli_args": [],
+    "enable_threads": true,
+    "enable_network": false
+  },
+  "result": {
+    "Success": {
+      "stdout": "",
+      "stderr": "",
+      "exit_code": 1
+    }
+  }
+}

tests/integration/cli/tests/snapshots/snapshot__snapshot_worker_panicking.snap

@@ -0,0 +1,22 @@
+---
+source: tests/integration/cli/tests/snapshot.rs
+assertion_line: 1336
+expression: snapshot
+---
+{
+  "spec": {
+    "name": "snapshot::test_snapshot_worker_panicking",
+    "use_packages": [],
+    "include_webcs": [],
+    "cli_args": [],
+    "enable_threads": true,
+    "enable_network": false
+  },
+  "result": {
+    "Success": {
+      "stdout": "",
+      "stderr": "thread '<unnamed>' panicked at src/bin/worker-panicking.rs:3:9:\nchild thread panicking\nnote: run with `RUST_BACKTRACE=1` environment variable to display a backtrace\n",
+      "exit_code": 173
+    }
+  }
+}

tests/integration/cli/tests/snapshots/snapshot__snapshot_worker_terminating_normally.snap

@@ -0,0 +1,22 @@
+---
+source: tests/integration/cli/tests/snapshot.rs
+assertion_line: 1327
+expression: snapshot
+---
+{
+  "spec": {
+    "name": "snapshot::test_snapshot_worker_terminating_normally",
+    "use_packages": [],
+    "include_webcs": [],
+    "cli_args": [],
+    "enable_threads": true,
+    "enable_network": false
+  },
+  "result": {
+    "Success": {
+      "stdout": "In child thread\nIn main thread\n",
+      "stderr": "",
+      "exit_code": 0
+    }
+  }
+}

tests/rust-wasi-tests/src/bin/exit-0-from-main.rs

@@ -0,0 +1,3 @@
+fn main() {
+    std::process::exit(0);
+}

tests/rust-wasi-tests/src/bin/exit-0-from-worker.rs

@@ -0,0 +1,3 @@
+fn main() {
+    std::thread::spawn(|| std::process::exit(0)).join().unwrap();
+}

tests/rust-wasi-tests/src/bin/exit-1-from-main.rs

@@ -0,0 +1,3 @@
+fn main() {
+    std::process::exit(1);
+}