1 #![cfg(not(miri))]
2 
3 use crate::async_functions::{CountPending, PollOnce};
4 use anyhow::anyhow;
5 use std::sync::Arc;
6 use std::sync::atomic::{AtomicBool, Ordering};
7 use wasmtime::*;
8 use wasmtime_test_macros::wasmtime_test;
9 
10 fn build_engine(config: &mut Config) -> Result<Arc<Engine>> {
11     config.async_support(true);
12     config.epoch_interruption(true);
13     Ok(Arc::new(Engine::new(&config)?))
14 }
15 
16 fn make_env<T: 'static>(engine: &Engine) -> Linker<T> {
17     let mut linker = Linker::new(engine);
18     let engine = engine.clone();
19 
20     linker
21         .func_new(
22             "",
23             "bump_epoch",
24             FuncType::new(&engine, None, None),
25             move |_caller, _params, _results| {
26                 engine.increment_epoch();
27                 Ok(())
28             },
29         )
30         .unwrap();
31 
32     linker
33 }
34 
35 enum InterruptMode {
36     Trap,
37     Callback(fn(StoreContextMut<usize>) -> Result<UpdateDeadline>),
38     Yield(u64),
39 }
40 
41 /// Run a test with the given wasm, giving an initial deadline of
42 /// `initial` ticks in the future, and either configuring the wasm to
43 /// yield and set a deadline `delta` ticks in the future if `delta` is
44 /// `Some(..)` or trapping if `delta` is `None`.
45 ///
46 /// Returns `Some((yields, store))` if function completed normally, giving
47 /// the number of yields that occurred, or `None` if a trap occurred.
48 async fn run_and_count_yields_or_trap<F: Fn(Arc<Engine>)>(
49     config: &mut Config,
50     wasm: &str,
51     initial: u64,
52     delta: InterruptMode,
53     setup_func: F,
54 ) -> Result<Option<(usize, usize)>> {
55     let engine = build_engine(config)?;
56     let linker = make_env::<usize>(&engine);
57     let module = Module::new(&engine, wasm)?;
58     let mut store = Store::new(&engine, 0);
59     store.set_epoch_deadline(initial);
60     match delta {
61         InterruptMode::Yield(delta) => {
62             store.epoch_deadline_async_yield_and_update(delta);
63         }
64         InterruptMode::Callback(func) => {
65             store.epoch_deadline_callback(func);
66         }
67         InterruptMode::Trap => {
68             store.epoch_deadline_trap();
69         }
70     }
71 
72     let engine_clone = engine.clone();
73     setup_func(engine_clone);
74 
75     let instance = linker.instantiate_async(&mut store, &module).await?;
76     let f = instance.get_func(&mut store, "run").unwrap();
77     let (result, yields) =
78         CountPending::new(Box::pin(f.call_async(&mut store, &[], &mut []))).await;
79     let store = store.data();
80     Ok(result.ok().map(|_| (yields, *store)))
81 }
82 
83 #[wasmtime_test(with = "#[tokio::test]")]
84 async fn epoch_yield_at_func_entry(config: &mut Config) -> Result<()> {
85     // Should yield at start of call to func $subfunc.
86     assert_eq!(
87         Some((1, 0)),
88         run_and_count_yields_or_trap(
89             config,
90             "
91             (module
92                 (import \"\" \"bump_epoch\" (func $bump))
93                 (func (export \"run\")
94                     call $bump  ;; bump epoch
95                     call $subfunc) ;; call func; will notice new epoch and yield
96                 (func $subfunc))
97             ",
98             1,
99             InterruptMode::Yield(1),
100             |_| {},
101         )
102         .await?
103     );
104     Ok(())
105 }
106 
107 #[wasmtime_test(with = "#[tokio::test]")]
108 async fn epoch_yield_at_loop_header(config: &mut Config) -> Result<()> {
109     // Should yield at top of loop, once per five iters.
110     assert_eq!(
111         Some((2, 0)),
112         run_and_count_yields_or_trap(
113             config,
114             "
115             (module
116                 (import \"\" \"bump_epoch\" (func $bump))
117                 (func (export \"run\")
118                     (local $i i32)
119                     (local.set $i (i32.const 10))
120                     (loop $l
121                         call $bump
122                         (br_if $l (local.tee $i (i32.sub (local.get $i) (i32.const 1)))))))
123             ",
124             0,
125             InterruptMode::Yield(5),
126             |_| {},
127         )
128         .await?
129     );
130     Ok(())
131 }
132 
133 #[wasmtime_test(with = "#[tokio::test]")]
134 async fn epoch_yield_immediate(config: &mut Config) -> Result<()> {
135     // We should see one yield immediately when the initial deadline
136     // is zero.
137     assert_eq!(
138         Some((1, 0)),
139         run_and_count_yields_or_trap(
140             config,
141             "
142             (module
143                 (import \"\" \"bump_epoch\" (func $bump))
144                 (func (export \"run\")))
145             ",
146             0,
147             InterruptMode::Yield(1),
148             |_| {},
149         )
150         .await?
151     );
152     Ok(())
153 }
154 
155 #[wasmtime_test(with = "#[tokio::test]")]
156 async fn epoch_yield_only_once(config: &mut Config) -> Result<()> {
157     // We should yield from the subfunction, and then when we return
158     // to the outer function and hit another loop header, we should
159     // not yield again (the double-check block will reload the correct
160     // epoch).
161     assert_eq!(
162         Some((1, 0)),
163         run_and_count_yields_or_trap(
164             config,
165             "
166             (module
167                 (import \"\" \"bump_epoch\" (func $bump))
168                 (func (export \"run\")
169                   (local $i i32)
170                   (call $subfunc)
171                   (local.set $i (i32.const 0))
172                   (loop $l
173                     (br_if $l (i32.eq (i32.const 10)
174                                       (local.tee $i (i32.add (i32.const 1) (local.get $i)))))))
175                 (func $subfunc
176                   (call $bump)))
177             ",
178             1,
179             InterruptMode::Yield(1),
180             |_| {},
181         )
182         .await?
183     );
184     Ok(())
185 }
186 
187 #[wasmtime_test(with = "#[tokio::test]")]
188 async fn epoch_interrupt_infinite_loop(config: &mut Config) -> Result<()> {
189     assert_eq!(
190         None,
191         run_and_count_yields_or_trap(
192             config,
193             "
194             (module
195                 (import \"\" \"bump_epoch\" (func $bump))
196                 (func (export \"run\")
197                   (loop $l
198                     (br $l))))
199             ",
200             1,
201             InterruptMode::Trap,
202             |engine| {
203                 std::thread::spawn(move || {
204                     std::thread::sleep(std::time::Duration::from_millis(50));
205                     engine.increment_epoch();
206                 });
207             },
208         )
209         .await?
210     );
211     Ok(())
212 }
213 
214 #[wasmtime_test(with = "#[tokio::test]")]
215 async fn epoch_interrupt_function_entries(config: &mut Config) -> Result<()> {
216     assert_eq!(
217         None,
218         run_and_count_yields_or_trap(
219             config,
220             "
221             (module
222                 (import \"\" \"bump_epoch\" (func $bump))
223                 (func (export \"run\")
224                   call $f1
225                   call $f1
226                   call $f1
227                   call $f1
228                   call $f1
229                   call $f1
230                   call $f1
231                   call $f1
232                   call $f1
233                   call $f1)
234                 (func $f1
235                   call $f2
236                   call $f2
237                   call $f2
238                   call $f2
239                   call $f2
240                   call $f2
241                   call $f2
242                   call $f2
243                   call $f2
244                   call $f2)
245                 (func $f2
246                   call $f3
247                   call $f3
248                   call $f3
249                   call $f3
250                   call $f3
251                   call $f3
252                   call $f3
253                   call $f3
254                   call $f3
255                   call $f3)
256                 (func $f3
257                   call $f4
258                   call $f4
259                   call $f4
260                   call $f4
261                   call $f4
262                   call $f4
263                   call $f4
264                   call $f4
265                   call $f4
266                   call $f4)
267                 (func $f4
268                   call $f5
269                   call $f5
270                   call $f5
271                   call $f5
272                   call $f5
273                   call $f5
274                   call $f5
275                   call $f5
276                   call $f5
277                   call $f5)
278                 (func $f5
279                   call $f6
280                   call $f6
281                   call $f6
282                   call $f6
283                   call $f6
284                   call $f6
285                   call $f6
286                   call $f6
287                   call $f6
288                   call $f6)
289                 (func $f6
290                   call $f7
291                   call $f7
292                   call $f7
293                   call $f7
294                   call $f7
295                   call $f7
296                   call $f7
297                   call $f7
298                   call $f7
299                   call $f7)
300                 (func $f7
301                   call $f8
302                   call $f8
303                   call $f8
304                   call $f8
305                   call $f8
306                   call $f8
307                   call $f8
308                   call $f8
309                   call $f8
310                   call $f8)
311                 (func $f8
312                   call $f9
313                   call $f9
314                   call $f9
315                   call $f9
316                   call $f9
317                   call $f9
318                   call $f9
319                   call $f9
320                   call $f9
321                   call $f9)
322                 (func $f9))
323             ",
324             1,
325             InterruptMode::Trap,
326             |engine| {
327                 std::thread::spawn(move || {
328                     std::thread::sleep(std::time::Duration::from_millis(50));
329                     engine.increment_epoch();
330                 });
331             },
332         )
333         .await?
334     );
335     Ok(())
336 }
337 
338 #[wasmtime_test(with = "#[tokio::test]")]
339 async fn epoch_callback_continue(config: &mut Config) -> Result<()> {
340     assert_eq!(
341         Some((0, 1)),
342         run_and_count_yields_or_trap(
343             config,
344             "
345             (module
346                 (import \"\" \"bump_epoch\" (func $bump))
347                 (func (export \"run\")
348                     call $bump  ;; bump epoch
349                     call $subfunc) ;; call func; will notice new epoch and yield
350                 (func $subfunc))
351             ",
352             1,
353             InterruptMode::Callback(|mut cx| {
354                 let s = cx.data_mut();
355                 *s += 1;
356                 Ok(UpdateDeadline::Continue(1))
357             }),
358             |_| {},
359         )
360         .await?
361     );
362     Ok(())
363 }
364 
365 #[wasmtime_test(with = "#[tokio::test]")]
366 async fn epoch_callback_yield(config: &mut Config) -> Result<()> {
367     assert_eq!(
368         Some((1, 1)),
369         run_and_count_yields_or_trap(
370             config,
371             "
372             (module
373                 (import \"\" \"bump_epoch\" (func $bump))
374                 (func (export \"run\")
375                     call $bump  ;; bump epoch
376                     call $subfunc) ;; call func; will notice new epoch and yield
377                 (func $subfunc))
378             ",
379             1,
380             InterruptMode::Callback(|mut cx| {
381                 let s = cx.data_mut();
382                 *s += 1;
383                 Ok(UpdateDeadline::Yield(1))
384             }),
385             |_| {},
386         )
387         .await?
388     );
389 
390     Ok(())
391 }
392 
393 #[wasmtime_test(with = "#[tokio::test]")]
394 async fn epoch_callback_yield_custom(config: &mut Config) -> Result<()> {
395     assert_eq!(
396         Some((1, 1)),
397         run_and_count_yields_or_trap(
398             config,
399             "
400             (module
401                 (import \"\" \"bump_epoch\" (func $bump))
402                 (func (export \"run\")
403                     call $bump  ;; bump epoch
404                     call $subfunc) ;; call func; will notice new epoch and yield
405                 (func $subfunc))
406             ",
407             1,
408             InterruptMode::Callback(|mut cx| {
409                 let s = cx.data_mut();
410                 *s += 1;
411                 let fut = Box::pin(tokio::task::yield_now());
412                 Ok(UpdateDeadline::YieldCustom(1, fut))
413             }),
414             |_| {},
415         )
416         .await?
417     );
418     Ok(())
419 }
420 
421 #[wasmtime_test(with = "#[tokio::test]")]
422 async fn epoch_callback_trap(config: &mut Config) -> Result<()> {
423     assert_eq!(
424         None,
425         run_and_count_yields_or_trap(
426             config,
427             "
428             (module
429                 (import \"\" \"bump_epoch\" (func $bump))
430                 (func (export \"run\")
431                     call $bump  ;; bump epoch
432                     call $subfunc) ;; call func; will notice new epoch and yield
433                 (func $subfunc))
434             ",
435             1,
436             InterruptMode::Callback(|_| Err(anyhow!("Failing in callback"))),
437             |_| {},
438         )
439         .await?
440     );
441     Ok(())
442 }
443 
444 #[wasmtime_test(with = "#[tokio::test]")]
445 async fn drop_future_on_epoch_yield(config: &mut Config) -> Result<()> {
446     let wasm = "
447     (module
448       (import \"\" \"bump_epoch\" (func $bump))
449       (import \"\" \"im_alive\" (func $im_alive))
450       (import \"\" \"oops\" (func $oops))
451       (func (export \"run\")
452         (call $im_alive)
453         (call $bump)
454         (call $subfunc)  ;; subfunc entry to do epoch check
455         (call $oops))
456       (func $subfunc))
457     ";
458 
459     let engine = build_engine(config)?;
460     let mut linker = make_env::<()>(&engine);
461 
462     // Create a few helpers for the Wasm to call.
463     let alive_flag = Arc::new(AtomicBool::new(false));
464     let alive_flag_clone = alive_flag.clone();
465     linker
466         .func_new(
467             "",
468             "oops",
469             FuncType::new(&engine, None, None),
470             move |_caller, _params, _results| {
471                 panic!("Should not have reached this point!");
472             },
473         )
474         .unwrap();
475     linker
476         .func_new(
477             "",
478             "im_alive",
479             FuncType::new(&engine, None, None),
480             move |_caller, _params, _results| {
481                 alive_flag_clone.store(true, Ordering::Release);
482                 Ok(())
483             },
484         )
485         .unwrap();
486 
487     let module = Module::new(&engine, wasm).unwrap();
488     let mut store = Store::new(&engine, ());
489 
490     store.set_epoch_deadline(1);
491     store.epoch_deadline_async_yield_and_update(1);
492 
493     let instance = linker.instantiate_async(&mut store, &module).await.unwrap();
494     let f = instance.get_func(&mut store, "run").unwrap();
495     let _ = PollOnce::new(Box::pin(f.call_async(&mut store, &[], &mut []))).await;
496 
497     assert_eq!(true, alive_flag.load(Ordering::Acquire));
498     Ok(())
499 }
500