review: catch host errors and exit all threads

Co-authored-by: Alex Crichton <[email protected]>

Author: abrown

crates/wasi-threads/src/lib.rs

@@ -2,10 +2,12 @@
 //!
 //! [`wasi-threads`]: https://github.com/WebAssembly/wasi-threads
 
-use anyhow::{anyhow, Result};
+use anyhow::{anyhow, bail, Result};
 use rand::Rng;
+use std::panic::{catch_unwind, AssertUnwindSafe};
 use std::sync::Arc;
-use wasmtime::{Caller, Linker, Module, SharedMemory, Store};
+use std::thread;
+use wasmtime::{Caller, Linker, Module, SharedMemory, Store, ValType};
 use wasmtime_wasi::maybe_exit_on_error;
 
 // This name is a function export designated by the wasi-threads specification:
@@ -36,46 +38,50 @@ impl<T: Clone + Send + 'static> WasiThreadsCtx<T> {
         let wasi_thread_id = random_thread_id();
         let builder = thread::Builder::new().name(format!("wasi-thread-{}", wasi_thread_id));
         builder.spawn(move || {
-            // Convenience function for printing failures, since the `Thread`
-            // has no way to report a failure to the outer context.
-            let fail = |msg: String| {
-                format!(
-                    "wasi-thread-{} exited unsuccessfully: {}",
-                    wasi_thread_id, msg
-                )
-            };
+            // Catch any panic failures in host code; e.g., if a WASI module
+            // were to crash, we want all threads to exit, not just this one.
+            let result = catch_unwind(AssertUnwindSafe(|| {
+                // Each new instance is created in its own store.
+                let mut store = Store::new(&module.engine(), host);
 
-            // Each new instance is created in its own store.
-            let mut store = Store::new(&module.engine(), host);
-            let instance = linker
-                .instantiate(&mut store, &module)
-                .expect(&fail("failed to instantiate".into()));
-            let thread_entry_point = instance
-                .get_typed_func::<(i32, i32), ()>(&mut store, WASI_ENTRY_POINT)
-                .expect(&fail(format!(
-                    "failed to find wasi-threads entry point function: {}",
-                    WASI_ENTRY_POINT
-                )));
+                // Ideally, we would have already checked much earlier (e.g.,
+                // `new`) whether the module can be instantiated. Because
+                // `Linker::instantiate_pre` requires a `Store` and that is only
+                // available now. TODO:
+                // https://github.com/bytecodealliance/wasmtime/issues/5675.
+                let instance = linker.instantiate(&mut store, &module).expect(&format!(
+                    "wasi-thread-{} exited unsuccessfully: failed to instantiate",
+                    wasi_thread_id
+                ));
+                let thread_entry_point = instance
+                    .get_typed_func::<(i32, i32), ()>(&mut store, WASI_ENTRY_POINT)
+                    .unwrap();
 
-            // Start the thread's entry point. Any traps or calls to
-            // `proc_exit`, by specification, should end execution for all
-            // threads. This code uses `process::exit` to do so, which is what
-            // the user expects from the CLI but probably not in a Wasmtime
-            // embedding.
-            log::trace!(
-                "spawned thread id = {}; calling start function `{}` with: {}",
-                wasi_thread_id,
-                WASI_ENTRY_POINT,
-                thread_start_arg
-            );
-            match thread_entry_point.call(&mut store, (wasi_thread_id, thread_start_arg)) {
-                Ok(_) => log::trace!("exiting thread id = {} normally", wasi_thread_id),
-                Err(e) => {
-                    log::trace!("exiting thread id = {} due to error", wasi_thread_id);
-                    let e = maybe_exit_on_error(e);
-                    eprintln!("Error: {:?}", e);
-                    std::process::exit(1);
+                // Start the thread's entry point. Any traps or calls to
+                // `proc_exit`, by specification, should end execution for all
+                // threads. This code uses `process::exit` to do so, which is what
+                // the user expects from the CLI but probably not in a Wasmtime
+                // embedding.
+                log::trace!(
+                    "spawned thread id = {}; calling start function `{}` with: {}",
+                    wasi_thread_id,
+                    WASI_ENTRY_POINT,
+                    thread_start_arg
+                );
+                match thread_entry_point.call(&mut store, (wasi_thread_id, thread_start_arg)) {
+                    Ok(_) => log::trace!("exiting thread id = {} normally", wasi_thread_id),
+                    Err(e) => {
+                        log::trace!("exiting thread id = {} due to error", wasi_thread_id);
+                        let e = maybe_exit_on_error(e);
+                        eprintln!("Error: {:?}", e);
+                        std::process::exit(1);
+                    }
                 }
+            }));
+
+            if let Err(e) = result {
+                eprintln!("Error: {:?}", e);
+                std::process::exit(1);
             }
         })?;