1 use crate::async_functions::{PollOnce, execute_across_threads};
2 use anyhow::Result;
3 use wasmtime::component::*;
4 use wasmtime::{Engine, Store, StoreContextMut, Trap};
5 use wasmtime_component_util::REALLOC_AND_FREE;
6 
7 /// This is super::func::thunks, except with an async store.
8 #[tokio::test]
9 #[cfg_attr(miri, ignore)]
10 async fn smoke() -> Result<()> {
11     let component = r#"
12         (component
13             (core module $m
14                 (func (export "thunk"))
15                 (func (export "thunk-trap") unreachable)
16             )
17             (core instance $i (instantiate $m))
18             (func (export "thunk")
19                 (canon lift (core func $i "thunk"))
20             )
21             (func (export "thunk-trap")
22                 (canon lift (core func $i "thunk-trap"))
23             )
24         )
25     "#;
26 
27     let engine = super::async_engine();
28     let component = Component::new(&engine, component)?;
29     let mut store = Store::new(&engine, ());
30     let instance = Linker::new(&engine)
31         .instantiate_async(&mut store, &component)
32         .await?;
33 
34     let thunk = instance.get_typed_func::<(), ()>(&mut store, "thunk")?;
35 
36     thunk.call_async(&mut store, ()).await?;
37     thunk.post_return_async(&mut store).await?;
38 
39     let err = instance
40         .get_typed_func::<(), ()>(&mut store, "thunk-trap")?
41         .call_async(&mut store, ())
42         .await
43         .unwrap_err();
44     assert_eq!(err.downcast::<Trap>()?, Trap::UnreachableCodeReached);
45 
46     Ok(())
47 }
48 
49 /// Handle an import function, created using component::Linker::func_wrap_async.
50 #[tokio::test]
51 #[cfg_attr(miri, ignore)]
52 async fn smoke_func_wrap() -> Result<()> {
53     let component = r#"
54         (component
55             (type $f (func))
56             (import "i" (func $f))
57 
58             (core module $m
59                 (import "imports" "i" (func $i))
60                 (func (export "thunk") call $i)
61             )
62 
63             (core func $f (canon lower (func $f)))
64             (core instance $i (instantiate $m
65                 (with "imports" (instance
66                     (export "i" (func $f))
67                 ))
68              ))
69             (func (export "thunk")
70                 (canon lift (core func $i "thunk"))
71             )
72         )
73     "#;
74 
75     let engine = super::async_engine();
76     let component = Component::new(&engine, component)?;
77     let mut store = Store::new(&engine, ());
78     let mut linker = Linker::new(&engine);
79     let mut root = linker.root();
80     root.func_wrap_async("i", |_: StoreContextMut<()>, _: ()| {
81         Box::new(async { Ok(()) })
82     })?;
83 
84     let instance = linker.instantiate_async(&mut store, &component).await?;
85 
86     let thunk = instance.get_typed_func::<(), ()>(&mut store, "thunk")?;
87 
88     thunk.call_async(&mut store, ()).await?;
89     thunk.post_return_async(&mut store).await?;
90 
91     Ok(())
92 }
93 
94 // This test stresses TLS management in combination with the `realloc` option
95 // for imported functions. This will create an async computation which invokes a
96 // component that invokes an imported function. The imported function returns a
97 // list which will require invoking malloc.
98 //
99 // As an added stressor all polls are sprinkled across threads through
100 // `execute_across_threads`. Yields are injected liberally by configuring 1
101 // fuel consumption to trigger a yield.
102 //
103 // Overall a yield should happen during malloc which should be an "interesting
104 // situation" with respect to the runtime.
105 #[tokio::test]
106 #[cfg_attr(miri, ignore)]
107 async fn resume_separate_thread() -> Result<()> {
108     let mut config = wasmtime_test_util::component::config();
109     config.async_support(true);
110     config.consume_fuel(true);
111     let engine = Engine::new(&config)?;
112     let component = format!(
113         r#"
114             (component
115                 (import "yield" (func $yield (result (list u8))))
116                 (core module $libc
117                     (memory (export "memory") 1)
118                     {REALLOC_AND_FREE}
119                 )
120                 (core instance $libc (instantiate $libc))
121 
122                 (core func $yield
123                     (canon lower
124                         (func $yield)
125                         (memory $libc "memory")
126                         (realloc (func $libc "realloc"))
127                     )
128                 )
129 
130                 (core module $m
131                     (import "" "yield" (func $yield (param i32)))
132                     (import "libc" "memory" (memory 0))
133                     (func $start
134                         i32.const 8
135                         call $yield
136                     )
137                     (start $start)
138                 )
139                 (core instance (instantiate $m
140                     (with "" (instance (export "yield" (func $yield))))
141                     (with "libc" (instance $libc))
142                 ))
143             )
144         "#
145     );
146     let component = Component::new(&engine, component)?;
147     let mut linker = Linker::new(&engine);
148     linker
149         .root()
150         .func_wrap_async("yield", |_: StoreContextMut<()>, _: ()| {
151             Box::new(async {
152                 tokio::task::yield_now().await;
153                 Ok((vec![1u8, 2u8],))
154             })
155         })?;
156 
157     execute_across_threads(async move {
158         let mut store = Store::new(&engine, ());
159         store.set_fuel(u64::MAX).unwrap();
160         store.fuel_async_yield_interval(Some(1)).unwrap();
161         linker.instantiate_async(&mut store, &component).await?;
162         Ok::<_, anyhow::Error>(())
163     })
164     .await?;
165     Ok(())
166 }
167 
168 // This test is intended to stress TLS management in the component model around
169 // the management of the `realloc` function. This creates an async computation
170 // representing the execution of a component model function where entry into the
171 // component uses `realloc` and then the component runs. This async computation
172 // is then polled iteratively with another "wasm activation" (in this case a
173 // core wasm function) on the stack. The poll-per-call should work and nothing
174 // should in theory have problems here.
175 //
176 // As an added stressor all polls are sprinkled across threads through
177 // `execute_across_threads`. Yields are injected liberally by configuring 1
178 // fuel consumption to trigger a yield.
179 //
180 // Overall a yield should happen during malloc which should be an "interesting
181 // situation" with respect to the runtime.
182 #[tokio::test]
183 #[cfg_attr(miri, ignore)]
184 async fn poll_through_wasm_activation() -> Result<()> {
185     let mut config = wasmtime_test_util::component::config();
186     config.async_support(true);
187     config.consume_fuel(true);
188     let engine = Engine::new(&config)?;
189     let component = format!(
190         r#"
191             (component
192                 (core module $m
193                     {REALLOC_AND_FREE}
194                     (memory (export "memory") 1)
195                     (func (export "run") (param i32 i32)
196                     )
197                 )
198                 (core instance $i (instantiate $m))
199                 (func (export "run") (param "x" (list u8))
200                     (canon lift (core func $i "run")
201                                 (memory $i "memory")
202                                 (realloc (func $i "realloc"))))
203             )
204         "#
205     );
206     let component = Component::new(&engine, component)?;
207     let linker = Linker::new(&engine);
208 
209     let invoke_component = {
210         let engine = engine.clone();
211         async move {
212             let mut store = Store::new(&engine, ());
213             store.set_fuel(u64::MAX).unwrap();
214             store.fuel_async_yield_interval(Some(1)).unwrap();
215             let instance = linker.instantiate_async(&mut store, &component).await?;
216             let func = instance.get_typed_func::<(Vec<u8>,), ()>(&mut store, "run")?;
217             func.call_async(&mut store, (vec![1, 2, 3],)).await?;
218             Ok::<_, anyhow::Error>(())
219         }
220     };
221 
222     execute_across_threads(async move {
223         let mut store = Store::new(&engine, Some(Box::pin(invoke_component)));
224         let poll_once = wasmtime::Func::wrap_async(&mut store, |mut cx, _: ()| {
225             let invoke_component = cx.data_mut().take().unwrap();
226             Box::new(async move {
227                 match PollOnce::new(invoke_component).await {
228                     Ok(result) => {
229                         result?;
230                         Ok(1)
231                     }
232                     Err(future) => {
233                         *cx.data_mut() = Some(future);
234                         Ok(0)
235                     }
236                 }
237             })
238         });
239         let poll_once = poll_once.typed::<(), i32>(&mut store)?;
240         while poll_once.call_async(&mut store, ()).await? != 1 {
241             // loop around to call again
242         }
243         Ok::<_, anyhow::Error>(())
244     })
245     .await?;
246     Ok(())
247 }
248 
249 /// Test async drop method for host resources.
250 #[tokio::test]
251 #[cfg_attr(miri, ignore)]
252 async fn drop_resource_async() -> Result<()> {
253     use std::sync::Arc;
254     use std::sync::Mutex;
255 
256     let engine = super::async_engine();
257     let c = Component::new(
258         &engine,
259         r#"
260             (component
261                 (import "t" (type $t (sub resource)))
262 
263                 (core func $drop (canon resource.drop $t))
264 
265                 (core module $m
266                     (import "" "drop" (func $drop (param i32)))
267                     (func (export "f") (param i32)
268                         (call $drop (local.get 0))
269                     )
270                 )
271                 (core instance $i (instantiate $m
272                     (with "" (instance
273                         (export "drop" (func $drop))
274                     ))
275                 ))
276 
277                 (func (export "f") (param "x" (own $t))
278                     (canon lift (core func $i "f")))
279             )
280         "#,
281     )?;
282 
283     struct MyType;
284 
285     let mut store = Store::new(&engine, ());
286     let mut linker = Linker::new(&engine);
287 
288     let drop_status = Arc::new(Mutex::new("not dropped"));
289     let ds = drop_status.clone();
290 
291     linker
292         .root()
293         .resource_async("t", ResourceType::host::<MyType>(), move |_, _| {
294             let ds = ds.clone();
295             Box::new(async move {
296                 *ds.lock().unwrap() = "before yield";
297                 tokio::task::yield_now().await;
298                 *ds.lock().unwrap() = "after yield";
299                 Ok(())
300             })
301         })?;
302     let i = linker.instantiate_async(&mut store, &c).await?;
303     let f = i.get_typed_func::<(Resource<MyType>,), ()>(&mut store, "f")?;
304 
305     execute_across_threads(async move {
306         let resource = Resource::new_own(100);
307         f.call_async(&mut store, (resource,)).await?;
308         f.post_return_async(&mut store).await?;
309         Ok::<_, anyhow::Error>(())
310     })
311     .await?;
312 
313     assert_eq!("after yield", *drop_status.lock().unwrap());
314 
315     Ok(())
316 }
317