wasmtime_wasi_threads/
lib.rs1use anyhow::{Result, anyhow};
6use std::panic::{AssertUnwindSafe, catch_unwind};
7use std::sync::Arc;
8use std::sync::atomic::{AtomicI32, Ordering};
9use std::thread;
10use wasmtime::{Caller, ExternType, InstancePre, Linker, Module, SharedMemory, Store};
11
12const WASI_ENTRY_POINT: &str = "wasi_thread_start";
15
16pub struct WasiThreadsCtx<T> {
17 instance_pre: Arc<InstancePre<T>>,
18 tid: AtomicI32,
19}
20
21impl<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 if !has_entry_point(instance_pre.module()) {
41 log::error!(
42 "failed to find a wasi-threads entry point function; expected an export with name: {WASI_ENTRY_POINT}"
43 );
44 return Ok(-1);
45 }
46 if !has_correct_signature(instance_pre.module()) {
47 log::error!(
48 "the exported entry point function has an incorrect signature: expected `(i32, i32) -> ()`"
49 );
50 return Ok(-1);
51 }
52
53 let wasi_thread_id = self.next_thread_id();
54 if wasi_thread_id.is_none() {
55 log::error!("ran out of valid thread IDs");
56 return Ok(-1);
57 }
58 let wasi_thread_id = wasi_thread_id.unwrap();
59
60 let builder = thread::Builder::new().name(format!("wasi-thread-{wasi_thread_id}"));
62 builder.spawn(move || {
63 let result = catch_unwind(AssertUnwindSafe(|| {
66 let mut store = Store::new(&instance_pre.module().engine(), host);
68
69 let instance = if instance_pre.module().engine().is_async() {
70 wasmtime_wasi::runtime::in_tokio(instance_pre.instantiate_async(&mut store))
71 } else {
72 instance_pre.instantiate(&mut store)
73 }
74 .unwrap();
75
76 let thread_entry_point = instance
77 .get_typed_func::<(i32, i32), ()>(&mut store, WASI_ENTRY_POINT)
78 .unwrap();
79
80 log::trace!(
86 "spawned thread id = {wasi_thread_id}; calling start function `{WASI_ENTRY_POINT}` with: {thread_start_arg}"
87 );
88 let res = if instance_pre.module().engine().is_async() {
89 wasmtime_wasi::runtime::in_tokio(
90 thread_entry_point
91 .call_async(&mut store, (wasi_thread_id, thread_start_arg)),
92 )
93 } else {
94 thread_entry_point.call(&mut store, (wasi_thread_id, thread_start_arg))
95 };
96 match res {
97 Ok(_) => log::trace!("exiting thread id = {wasi_thread_id} normally"),
98 Err(e) => {
99 log::trace!("exiting thread id = {wasi_thread_id} due to error");
100 let e = wasi_common::maybe_exit_on_error(e);
101 eprintln!("Error: {e:?}");
102 std::process::exit(1);
103 }
104 }
105 }));
106
107 if let Err(e) = result {
108 eprintln!("wasi-thread-{wasi_thread_id} panicked: {e:?}");
109 std::process::exit(1);
110 }
111 })?;
112
113 Ok(wasi_thread_id)
114 }
115
116 fn next_thread_id(&self) -> Option<i32> {
122 match self
123 .tid
124 .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| match v {
125 ..=0x1ffffffe => Some(v + 1),
126 _ => None,
127 }) {
128 Ok(v) => Some(v + 1),
129 Err(_) => None,
130 }
131 }
132}
133
134pub fn add_to_linker<T: Clone + Send + 'static>(
140 linker: &mut wasmtime::Linker<T>,
141 store: &wasmtime::Store<T>,
142 module: &Module,
143 get_cx: impl Fn(&mut T) -> &WasiThreadsCtx<T> + Send + Sync + Copy + 'static,
144) -> anyhow::Result<()> {
145 linker.func_wrap(
146 "wasi",
147 "thread-spawn",
148 move |mut caller: Caller<'_, T>, start_arg: i32| -> i32 {
149 log::trace!("new thread requested via `wasi::thread_spawn` call");
150 let host = caller.data().clone();
151 let ctx = get_cx(caller.data_mut());
152 match ctx.spawn(host, start_arg) {
153 Ok(thread_id) => {
154 assert!(thread_id >= 0, "thread_id = {thread_id}");
155 thread_id
156 }
157 Err(e) => {
158 log::error!("failed to spawn thread: {e}");
159 -1
160 }
161 }
162 },
163 )?;
164
165 for import in module.imports() {
168 if let Some(m) = import.ty().memory() {
169 if m.is_shared() {
170 let mem = SharedMemory::new(module.engine(), m.clone())?;
171 linker.define(store, import.module(), import.name(), mem.clone())?;
172 } else {
173 return Err(anyhow!(
174 "memory was not shared; a `wasi-threads` must import \
175 a shared memory as \"memory\""
176 ));
177 }
178 }
179 }
180 Ok(())
181}
182
183fn has_entry_point(module: &Module) -> bool {
185 module.get_export(WASI_ENTRY_POINT).is_some()
186}
187
188fn has_correct_signature(module: &Module) -> bool {
190 match module.get_export(WASI_ENTRY_POINT) {
191 Some(ExternType::Func(ty)) => {
192 ty.params().len() == 2
193 && ty.params().nth(0).unwrap().is_i32()
194 && ty.params().nth(1).unwrap().is_i32()
195 && ty.results().len() == 0
196 }
197 _ => false,
198 }
199}