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