wasmer_wasix/runtime/task_manager/
mod.rs

1// TODO: should be behind a different , tokio specific feature flag.
2#[cfg(feature = "sys-thread")]
3pub mod tokio;
4
5use std::ops::Deref;
6use std::task::{Context, Poll};
7use std::{pin::Pin, time::Duration};
8
9use bytes::Bytes;
10use derive_more::Debug;
11use futures::future::BoxFuture;
12use futures::{Future, TryFutureExt};
13use wasmer::{AsStoreMut, Memory, MemoryType, Module, SharedMemory, Store, StoreMut};
14use wasmer_wasix_types::wasi::{Errno, ExitCode};
15
16use crate::os::task::thread::WasiThreadError;
17use crate::{StoreSnapshot, WasiEnv, WasiFunctionEnv, WasiThread, capture_store_snapshot};
18use crate::{state::PreparedInstanceGroupData, syscalls::AsyncifyFuture};
19
20pub use virtual_mio::waker::*;
21
22#[derive(Debug)]
23pub enum SpawnType {
24    CreateMemory,
25    CreateMemoryOfType(MemoryType),
26    AttachMemory(SharedMemory),
27    #[debug("NewLinkerInstanceGroup(..)")]
28    NewLinkerInstanceGroup(PreparedInstanceGroupData),
29}
30
31/// Describes whether a new memory should be created (and, in case, its type) or if it was already
32/// created and the store it belongs to.
33///
34/// # Note
35///
36/// This type is necessary for now because we can't pass a [`wasmer::StoreRef`] between threads, so this
37/// conceptually is a Send-able [`SpawnMemoryTypeOrStore`].
38pub enum SpawnMemoryTypeOrStore {
39    New,
40    Type(wasmer::MemoryType),
41    StoreAndMemory(wasmer::Store, Memory),
42}
43
44pub type WasmResumeTask = dyn FnOnce(WasiFunctionEnv, Store, Bytes) + Send + 'static;
45
46pub type WasmResumeTrigger = dyn FnOnce() -> Pin<Box<dyn Future<Output = Result<Bytes, ExitCode>> + Send + 'static>>
47    + Send
48    + Sync;
49
50/// The properties passed to the task
51#[derive(derive_more::Debug)]
52pub struct TaskWasmRunProperties {
53    pub ctx: WasiFunctionEnv,
54    pub store: Store,
55    /// The result of the asynchronous trigger serialized into bytes using the bincode serializer
56    /// When no trigger is associated with the run operation (i.e. spawning threads) then this will be None.
57    /// (if the trigger returns an ExitCode then the WASM process will be terminated without resuming)
58    pub trigger_result: Option<Result<Bytes, ExitCode>>,
59    /// The instance will be recycled back to this function when the WASM run has finished
60    #[debug(ignore)]
61    pub recycle: Option<Box<TaskWasmRecycle>>,
62}
63
64pub type TaskWasmPreRun = dyn (for<'a> FnOnce(
65        &'a mut WasiFunctionEnv,
66        &'a mut Store,
67    ) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>>)
68    + Send;
69
70/// Callback that will be invoked
71pub type TaskWasmRun = dyn FnOnce(TaskWasmRunProperties) + Send + 'static;
72
73/// Callback that will be invoked
74pub type TaskExecModule = dyn FnOnce(Module) + Send + 'static;
75
76/// The properties passed to the task
77#[derive(Debug)]
78pub struct TaskWasmRecycleProperties {
79    pub env: WasiEnv,
80    pub memory: Memory,
81    pub store: Store,
82}
83
84/// Callback that will be invoked
85pub type TaskWasmRecycle = dyn FnOnce(TaskWasmRecycleProperties) + Send + 'static;
86
87pub struct TaskWasmCallbacks {
88    pub run: Box<TaskWasmRun>,
89    pub recycle: Option<Box<TaskWasmRecycle>>,
90    pub pre_run: Option<Box<TaskWasmPreRun>>,
91    pub trigger: Option<Box<WasmResumeTrigger>>,
92}
93
94/// Represents a WASM task that will be executed on a dedicated thread
95pub struct TaskWasm {
96    pub callbacks: TaskWasmCallbacks,
97    pub env: WasiEnv,
98    pub module: Module,
99    pub globals: Option<StoreSnapshot>,
100    pub spawn_type: SpawnType,
101    pub update_layout: bool,
102    pub call_initialize: bool,
103}
104
105impl TaskWasm {
106    pub fn new(
107        run: Box<TaskWasmRun>,
108        env: WasiEnv,
109        module: Module,
110        update_layout: bool,
111        call_initialize: bool,
112    ) -> Self {
113        let shared_memory = module.imports().memories().next().map(|a| *a.ty());
114        Self {
115            callbacks: TaskWasmCallbacks {
116                run,
117                recycle: None,
118                pre_run: None,
119                trigger: None,
120            },
121            env,
122            module,
123            globals: None,
124            spawn_type: match shared_memory {
125                Some(ty) => SpawnType::CreateMemoryOfType(ty),
126                None => SpawnType::CreateMemory,
127            },
128            update_layout,
129            call_initialize,
130        }
131    }
132
133    pub fn with_memory(mut self, spawn_type: SpawnType) -> Self {
134        self.spawn_type = spawn_type;
135        self
136    }
137
138    pub fn with_optional_memory(mut self, spawn_type: Option<SpawnType>) -> Self {
139        if let Some(spawn_type) = spawn_type {
140            self.spawn_type = spawn_type;
141        }
142        self
143    }
144
145    pub fn with_globals(mut self, snapshot: StoreSnapshot) -> Self {
146        self.globals.replace(snapshot);
147        self
148    }
149
150    pub fn with_trigger(mut self, trigger: Box<WasmResumeTrigger>) -> Self {
151        self.callbacks.trigger.replace(trigger);
152        self
153    }
154
155    pub fn with_recycle(mut self, recycle: Box<TaskWasmRecycle>) -> Self {
156        self.callbacks.recycle.replace(recycle);
157        self
158    }
159
160    pub fn with_pre_run(mut self, pre_run: Box<TaskWasmPreRun>) -> Self {
161        self.callbacks.pre_run.replace(pre_run);
162        self
163    }
164}
165
166/// A task executor backed by a thread pool.
167///
168/// ## Thread Safety
169///
170/// Due to [#4158], it is possible to pass non-thread safe objects across
171/// threads by capturing them in the task passed to
172/// [`VirtualTaskManager::task_shared()`] or
173/// [`VirtualTaskManager::task_dedicated()`].
174///
175/// If your task needs access to a [`wasmer::Module`], [`wasmer::Memory`], or
176/// [`wasmer::Instance`], it should explicitly transfer the objects using
177/// either [`VirtualTaskManager::task_wasm()`] when in syscall context or
178/// [`VirtualTaskManager::spawn_with_module()`] for higher level code.
179///
180/// [#4158]: https://github.com/wasmerio/wasmer/issues/4158
181#[allow(unused_variables)]
182pub trait VirtualTaskManager: std::fmt::Debug + Send + Sync + 'static {
183    /// Build a new Webassembly memory.
184    ///
185    /// May return `None` if the memory can just be auto-constructed.
186    fn build_memory(
187        &self,
188        mut store: &mut StoreMut,
189        spawn_type: SpawnType,
190    ) -> Result<Option<Memory>, WasiThreadError> {
191        match spawn_type {
192            SpawnType::CreateMemoryOfType(mut ty) => {
193                ty.shared = true;
194
195                // Note: If memory is shared, maximum needs to be set in the
196                // browser otherwise creation will fail.
197                let _ = ty.maximum.get_or_insert(wasmer_types::Pages::max_value());
198
199                let mem = Memory::new(&mut store, ty).map_err(|err| {
200                    tracing::error!(
201                        error = &err as &dyn std::error::Error,
202                        memory_type=?ty,
203                        "could not create memory",
204                    );
205                    WasiThreadError::MemoryCreateFailed(err)
206                })?;
207                Ok(Some(mem))
208            }
209            SpawnType::AttachMemory(mem) => mem
210                .try_attach(store)
211                .map(Some)
212                .map_err(WasiThreadError::MemoryCreateFailed),
213            SpawnType::CreateMemory | SpawnType::NewLinkerInstanceGroup(..) => Ok(None),
214        }
215    }
216
217    /// Pause the current thread of execution.
218    ///
219    /// This is typically invoked whenever a WASM thread goes idle. Besides
220    /// acting as a platform-agnostic [`std::thread::sleep()`], this also gives
221    /// the runtime a chance to do asynchronous work like pumping an event
222    /// loop.
223    fn sleep_now(
224        &self,
225        time: Duration,
226    ) -> Pin<Box<dyn Future<Output = ()> + Send + Sync + 'static>>;
227
228    /// Run an asynchronous operation on the thread pool.
229    ///
230    /// This task must not block execution or it could cause deadlocks.
231    ///
232    /// See the "Thread Safety" documentation on [`VirtualTaskManager`] for
233    /// limitations on what a `task` can and can't contain.
234    fn task_shared(
235        &self,
236        task: Box<dyn FnOnce() -> BoxFuture<'static, ()> + Send + 'static>,
237    ) -> Result<(), WasiThreadError>;
238
239    /// Run a blocking WebAssembly operation on the thread pool.
240    ///
241    /// This is primarily used inside the context of a syscall and allows
242    /// the transfer of things like [`wasmer::Module`] across threads.
243    fn task_wasm(&self, task: TaskWasm) -> Result<(), WasiThreadError>;
244
245    /// Run a blocking operation on the thread pool.
246    ///
247    /// It is okay for this task to block execution and any async futures within
248    /// its scope.
249    fn task_dedicated(
250        &self,
251        task: Box<dyn FnOnce() + Send + 'static>,
252    ) -> Result<(), WasiThreadError>;
253
254    /// Returns the amount of parallelism that is possible on this platform.
255    fn thread_parallelism(&self) -> Result<usize, WasiThreadError>;
256
257    /// Schedule a blocking task to run on the threadpool, explicitly
258    /// transferring a [`Module`] to the task.
259    ///
260    /// This should be preferred over [`VirtualTaskManager::task_dedicated()`]
261    /// where possible because [`wasmer::Module`] is actually `!Send` in the
262    /// browser and can only be transferred to background threads via
263    /// an explicit `postMessage()`. See [#4158] for more details.
264    ///
265    /// This is very similar to [`VirtualTaskManager::task_wasm()`], but
266    /// intended for use outside of a syscall context. For example, when you are
267    /// running in the browser and want to run a WebAssembly module in the
268    /// background.
269    ///
270    /// [#4158]: https://github.com/wasmerio/wasmer/issues/4158
271    fn spawn_with_module(
272        &self,
273        module: Module,
274        task: Box<dyn FnOnce(Module) + Send + 'static>,
275    ) -> Result<(), WasiThreadError> {
276        // Note: Ideally, this function and task_wasm() would be superseded by
277        // a more general mechanism for transferring non-thread safe values
278        // to the thread pool.
279        self.task_dedicated(Box::new(move || task(module)))
280    }
281}
282
283impl<D, T> VirtualTaskManager for D
284where
285    D: Deref<Target = T> + std::fmt::Debug + Send + Sync + 'static,
286    T: VirtualTaskManager + ?Sized,
287{
288    fn build_memory(
289        &self,
290        store: &mut StoreMut,
291        spawn_type: SpawnType,
292    ) -> Result<Option<Memory>, WasiThreadError> {
293        (**self).build_memory(store, spawn_type)
294    }
295
296    fn sleep_now(
297        &self,
298        time: Duration,
299    ) -> Pin<Box<dyn Future<Output = ()> + Send + Sync + 'static>> {
300        (**self).sleep_now(time)
301    }
302
303    fn task_shared(
304        &self,
305        task: Box<dyn FnOnce() -> BoxFuture<'static, ()> + Send + 'static>,
306    ) -> Result<(), WasiThreadError> {
307        (**self).task_shared(task)
308    }
309
310    fn task_wasm(&self, task: TaskWasm) -> Result<(), WasiThreadError> {
311        (**self).task_wasm(task)
312    }
313
314    fn task_dedicated(
315        &self,
316        task: Box<dyn FnOnce() + Send + 'static>,
317    ) -> Result<(), WasiThreadError> {
318        (**self).task_dedicated(task)
319    }
320
321    fn thread_parallelism(&self) -> Result<usize, WasiThreadError> {
322        (**self).thread_parallelism()
323    }
324
325    fn spawn_with_module(
326        &self,
327        module: Module,
328        task: Box<dyn FnOnce(Module) + Send + 'static>,
329    ) -> Result<(), WasiThreadError> {
330        (**self).spawn_with_module(module, task)
331    }
332}
333
334impl dyn VirtualTaskManager {
335    /// Starts an WebAssembly task will run on a dedicated thread
336    /// pulled from the worker pool that has a stateful thread local variable
337    /// After the poller has succeeded
338    #[doc(hidden)]
339    pub unsafe fn resume_wasm_after_poller(
340        &self,
341        task: Box<WasmResumeTask>,
342        ctx: WasiFunctionEnv,
343        mut store: Store,
344        trigger: Pin<Box<AsyncifyFuture>>,
345    ) -> Result<(), WasiThreadError> {
346        // This poller will process any signals when the main working function is idle
347        struct AsyncifyPollerOwned {
348            thread: WasiThread,
349            trigger: Pin<Box<AsyncifyFuture>>,
350        }
351        impl Future for AsyncifyPollerOwned {
352            type Output = Result<Bytes, ExitCode>;
353            fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
354                let work = self.trigger.as_mut();
355                Poll::Ready(if let Poll::Ready(res) = work.poll(cx) {
356                    Ok(res)
357                } else if let Some(forced_exit) = self.thread.try_join() {
358                    return Poll::Ready(Err(forced_exit.unwrap_or_else(|err| {
359                        tracing::debug!("exit runtime error - {}", err);
360                        Errno::Child.into()
361                    })));
362                } else {
363                    return Poll::Pending;
364                })
365            }
366        }
367
368        let snapshot = capture_store_snapshot(&mut store.as_store_mut());
369        let env = ctx.data(&store);
370        let env_inner = env.inner();
371        let handles = env_inner
372            .static_module_instance_handles()
373            .ok_or(WasiThreadError::Unsupported)?;
374        let module = handles.module_clone();
375        let memory = handles.memory_clone();
376        let thread = env.thread.clone();
377        let env = env.clone();
378
379        let thread_inner = thread.clone();
380        self.task_wasm(
381            TaskWasm::new(
382                Box::new(move |props| {
383                    let result = props
384                        .trigger_result
385                        .expect("If there is no result then its likely the trigger did not run");
386                    let result = match result {
387                        Ok(r) => r,
388                        Err(exit_code) => {
389                            thread.set_status_finished(Ok(exit_code));
390                            return;
391                        }
392                    };
393                    task(props.ctx, props.store, result)
394                }),
395                env.clone(),
396                module,
397                false,
398                false,
399            )
400            .with_memory(SpawnType::AttachMemory(
401                memory.as_shared(&store).ok_or_else(|| {
402                    tracing::error!("failed to get shared memory for asyncify poller");
403                    WasiThreadError::Unsupported
404                })?,
405            ))
406            .with_globals(snapshot)
407            .with_trigger(Box::new(move || {
408                Box::pin(async move {
409                    let mut poller = AsyncifyPollerOwned {
410                        thread: thread_inner,
411                        trigger,
412                    };
413                    let res = Pin::new(&mut poller).await;
414                    let res = match res {
415                        Ok(res) => res,
416                        Err(exit_code) => {
417                            env.thread.set_status_finished(Ok(exit_code));
418                            return Err(exit_code);
419                        }
420                    };
421
422                    tracing::trace!("deep sleep woken - res.len={}", res.len());
423                    Ok(res)
424                })
425            })),
426        )
427    }
428}
429
430/// Generic utility methods for VirtualTaskManager
431pub trait VirtualTaskManagerExt {
432    /// Runs the work in the background via the task managers shared background
433    /// threads while blocking the current execution until it finishes
434    fn spawn_and_block_on<A>(
435        &self,
436        task: impl Future<Output = A> + Send + 'static,
437    ) -> Result<A, anyhow::Error>
438    where
439        A: Send + 'static;
440
441    fn spawn_await<O, F>(
442        &self,
443        f: F,
444    ) -> Box<dyn Future<Output = Result<O, Box<dyn std::error::Error>>> + Unpin + Send + 'static>
445    where
446        O: Send + 'static,
447        F: FnOnce() -> O + Send + 'static;
448}
449
450impl<D, T> VirtualTaskManagerExt for D
451where
452    D: Deref<Target = T>,
453    T: VirtualTaskManager + ?Sized,
454{
455    /// Runs the work in the background via the task managers shared background
456    /// threads while blocking the current execution until it finishes
457    fn spawn_and_block_on<A>(
458        &self,
459        task: impl Future<Output = A> + Send + 'static,
460    ) -> Result<A, anyhow::Error>
461    where
462        A: Send + 'static,
463    {
464        let (tx, rx) = ::tokio::sync::oneshot::channel();
465        let work = Box::pin(async move {
466            let ret = task.await;
467            tx.send(ret).ok();
468        });
469        self.task_shared(Box::new(move || work)).unwrap();
470        rx.blocking_recv()
471            .map_err(|_| anyhow::anyhow!("task execution failed - result channel dropped"))
472    }
473
474    fn spawn_await<O, F>(
475        &self,
476        f: F,
477    ) -> Box<dyn Future<Output = Result<O, Box<dyn std::error::Error>>> + Unpin + Send + 'static>
478    where
479        O: Send + 'static,
480        F: FnOnce() -> O + Send + 'static,
481    {
482        let (sender, receiver) = ::tokio::sync::oneshot::channel();
483
484        self.task_dedicated(Box::new(move || {
485            let result = f();
486            let _ = sender.send(result);
487        }))
488        .unwrap();
489
490        Box::new(receiver.map_err(|e| Box::new(e).into()))
491    }
492}
wasmer_wasix/runtime/task_manager/mod.rs

wasmer_wasix/runtime/task_manager/
mod.rs
