1 //! Implement [`wasi-threads`].
2 //!
3 //! [`wasi-threads`]: https://github.com/WebAssembly/wasi-threads
4 
5 use anyhow::{anyhow, Result};
6 use std::panic::{catch_unwind, AssertUnwindSafe};
7 use std::sync::atomic::{AtomicI32, Ordering};
8 use std::sync::Arc;
9 use std::thread;
10 use wasmtime::{Caller, ExternType, InstancePre, Linker, Module, SharedMemory, Store, ValType};
11 
12 // This name is a function export designated by the wasi-threads specification:
13 // https://github.com/WebAssembly/wasi-threads/#detailed-design-discussion
14 const WASI_ENTRY_POINT: &str = "wasi_thread_start";
15 
16 pub struct WasiThreadsCtx<T> {
17     instance_pre: Arc<InstancePre<T>>,
18     tid: AtomicI32,
19 }
20 
21 impl<T: Clone + Send + 'static> WasiThreadsCtx<T> {
22     pub fn new(module: Module, linker: Arc<Linker<T>>) -> Result<Self> {
23         let instance_pre = Arc::new(linker.instantiate_pre(&module)?);
24         let tid = AtomicI32::new(0);
25         Ok(Self { instance_pre, tid })
26     }
27 
28     pub fn spawn(&self, host: T, thread_start_arg: i32) -> Result<i32> {
29         let instance_pre = self.instance_pre.clone();
30 
31         // Check that the thread entry point is present. Why here? If we check
32         // for this too early, then we cannot accept modules that do not have an
33         // entry point but never spawn a thread. As pointed out in
34         // https://github.com/bytecodealliance/wasmtime/issues/6153, checking
35         // the entry point here allows wasi-threads to be compatible with more
36         // modules.
37         //
38         // As defined in the wasi-threads specification, returning a negative
39         // result here indicates to the guest module that the spawn failed.
40         if !has_entry_point(instance_pre.module()) {
41             log::error!("failed to find a wasi-threads entry point function; expected an export with name: {WASI_ENTRY_POINT}");
42             return Ok(-1);
43         }
44         if !has_correct_signature(instance_pre.module()) {
45             log::error!("the exported entry point function has an incorrect signature: expected `(i32, i32) -> ()`");
46             return Ok(-1);
47         }
48 
49         let wasi_thread_id = self.next_thread_id();
50         if wasi_thread_id.is_none() {
51             log::error!("ran out of valid thread IDs");
52             return Ok(-1);
53         }
54         let wasi_thread_id = wasi_thread_id.unwrap();
55 
56         // Start a Rust thread running a new instance of the current module.
57         let builder = thread::Builder::new().name(format!("wasi-thread-{}", wasi_thread_id));
58         builder.spawn(move || {
59             // Catch any panic failures in host code; e.g., if a WASI module
60             // were to crash, we want all threads to exit, not just this one.
61             let result = catch_unwind(AssertUnwindSafe(|| {
62                 // Each new instance is created in its own store.
63                 let mut store = Store::new(&instance_pre.module().engine(), host);
64                 let instance = instance_pre.instantiate(&mut store).unwrap();
65                 let thread_entry_point = instance
66                     .get_typed_func::<(i32, i32), ()>(&mut store, WASI_ENTRY_POINT)
67                     .unwrap();
68 
69                 // Start the thread's entry point. Any traps or calls to
70                 // `proc_exit`, by specification, should end execution for all
71                 // threads. This code uses `process::exit` to do so, which is
72                 // what the user expects from the CLI but probably not in a
73                 // Wasmtime embedding.
74                 log::trace!(
75                     "spawned thread id = {}; calling start function `{}` with: {}",
76                     wasi_thread_id,
77                     WASI_ENTRY_POINT,
78                     thread_start_arg
79                 );
80                 match thread_entry_point.call(&mut store, (wasi_thread_id, thread_start_arg)) {
81                     Ok(_) => log::trace!("exiting thread id = {} normally", wasi_thread_id),
82                     Err(e) => {
83                         log::trace!("exiting thread id = {} due to error", wasi_thread_id);
84                         let e = wasi_common::maybe_exit_on_error(e);
85                         eprintln!("Error: {:?}", e);
86                         std::process::exit(1);
87                     }
88                 }
89             }));
90 
91             if let Err(e) = result {
92                 eprintln!("wasi-thread-{} panicked: {:?}", wasi_thread_id, e);
93                 std::process::exit(1);
94             }
95         })?;
96 
97         Ok(wasi_thread_id)
98     }
99 
100     /// Helper for generating valid WASI thread IDs (TID).
101     ///
102     /// Callers of `wasi_thread_spawn` expect a TID in range of 0 < TID <= 0x1FFFFFFF
103     /// to indicate a successful spawning of the thread whereas a negative
104     /// return value indicates an failure to spawn.
105     fn next_thread_id(&self) -> Option<i32> {
106         match self
107             .tid
108             .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| match v {
109                 ..=0x1ffffffe => Some(v + 1),
110                 _ => None,
111             }) {
112             Ok(v) => Some(v + 1),
113             Err(_) => None,
114         }
115     }
116 }
117 
118 /// Manually add the WASI `thread_spawn` function to the linker.
119 ///
120 /// It is unclear what namespace the `wasi-threads` proposal should live under:
121 /// it is not clear if it should be included in any of the `preview*` releases
122 /// so for the time being its module namespace is simply `"wasi"` (TODO).
123 pub fn add_to_linker<T: Clone + Send + 'static>(
124     linker: &mut wasmtime::Linker<T>,
125     store: &wasmtime::Store<T>,
126     module: &Module,
127     get_cx: impl Fn(&mut T) -> &WasiThreadsCtx<T> + Send + Sync + Copy + 'static,
128 ) -> anyhow::Result<()> {
129     linker.func_wrap(
130         "wasi",
131         "thread-spawn",
132         move |mut caller: Caller<'_, T>, start_arg: i32| -> i32 {
133             log::trace!("new thread requested via `wasi::thread_spawn` call");
134             let host = caller.data().clone();
135             let ctx = get_cx(caller.data_mut());
136             match ctx.spawn(host, start_arg) {
137                 Ok(thread_id) => {
138                     assert!(thread_id >= 0, "thread_id = {}", thread_id);
139                     thread_id
140                 }
141                 Err(e) => {
142                     log::error!("failed to spawn thread: {}", e);
143                     -1
144                 }
145             }
146         },
147     )?;
148 
149     // Find the shared memory import and satisfy it with a newly-created shared
150     // memory import.
151     for import in module.imports() {
152         if let Some(m) = import.ty().memory() {
153             if m.is_shared() {
154                 let mem = SharedMemory::new(module.engine(), m.clone())?;
155                 linker.define(store, import.module(), import.name(), mem.clone())?;
156             } else {
157                 return Err(anyhow!(
158                     "memory was not shared; a `wasi-threads` must import \
159                      a shared memory as \"memory\""
160                 ));
161             }
162         }
163     }
164     Ok(())
165 }
166 
167 /// Check if wasi-threads' `wasi_thread_start` export is present.
168 fn has_entry_point(module: &Module) -> bool {
169     module.get_export(WASI_ENTRY_POINT).is_some()
170 }
171 
172 /// Check if the entry function has the correct signature `(i32, i32) -> ()`.
173 fn has_correct_signature(module: &Module) -> bool {
174     use ValType::*;
175     match module.get_export(WASI_ENTRY_POINT) {
176         Some(ExternType::Func(ty)) => ty.params().eq([I32, I32]) && ty.results().len() == 0,
177         _ => false,
178     }
179 }
180