wasmer_wasix/runtime/task_manager/
tokio.rs

1use std::sync::Mutex;
2use std::{num::NonZeroUsize, pin::Pin, sync::Arc, time::Duration};
3
4use futures::{Future, future::BoxFuture};
5use tokio::runtime::{Handle, Runtime};
6use virtual_mio::block_on;
7
8use crate::runtime::{SpawnType, task_manager::TaskWasmCallbacks};
9use crate::{WasiFunctionEnv, os::task::thread::WasiThreadError};
10
11use super::{SpawnMemoryTypeOrStore, TaskWasm, TaskWasmRunProperties, VirtualTaskManager};
12
13#[derive(Debug, Clone)]
14pub enum RuntimeOrHandle {
15    Handle(Handle),
16    Runtime(Handle, Arc<Mutex<Option<Runtime>>>),
17}
18impl From<Handle> for RuntimeOrHandle {
19    fn from(value: Handle) -> Self {
20        Self::Handle(value)
21    }
22}
23impl From<Runtime> for RuntimeOrHandle {
24    fn from(value: Runtime) -> Self {
25        Self::Runtime(value.handle().clone(), Arc::new(Mutex::new(Some(value))))
26    }
27}
28
29impl Drop for RuntimeOrHandle {
30    fn drop(&mut self) {
31        if let Self::Runtime(_, runtime) = self
32            && let Some(h) = runtime.lock().unwrap().take()
33        {
34            h.shutdown_timeout(Duration::from_secs(0))
35        }
36    }
37}
38
39impl RuntimeOrHandle {
40    pub fn handle(&self) -> &Handle {
41        match self {
42            Self::Handle(h) => h,
43            Self::Runtime(h, _) => h,
44        }
45    }
46}
47
48#[derive(Clone)]
49pub struct ThreadPool {
50    inner: rusty_pool::ThreadPool,
51}
52
53impl std::ops::Deref for ThreadPool {
54    type Target = rusty_pool::ThreadPool;
55
56    fn deref(&self) -> &Self::Target {
57        &self.inner
58    }
59}
60
61impl std::fmt::Debug for ThreadPool {
62    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
63        f.debug_struct("ThreadPool")
64            .field("name", &self.get_name())
65            .field("current_worker_count", &self.get_current_worker_count())
66            .field("idle_worker_count", &self.get_idle_worker_count())
67            .finish()
68    }
69}
70
71/// A task manager that uses tokio to spawn tasks.
72#[derive(Clone, Debug)]
73pub struct TokioTaskManager {
74    rt: RuntimeOrHandle,
75    pool: Arc<ThreadPool>,
76}
77
78impl TokioTaskManager {
79    pub fn new<I>(rt: I) -> Self
80    where
81        I: Into<RuntimeOrHandle>,
82    {
83        let concurrency = std::thread::available_parallelism()
84            .unwrap_or(NonZeroUsize::new(1).unwrap())
85            .get();
86        let max_threads = 200usize.max(concurrency * 100);
87
88        Self {
89            rt: rt.into(),
90            pool: Arc::new(ThreadPool {
91                inner: rusty_pool::Builder::new()
92                    .name("TokioTaskManager Thread Pool".to_string())
93                    .core_size(max_threads)
94                    .max_size(max_threads)
95                    .build(),
96            }),
97        }
98    }
99
100    pub fn runtime_handle(&self) -> tokio::runtime::Handle {
101        self.rt.handle().clone()
102    }
103
104    pub fn pool_handle(&self) -> Arc<ThreadPool> {
105        self.pool.clone()
106    }
107}
108
109impl Default for TokioTaskManager {
110    fn default() -> Self {
111        Self::new(Handle::current())
112    }
113}
114
115impl VirtualTaskManager for TokioTaskManager {
116    /// See [`VirtualTaskManager::sleep_now`].
117    fn sleep_now(&self, time: Duration) -> Pin<Box<dyn Future<Output = ()> + Send + Sync>> {
118        let handle = self.runtime_handle();
119        Box::pin(async move {
120            SleepNow::default()
121                .enter(handle, time)
122                .await
123                .ok()
124                .unwrap_or(())
125        })
126    }
127
128    /// See [`VirtualTaskManager::task_shared`].
129    fn task_shared(
130        &self,
131        task: Box<dyn FnOnce() -> BoxFuture<'static, ()> + Send + 'static>,
132    ) -> Result<(), WasiThreadError> {
133        self.rt.handle().spawn(async move {
134            let fut = task();
135            fut.await
136        });
137        Ok(())
138    }
139
140    /// See [`VirtualTaskManager::task_wasm`].
141    fn task_wasm(&self, task: TaskWasm) -> Result<(), WasiThreadError> {
142        fn env_and_store(
143            task: TaskWasm,
144        ) -> Result<(WasiFunctionEnv, wasmer::Store, TaskWasmCallbacks), WasiThreadError> {
145            let (make_memory, instance_group_data) = match task.spawn_type {
146                SpawnType::CreateMemory => (SpawnMemoryTypeOrStore::New, None),
147                SpawnType::NewLinkerInstanceGroup(instance_group_data) => {
148                    (SpawnMemoryTypeOrStore::New, Some(instance_group_data))
149                }
150                SpawnType::CreateMemoryOfType(t) => (SpawnMemoryTypeOrStore::Type(t), None),
151                SpawnType::AttachMemory(mem) => {
152                    let mut store = task.env.runtime().new_store();
153                    let memory = mem
154                        .try_attach(&mut store)
155                        .map_err(WasiThreadError::MemoryCreateFailed)?;
156                    (SpawnMemoryTypeOrStore::StoreAndMemory(store, memory), None)
157                }
158            };
159
160            let (env, store) = WasiFunctionEnv::new_with_store(
161                task.module,
162                task.env,
163                task.globals,
164                make_memory,
165                task.update_layout,
166                task.call_initialize,
167                instance_group_data,
168            )?;
169            Ok((env, store, task.callbacks))
170        }
171
172        let (sx, rx) = std::sync::mpsc::channel();
173
174        if task.callbacks.trigger.is_some() {
175            tracing::trace!("spawning task_wasm trigger in async pool");
176            self.pool.execute(move || {
177                let (mut ctx, mut store, callbacks) = match env_and_store(task) {
178                    Ok(x) => {
179                        sx.send(Ok(())).unwrap();
180                        x
181                    }
182                    Err(c) => {
183                        tracing::error!("failed to prepare environment for task_wasm trigger: {c}");
184                        sx.send(Err(c)).unwrap();
185                        return;
186                    }
187                };
188
189                let result = {
190                    let mut trigger = (callbacks.trigger.unwrap())();
191                    let pre_run = callbacks.pre_run;
192                    let ctx = &mut ctx;
193                    let store = &mut store;
194                    block_on(async move {
195                        // We wait for either the trigger or for a snapshot to take place
196                        let result = loop {
197                            let env = ctx.data(store);
198                            break tokio::select! {
199                                r = &mut trigger => r,
200                                _ = env.thread.wait_for_signal() => {
201                                    tracing::debug!("wait-for-signal(triggered)");
202                                    let mut ctx = ctx.env.clone().into_mut(store);
203                                    if let Err(err) =
204                                        crate::WasiEnv::do_pending_link_operations(
205                                            &mut ctx,
206                                            false
207                                        ).and_then(|()|
208                                            crate::WasiEnv::process_signals_and_exit(&mut ctx)
209                                        )
210                                    {
211                                        match err {
212                                            crate::WasiError::Exit(code) => Err(code),
213                                            err => {
214                                                tracing::error!("failed to process signals - {}", err);
215                                                continue;
216                                            }
217                                        }
218                                    } else {
219                                        continue;
220                                    }
221                                }
222                                _ = crate::wait_for_snapshot(env) => {
223                                    tracing::debug!("wait-for-snapshot(triggered)");
224                                    let mut ctx = ctx.env.clone().into_mut(store);
225                                    crate::os::task::WasiProcessInner::do_checkpoints_from_outside(&mut ctx);
226                                    continue;
227                                }
228                            };
229                        };
230
231                        if let Some(pre_run) = pre_run {
232                            pre_run(ctx, store).await;
233                        }
234
235                        result
236                    })
237                };
238
239                // Invoke the callback
240                (callbacks.run)(TaskWasmRunProperties {
241                    ctx,
242                    store,
243                    trigger_result: Some(result),
244                    recycle: callbacks.recycle,
245                });
246            });
247        } else {
248            tracing::trace!("spawning task_wasm in blocking thread");
249
250            // Run the callback on a dedicated thread
251            self.pool.execute(move || {
252                tracing::trace!("task_wasm started in blocking thread");
253                let (mut ctx, mut store, callbacks) = match env_and_store(task) {
254                    Ok(x) => {
255                        sx.send(Ok(())).unwrap();
256                        x
257                    }
258                    Err(c) => {
259                        sx.send(Err(c)).unwrap();
260                        return;
261                    }
262                };
263
264                if let Some(pre_run) = callbacks.pre_run {
265                    block_on(pre_run(&mut ctx, &mut store));
266                }
267
268                // Invoke the callback
269                (callbacks.run)(TaskWasmRunProperties {
270                    ctx,
271                    store,
272                    trigger_result: None,
273                    recycle: callbacks.recycle,
274                });
275            });
276        }
277
278        rx.recv()
279            .map_err(|_| WasiThreadError::InvalidWasmContext)??;
280
281        Ok(())
282    }
283
284    /// See [`VirtualTaskManager::task_dedicated`].
285    fn task_dedicated(
286        &self,
287        task: Box<dyn FnOnce() + Send + 'static>,
288    ) -> Result<(), WasiThreadError> {
289        self.pool.execute(move || {
290            task();
291        });
292        Ok(())
293    }
294
295    /// See [`VirtualTaskManager::thread_parallelism`].
296    fn thread_parallelism(&self) -> Result<usize, WasiThreadError> {
297        Ok(std::thread::available_parallelism()
298            .map(usize::from)
299            .unwrap_or(8))
300    }
301}
302
303// Used by [`VirtualTaskManager::sleep_now`] to abort a sleep task when drop.
304#[derive(Default)]
305pub struct SleepNow {
306    abort_handle: Option<tokio::task::AbortHandle>,
307}
308
309impl SleepNow {
310    pub async fn enter(
311        &mut self,
312        handle: tokio::runtime::Handle,
313        time: Duration,
314    ) -> Result<(), tokio::task::JoinError> {
315        let handle = handle.spawn(async move {
316            if time == Duration::ZERO {
317                tokio::task::yield_now().await;
318            } else {
319                tokio::time::sleep(time).await;
320            }
321        });
322        self.abort_handle = Some(handle.abort_handle());
323        handle.await
324    }
325}
326
327impl Drop for SleepNow {
328    fn drop(&mut self) {
329        if let Some(h) = self.abort_handle.as_ref() {
330            h.abort()
331        }
332    }
333}
wasmer_wasix/runtime/task_manager/tokio.rs

wasmer_wasix/runtime/task_manager/
tokio.rs
