1 #![cfg(not(miri))] 2 3 use anyhow::{anyhow, bail}; 4 use std::future::Future; 5 use std::pin::Pin; 6 use std::sync::{Arc, Mutex}; 7 use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker}; 8 use wasmtime::*; 9 10 fn async_store() -> Store<()> { 11 Store::new(&Engine::new(Config::new().async_support(true)).unwrap(), ()) 12 } 13 14 async fn run_smoke_test(store: &mut Store<()>, func: Func) { 15 func.call_async(&mut *store, &[], &mut []).await.unwrap(); 16 func.call_async(&mut *store, &[], &mut []).await.unwrap(); 17 } 18 19 async fn run_smoke_typed_test(store: &mut Store<()>, func: Func) { 20 let func = func.typed::<(), ()>(&store).unwrap(); 21 func.call_async(&mut *store, ()).await.unwrap(); 22 func.call_async(&mut *store, ()).await.unwrap(); 23 } 24 25 #[tokio::test] 26 async fn smoke() { 27 let mut store = async_store(); 28 let func_ty = FuncType::new(store.engine(), None, None); 29 let func = Func::new_async(&mut store, func_ty, move |_caller, _params, _results| { 30 Box::new(async { Ok(()) }) 31 }); 32 run_smoke_test(&mut store, func).await; 33 run_smoke_typed_test(&mut store, func).await; 34 35 let func = Func::wrap_async(&mut store, move |_caller, _: ()| Box::new(async { Ok(()) })); 36 run_smoke_test(&mut store, func).await; 37 run_smoke_typed_test(&mut store, func).await; 38 } 39 40 #[tokio::test] 41 async fn smoke_host_func() -> Result<()> { 42 let mut store = async_store(); 43 let mut linker = Linker::new(store.engine()); 44 45 linker.func_new_async( 46 "", 47 "first", 48 FuncType::new(store.engine(), None, None), 49 move |_caller, _params, _results| Box::new(async { Ok(()) }), 50 )?; 51 52 linker.func_wrap_async("", "second", move |_caller, _: ()| { 53 Box::new(async { Ok(()) }) 54 })?; 55 56 let func = linker 57 .get(&mut store, "", "first") 58 .unwrap() 59 .into_func() 60 .unwrap(); 61 run_smoke_test(&mut store, func).await; 62 run_smoke_typed_test(&mut store, func).await; 63 64 let func = linker 65 .get(&mut store, "", "second") 66 .unwrap() 67 .into_func() 68 .unwrap(); 69 run_smoke_test(&mut store, func).await; 70 run_smoke_typed_test(&mut store, func).await; 71 72 Ok(()) 73 } 74 75 #[tokio::test] 76 async fn smoke_with_suspension() { 77 let mut store = async_store(); 78 let func_ty = FuncType::new(store.engine(), None, None); 79 let func = Func::new_async(&mut store, func_ty, move |_caller, _params, _results| { 80 Box::new(async { 81 tokio::task::yield_now().await; 82 Ok(()) 83 }) 84 }); 85 run_smoke_test(&mut store, func).await; 86 run_smoke_typed_test(&mut store, func).await; 87 88 let func = Func::wrap_async(&mut store, move |_caller, _: ()| { 89 Box::new(async { 90 tokio::task::yield_now().await; 91 Ok(()) 92 }) 93 }); 94 run_smoke_test(&mut store, func).await; 95 run_smoke_typed_test(&mut store, func).await; 96 } 97 98 #[tokio::test] 99 async fn smoke_host_func_with_suspension() -> Result<()> { 100 let mut store = async_store(); 101 let mut linker = Linker::new(store.engine()); 102 103 linker.func_new_async( 104 "", 105 "first", 106 FuncType::new(store.engine(), None, None), 107 move |_caller, _params, _results| { 108 Box::new(async { 109 tokio::task::yield_now().await; 110 Ok(()) 111 }) 112 }, 113 )?; 114 115 linker.func_wrap_async("", "second", move |_caller, _: ()| { 116 Box::new(async { 117 tokio::task::yield_now().await; 118 Ok(()) 119 }) 120 })?; 121 122 let func = linker 123 .get(&mut store, "", "first") 124 .unwrap() 125 .into_func() 126 .unwrap(); 127 run_smoke_test(&mut store, func).await; 128 run_smoke_typed_test(&mut store, func).await; 129 130 let func = linker 131 .get(&mut store, "", "second") 132 .unwrap() 133 .into_func() 134 .unwrap(); 135 run_smoke_test(&mut store, func).await; 136 run_smoke_typed_test(&mut store, func).await; 137 138 Ok(()) 139 } 140 141 #[tokio::test] 142 async fn recursive_call() { 143 let mut store = async_store(); 144 let func_ty = FuncType::new(store.engine(), None, None); 145 let async_wasm_func = Func::new_async(&mut store, func_ty, |_caller, _params, _results| { 146 Box::new(async { 147 tokio::task::yield_now().await; 148 Ok(()) 149 }) 150 }); 151 152 // Create an imported function which recursively invokes another wasm 153 // function asynchronously, although this one is just our own host function 154 // which suffices for this test. 155 let func_ty = FuncType::new(store.engine(), None, None); 156 let func2 = Func::new_async(&mut store, func_ty, move |mut caller, _params, _results| { 157 Box::new(async move { 158 async_wasm_func 159 .call_async(&mut caller, &[], &mut []) 160 .await?; 161 Ok(()) 162 }) 163 }); 164 165 // Create an instance which calls an async import twice. 166 let module = Module::new( 167 store.engine(), 168 " 169 (module 170 (import \"\" \"\" (func)) 171 (func (export \"\") 172 ;; call imported function which recursively does an async 173 ;; call 174 call 0 175 ;; do it again, and our various pointers all better align 176 call 0)) 177 ", 178 ) 179 .unwrap(); 180 181 let instance = Instance::new_async(&mut store, &module, &[func2.into()]) 182 .await 183 .unwrap(); 184 let func = instance.get_func(&mut store, "").unwrap(); 185 func.call_async(&mut store, &[], &mut []).await.unwrap(); 186 } 187 188 #[tokio::test] 189 async fn suspend_while_suspending() { 190 let mut store = async_store(); 191 192 // Create a synchronous function which calls our asynchronous function and 193 // runs it locally. This shouldn't generally happen but we know everything 194 // is synchronous in this test so it's fine for us to do this. 195 // 196 // The purpose of this test is intended to stress various cases in how 197 // we manage pointers in ways that are not necessarily common but are still 198 // possible in safe code. 199 let func_ty = FuncType::new(store.engine(), None, None); 200 let async_thunk = Func::new_async(&mut store, func_ty, |_caller, _params, _results| { 201 Box::new(async { Ok(()) }) 202 }); 203 let func_ty = FuncType::new(store.engine(), None, None); 204 let sync_call_async_thunk = 205 Func::new(&mut store, func_ty, move |mut caller, _params, _results| { 206 let mut future = Box::pin(async_thunk.call_async(&mut caller, &[], &mut [])); 207 let poll = future 208 .as_mut() 209 .poll(&mut Context::from_waker(&noop_waker())); 210 assert!(poll.is_ready()); 211 Ok(()) 212 }); 213 214 // A small async function that simply awaits once to pump the loops and 215 // then finishes. 216 let func_ty = FuncType::new(store.engine(), None, None); 217 let async_import = Func::new_async(&mut store, func_ty, move |_caller, _params, _results| { 218 Box::new(async move { 219 tokio::task::yield_now().await; 220 Ok(()) 221 }) 222 }); 223 224 let module = Module::new( 225 store.engine(), 226 " 227 (module 228 (import \"\" \"\" (func $sync_call_async_thunk)) 229 (import \"\" \"\" (func $async_import)) 230 (func (export \"\") 231 ;; Set some store-local state and pointers 232 call $sync_call_async_thunk 233 ;; .. and hopefully it's all still configured correctly 234 call $async_import)) 235 ", 236 ) 237 .unwrap(); 238 let instance = Instance::new_async( 239 &mut store, 240 &module, 241 &[sync_call_async_thunk.into(), async_import.into()], 242 ) 243 .await 244 .unwrap(); 245 let func = instance.get_func(&mut store, "").unwrap(); 246 func.call_async(&mut store, &[], &mut []).await.unwrap(); 247 } 248 249 #[tokio::test] 250 async fn cancel_during_run() { 251 let mut store = Store::new(&Engine::new(Config::new().async_support(true)).unwrap(), 0); 252 253 let func_ty = FuncType::new(store.engine(), None, None); 254 let async_thunk = Func::new_async(&mut store, func_ty, move |mut caller, _params, _results| { 255 assert_eq!(*caller.data(), 0); 256 *caller.data_mut() = 1; 257 let dtor = SetOnDrop(caller); 258 Box::new(async move { 259 // SetOnDrop is not destroyed when dropping the reference of it 260 // here. Instead, it is moved into the future where it's forced 261 // to live in and will be destroyed at the end of the future. 262 let _ = &dtor; 263 tokio::task::yield_now().await; 264 Ok(()) 265 }) 266 }); 267 // Shouldn't have called anything yet... 268 assert_eq!(*store.data(), 0); 269 270 // Create our future, but as per async conventions this still doesn't 271 // actually do anything. No wasm or host function has been called yet. 272 let future = Box::pin(async_thunk.call_async(&mut store, &[], &mut [])); 273 274 // Push the future forward one tick, which actually runs the host code in 275 // our async func. Our future is designed to be pending once, however. 276 let future = PollOnce::new(future).await; 277 278 // Now that our future is running (on a separate, now-suspended fiber), drop 279 // the future and that should deallocate all the Rust bits as well. 280 drop(future); 281 assert_eq!(*store.data(), 2); 282 283 struct SetOnDrop<'a>(Caller<'a, usize>); 284 285 impl Drop for SetOnDrop<'_> { 286 fn drop(&mut self) { 287 assert_eq!(*self.0.data(), 1); 288 *self.0.data_mut() = 2; 289 } 290 } 291 } 292 293 #[tokio::test] 294 async fn iloop_with_fuel() { 295 let engine = Engine::new(Config::new().async_support(true).consume_fuel(true)).unwrap(); 296 let mut store = Store::new(&engine, ()); 297 store.set_fuel(10_000).unwrap(); 298 store.fuel_async_yield_interval(Some(100)).unwrap(); 299 let module = Module::new( 300 &engine, 301 " 302 (module 303 (func (loop br 0)) 304 (start 0) 305 ) 306 ", 307 ) 308 .unwrap(); 309 let instance = Instance::new_async(&mut store, &module, &[]); 310 311 // This should yield a bunch of times but eventually finish 312 let (_, pending) = CountPending::new(Box::pin(instance)).await; 313 assert_eq!(pending, 99); 314 } 315 316 #[tokio::test] 317 async fn fuel_eventually_finishes() { 318 let engine = Engine::new(Config::new().async_support(true).consume_fuel(true)).unwrap(); 319 let mut store = Store::new(&engine, ()); 320 store.set_fuel(u64::MAX).unwrap(); 321 store.fuel_async_yield_interval(Some(10)).unwrap(); 322 let module = Module::new( 323 &engine, 324 " 325 (module 326 (func 327 (local i32) 328 i32.const 100 329 local.set 0 330 (loop 331 local.get 0 332 i32.const -1 333 i32.add 334 local.tee 0 335 br_if 0) 336 ) 337 (start 0) 338 ) 339 ", 340 ) 341 .unwrap(); 342 let instance = Instance::new_async(&mut store, &module, &[]); 343 instance.await.unwrap(); 344 } 345 346 #[tokio::test] 347 async fn async_with_pooling_stacks() { 348 let mut pool = crate::small_pool_config(); 349 pool.total_stacks(1) 350 .max_memory_size(1 << 16) 351 .table_elements(0); 352 let mut config = Config::new(); 353 config.async_support(true); 354 config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)); 355 config.dynamic_memory_guard_size(0); 356 config.static_memory_guard_size(0); 357 config.static_memory_maximum_size(1 << 16); 358 359 let engine = Engine::new(&config).unwrap(); 360 let mut store = Store::new(&engine, ()); 361 let func_ty = FuncType::new(store.engine(), None, None); 362 let func = Func::new_async(&mut store, func_ty, move |_caller, _params, _results| { 363 Box::new(async { Ok(()) }) 364 }); 365 366 run_smoke_test(&mut store, func).await; 367 run_smoke_typed_test(&mut store, func).await; 368 } 369 370 #[tokio::test] 371 async fn async_host_func_with_pooling_stacks() -> Result<()> { 372 let mut pooling = crate::small_pool_config(); 373 pooling 374 .total_stacks(1) 375 .max_memory_size(1 << 16) 376 .table_elements(0); 377 let mut config = Config::new(); 378 config.async_support(true); 379 config.allocation_strategy(InstanceAllocationStrategy::Pooling(pooling)); 380 config.dynamic_memory_guard_size(0); 381 config.static_memory_guard_size(0); 382 config.static_memory_maximum_size(1 << 16); 383 384 let mut store = Store::new(&Engine::new(&config)?, ()); 385 let mut linker = Linker::new(store.engine()); 386 linker.func_new_async( 387 "", 388 "", 389 FuncType::new(store.engine(), None, None), 390 move |_caller, _params, _results| Box::new(async { Ok(()) }), 391 )?; 392 393 let func = linker.get(&mut store, "", "").unwrap().into_func().unwrap(); 394 run_smoke_test(&mut store, func).await; 395 run_smoke_typed_test(&mut store, func).await; 396 Ok(()) 397 } 398 399 #[tokio::test] 400 async fn async_mpk_protection() -> Result<()> { 401 let _ = env_logger::try_init(); 402 403 // Construct a pool with MPK protection enabled; note that the MPK 404 // protection is configured in `small_pool_config`. 405 let mut pooling = crate::small_pool_config(); 406 pooling 407 .total_memories(10) 408 .total_stacks(2) 409 .max_memory_size(1 << 16) 410 .table_elements(0); 411 let mut config = Config::new(); 412 config.async_support(true); 413 config.allocation_strategy(InstanceAllocationStrategy::Pooling(pooling)); 414 config.static_memory_maximum_size(1 << 26); 415 config.epoch_interruption(true); 416 let engine = Engine::new(&config)?; 417 418 // Craft a module that loops for several iterations and checks whether it 419 // has access to its memory range (0x0-0x10000). 420 const WAT: &str = " 421 (module 422 (func $start 423 (local $i i32) 424 (local.set $i (i32.const 3)) 425 (loop $cont 426 (drop (i32.load (i32.const 0))) 427 (drop (i32.load (i32.const 0xfffc))) 428 (br_if $cont (local.tee $i (i32.sub (local.get $i) (i32.const 1)))))) 429 (memory 1) 430 (start $start)) 431 "; 432 433 // Start two instances of the module in separate fibers, `a` and `b`. 434 async fn run_instance(engine: &Engine, name: &str) -> Instance { 435 let mut store = Store::new(&engine, ()); 436 store.set_epoch_deadline(0); 437 store.epoch_deadline_async_yield_and_update(0); 438 let module = Module::new(store.engine(), WAT).unwrap(); 439 println!("[{name}] building instance"); 440 Instance::new_async(&mut store, &module, &[]).await.unwrap() 441 } 442 let mut a = Box::pin(run_instance(&engine, "a")); 443 let mut b = Box::pin(run_instance(&engine, "b")); 444 445 // Alternately poll each instance until completion. This should exercise 446 // fiber suspensions requiring the `Store` to appropriately save and restore 447 // the PKRU context between suspensions (see `AsyncCx::block_on`). 448 for i in 0..10 { 449 if i % 2 == 0 { 450 match PollOnce::new(a).await { 451 Ok(_) => { 452 println!("[a] done"); 453 break; 454 } 455 Err(a_) => { 456 println!("[a] not done"); 457 a = a_; 458 } 459 } 460 } else { 461 match PollOnce::new(b).await { 462 Ok(_) => { 463 println!("[b] done"); 464 break; 465 } 466 Err(b_) => { 467 println!("[b] not done"); 468 b = b_; 469 } 470 } 471 } 472 } 473 474 Ok(()) 475 } 476 477 /// This will execute the `future` provided to completion and each invocation of 478 /// `poll` for the future will be executed on a separate thread. 479 pub async fn execute_across_threads<F>(future: F) -> F::Output 480 where 481 F: Future + Send + 'static, 482 F::Output: Send, 483 { 484 let mut future = Box::pin(future); 485 loop { 486 let once = PollOnce::new(future); 487 let handle = tokio::runtime::Handle::current(); 488 let result = std::thread::spawn(move || handle.block_on(once)) 489 .join() 490 .unwrap(); 491 match result { 492 Ok(val) => break val, 493 Err(f) => future = f, 494 } 495 } 496 } 497 498 #[tokio::test] 499 async fn resume_separate_thread() { 500 // This test will poll the following future on two threads. Simulating a 501 // trap requires accessing TLS info, so that should be preserved correctly. 502 execute_across_threads(async { 503 let mut store = async_store(); 504 let module = Module::new( 505 store.engine(), 506 " 507 (module 508 (import \"\" \"\" (func)) 509 (start 0) 510 ) 511 ", 512 ) 513 .unwrap(); 514 let func = Func::wrap_async(&mut store, |_, _: ()| { 515 Box::new(async { 516 tokio::task::yield_now().await; 517 Err::<(), _>(anyhow!("test")) 518 }) 519 }); 520 let result = Instance::new_async(&mut store, &module, &[func.into()]).await; 521 assert!(result.is_err()); 522 }) 523 .await; 524 } 525 526 #[tokio::test] 527 async fn resume_separate_thread2() { 528 // This test will poll the following future on two threads. Catching a 529 // signal requires looking up TLS information to determine whether it's a 530 // trap to handle or not, so that must be preserved correctly across threads. 531 execute_across_threads(async { 532 let mut store = async_store(); 533 let module = Module::new( 534 store.engine(), 535 " 536 (module 537 (import \"\" \"\" (func)) 538 (func $start 539 call 0 540 unreachable) 541 (start $start) 542 ) 543 ", 544 ) 545 .unwrap(); 546 let func = Func::wrap_async(&mut store, |_, _: ()| { 547 Box::new(async { 548 tokio::task::yield_now().await; 549 }) 550 }); 551 let result = Instance::new_async(&mut store, &module, &[func.into()]).await; 552 assert!(result.is_err()); 553 }) 554 .await; 555 } 556 557 #[tokio::test] 558 async fn resume_separate_thread3() { 559 let _ = env_logger::try_init(); 560 561 // This test doesn't actually do anything with cross-thread polls, but 562 // instead it deals with scheduling futures at "odd" times. 563 // 564 // First we'll set up a *synchronous* call which will initialize TLS info. 565 // This call is simply to a host-defined function, but it still has the same 566 // "enter into wasm" semantics since it's just calling a trampoline. In this 567 // situation we'll set up the TLS info so it's in place while the body of 568 // the function executes... 569 let mut store = Store::new(&Engine::default(), None); 570 let f = Func::wrap(&mut store, move |mut caller: Caller<'_, _>| -> Result<()> { 571 // ... and the execution of this host-defined function (while the TLS 572 // info is initialized), will set up a recursive call into wasm. This 573 // recursive call will be done asynchronously so we can suspend it 574 // halfway through. 575 let f = async { 576 let mut store = async_store(); 577 let module = Module::new( 578 store.engine(), 579 " 580 (module 581 (import \"\" \"\" (func)) 582 (start 0) 583 ) 584 ", 585 ) 586 .unwrap(); 587 let func = Func::wrap_async(&mut store, |_, _: ()| { 588 Box::new(async { 589 tokio::task::yield_now().await; 590 }) 591 }); 592 drop(Instance::new_async(&mut store, &module, &[func.into()]).await); 593 unreachable!() 594 }; 595 let mut future = Box::pin(f); 596 let poll = future 597 .as_mut() 598 .poll(&mut Context::from_waker(&noop_waker())); 599 assert!(poll.is_pending()); 600 601 // ... so at this point our call into wasm is suspended. The call into 602 // wasm will have overwritten TLS info, and we sure hope that the 603 // information is restored at this point. Note that we squirrel away the 604 // future somewhere else to get dropped later. If we were to drop it 605 // here then we would reenter the future's suspended stack to clean it 606 // up, which would do more alterations of TLS information we're not 607 // testing here. 608 *caller.data_mut() = Some(future); 609 610 // ... all in all this function will need access to the original TLS 611 // information to raise the trap. This TLS information should be 612 // restored even though the asynchronous execution is suspended. 613 bail!("") 614 }); 615 assert!(f.call(&mut store, &[], &mut []).is_err()); 616 } 617 618 #[tokio::test] 619 async fn recursive_async() -> Result<()> { 620 let _ = env_logger::try_init(); 621 let mut store = async_store(); 622 let m = Module::new( 623 store.engine(), 624 "(module 625 (func (export \"overflow\") call 0) 626 (func (export \"normal\")) 627 )", 628 )?; 629 let i = Instance::new_async(&mut store, &m, &[]).await?; 630 let overflow = i.get_typed_func::<(), ()>(&mut store, "overflow")?; 631 let normal = i.get_typed_func::<(), ()>(&mut store, "normal")?; 632 let f2 = Func::wrap_async(&mut store, move |mut caller, _: ()| { 633 let normal = normal.clone(); 634 let overflow = overflow.clone(); 635 Box::new(async move { 636 // recursive async calls shouldn't immediately stack overflow... 637 normal.call_async(&mut caller, ()).await?; 638 639 // ... but calls that actually stack overflow should indeed stack 640 // overflow 641 let err = overflow 642 .call_async(&mut caller, ()) 643 .await 644 .unwrap_err() 645 .downcast::<Trap>()?; 646 assert_eq!(err, Trap::StackOverflow); 647 Ok(()) 648 }) 649 }); 650 f2.call_async(&mut store, &[], &mut []).await?; 651 Ok(()) 652 } 653 654 #[tokio::test] 655 async fn linker_module_command() -> Result<()> { 656 let mut store = async_store(); 657 let mut linker = Linker::new(store.engine()); 658 659 let module1 = Module::new( 660 store.engine(), 661 r#" 662 (module 663 (global $g (mut i32) (i32.const 0)) 664 665 (func (export "_start")) 666 667 (func (export "g") (result i32) 668 global.get $g 669 i32.const 1 670 global.set $g) 671 ) 672 "#, 673 )?; 674 675 let module2 = Module::new( 676 store.engine(), 677 r#" 678 (module 679 (import "" "g" (func (result i32))) 680 681 (func (export "get") (result i32) 682 call 0) 683 ) 684 "#, 685 )?; 686 687 linker.module_async(&mut store, "", &module1).await?; 688 let instance = linker.instantiate_async(&mut store, &module2).await?; 689 let f = instance.get_typed_func::<(), i32>(&mut store, "get")?; 690 assert_eq!(f.call_async(&mut store, ()).await?, 0); 691 assert_eq!(f.call_async(&mut store, ()).await?, 0); 692 693 Ok(()) 694 } 695 696 #[tokio::test] 697 async fn linker_module_reactor() -> Result<()> { 698 let mut store = async_store(); 699 let mut linker = Linker::new(store.engine()); 700 let module1 = Module::new( 701 store.engine(), 702 r#" 703 (module 704 (global $g (mut i32) (i32.const 0)) 705 706 (func (export "g") (result i32) 707 global.get $g 708 i32.const 1 709 global.set $g) 710 ) 711 "#, 712 )?; 713 let module2 = Module::new( 714 store.engine(), 715 r#" 716 (module 717 (import "" "g" (func (result i32))) 718 719 (func (export "get") (result i32) 720 call 0) 721 ) 722 "#, 723 )?; 724 725 linker.module_async(&mut store, "", &module1).await?; 726 let instance = linker.instantiate_async(&mut store, &module2).await?; 727 let f = instance.get_typed_func::<(), i32>(&mut store, "get")?; 728 assert_eq!(f.call_async(&mut store, ()).await?, 0); 729 assert_eq!(f.call_async(&mut store, ()).await?, 1); 730 731 Ok(()) 732 } 733 734 pub struct CountPending<F> { 735 future: F, 736 yields: usize, 737 } 738 739 impl<F> CountPending<F> { 740 pub fn new(future: F) -> CountPending<F> { 741 CountPending { future, yields: 0 } 742 } 743 } 744 745 impl<F> Future for CountPending<F> 746 where 747 F: Future + Unpin, 748 { 749 type Output = (F::Output, usize); 750 751 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { 752 match Pin::new(&mut self.future).poll(cx) { 753 Poll::Pending => { 754 self.yields += 1; 755 Poll::Pending 756 } 757 Poll::Ready(e) => Poll::Ready((e, self.yields)), 758 } 759 } 760 } 761 762 pub struct PollOnce<F>(Option<F>); 763 764 impl<F> PollOnce<F> { 765 pub fn new(future: F) -> PollOnce<F> { 766 PollOnce(Some(future)) 767 } 768 } 769 770 impl<F> Future for PollOnce<F> 771 where 772 F: Future + Unpin, 773 { 774 type Output = Result<F::Output, F>; 775 776 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { 777 let mut future = self.0.take().unwrap(); 778 match Pin::new(&mut future).poll(cx) { 779 Poll::Pending => Poll::Ready(Err(future)), 780 Poll::Ready(val) => Poll::Ready(Ok(val)), 781 } 782 } 783 } 784 785 fn noop_waker() -> Waker { 786 const VTABLE: RawWakerVTable = 787 RawWakerVTable::new(|ptr| RawWaker::new(ptr, &VTABLE), |_| {}, |_| {}, |_| {}); 788 const RAW: RawWaker = RawWaker::new(0 as *const (), &VTABLE); 789 unsafe { Waker::from_raw(RAW) } 790 } 791 792 #[tokio::test] 793 async fn non_stacky_async_activations() -> Result<()> { 794 let mut config = Config::new(); 795 config.async_support(true); 796 let engine = Engine::new(&config)?; 797 let mut store1: Store<Option<Pin<Box<dyn Future<Output = Result<()>> + Send>>>> = 798 Store::new(&engine, None); 799 let mut linker1 = Linker::new(&engine); 800 801 let module1 = Module::new( 802 &engine, 803 r#" 804 (module $m1 805 (import "" "host_capture_stack" (func $host_capture_stack)) 806 (import "" "start_async_instance" (func $start_async_instance)) 807 (func $capture_stack (export "capture_stack") 808 call $host_capture_stack 809 ) 810 (func $run_sync (export "run_sync") 811 call $start_async_instance 812 ) 813 ) 814 "#, 815 )?; 816 817 let module2 = Module::new( 818 &engine, 819 r#" 820 (module $m2 821 (import "" "yield" (func $yield)) 822 823 (func $run_async (export "run_async") 824 call $yield 825 ) 826 ) 827 "#, 828 )?; 829 830 let stacks = Arc::new(Mutex::new(vec![])); 831 fn capture_stack(stacks: &Arc<Mutex<Vec<WasmBacktrace>>>, store: impl AsContext) { 832 let mut stacks = stacks.lock().unwrap(); 833 stacks.push(wasmtime::WasmBacktrace::force_capture(store)); 834 } 835 836 linker1.func_wrap_async("", "host_capture_stack", { 837 let stacks = stacks.clone(); 838 move |caller, _: ()| { 839 capture_stack(&stacks, &caller); 840 Box::new(async { Ok(()) }) 841 } 842 })?; 843 844 linker1.func_wrap_async("", "start_async_instance", { 845 let stacks = stacks.clone(); 846 move |mut caller, _: ()| { 847 let stacks = stacks.clone(); 848 capture_stack(&stacks, &caller); 849 850 let module2 = module2.clone(); 851 let mut store2 = Store::new(caller.engine(), ()); 852 let mut linker2 = Linker::new(caller.engine()); 853 linker2 854 .func_wrap_async("", "yield", { 855 let stacks = stacks.clone(); 856 move |caller, _: ()| { 857 let stacks = stacks.clone(); 858 Box::new(async move { 859 capture_stack(&stacks, &caller); 860 tokio::task::yield_now().await; 861 capture_stack(&stacks, &caller); 862 Ok(()) 863 }) 864 } 865 }) 866 .unwrap(); 867 868 Box::new(async move { 869 let future = PollOnce::new(Box::pin({ 870 let stacks = stacks.clone(); 871 async move { 872 let instance2 = linker2.instantiate_async(&mut store2, &module2).await?; 873 874 instance2 875 .get_func(&mut store2, "run_async") 876 .unwrap() 877 .call_async(&mut store2, &[], &mut []) 878 .await?; 879 880 capture_stack(&stacks, &store2); 881 Ok(()) 882 } 883 }) as _) 884 .await 885 .err() 886 .unwrap(); 887 capture_stack(&stacks, &caller); 888 *caller.data_mut() = Some(future); 889 Ok(()) 890 }) 891 } 892 })?; 893 894 let instance1 = linker1.instantiate_async(&mut store1, &module1).await?; 895 instance1 896 .get_typed_func::<(), ()>(&mut store1, "run_sync")? 897 .call_async(&mut store1, ()) 898 .await?; 899 let future = store1.data_mut().take().unwrap(); 900 future.await?; 901 902 instance1 903 .get_typed_func::<(), ()>(&mut store1, "capture_stack")? 904 .call_async(&mut store1, ()) 905 .await?; 906 907 let stacks = stacks.lock().unwrap(); 908 eprintln!("stacks = {stacks:#?}"); 909 910 assert_eq!(stacks.len(), 6); 911 for (actual, expected) in stacks.iter().zip(vec![ 912 vec!["run_sync"], 913 vec!["run_async"], 914 vec!["run_sync"], 915 vec!["run_async"], 916 vec![], 917 vec!["capture_stack"], 918 ]) { 919 eprintln!("expected = {expected:?}"); 920 eprintln!("actual = {actual:?}"); 921 assert_eq!(actual.frames().len(), expected.len()); 922 for (actual, expected) in actual.frames().iter().zip(expected) { 923 assert_eq!(actual.func_name(), Some(expected)); 924 } 925 } 926 927 Ok(()) 928 } 929 930 #[tokio::test] 931 async fn gc_preserves_externref_on_historical_async_stacks() -> Result<()> { 932 let _ = env_logger::try_init(); 933 934 let mut config = Config::new(); 935 config.async_support(true); 936 let engine = Engine::new(&config)?; 937 938 let module = Module::new( 939 &engine, 940 r#" 941 (module $m1 942 (import "" "gc" (func $gc)) 943 (import "" "recurse" (func $recurse (param i32))) 944 (import "" "test" (func $test (param i32 externref))) 945 (func (export "run") (param i32 externref) 946 local.get 0 947 if 948 local.get 0 949 i32.const -1 950 i32.add 951 call $recurse 952 else 953 call $gc 954 end 955 956 local.get 0 957 local.get 1 958 call $test 959 ) 960 ) 961 "#, 962 )?; 963 964 type F = TypedFunc<(i32, Option<Rooted<ExternRef>>), ()>; 965 966 let mut store = Store::new(&engine, None); 967 let mut linker = Linker::<Option<F>>::new(&engine); 968 linker.func_wrap("", "gc", |mut cx: Caller<'_, _>| cx.gc())?; 969 linker.func_wrap( 970 "", 971 "test", 972 |cx: Caller<'_, _>, val: i32, handle: Option<Rooted<ExternRef>>| -> Result<()> { 973 assert_eq!(handle.unwrap().data(&cx)?.downcast_ref(), Some(&val)); 974 Ok(()) 975 }, 976 )?; 977 linker.func_wrap_async( 978 "", 979 "recurse", 980 |mut cx: Caller<'_, Option<F>>, (val,): (i32,)| { 981 let func = cx.data().clone().unwrap(); 982 Box::new(async move { 983 let r = Some(ExternRef::new(&mut cx, val)?); 984 Ok(func.call_async(&mut cx, (val, r)).await) 985 }) 986 }, 987 )?; 988 let instance = linker.instantiate_async(&mut store, &module).await?; 989 let func: F = instance.get_typed_func(&mut store, "run")?; 990 *store.data_mut() = Some(func.clone()); 991 992 let r = Some(ExternRef::new(&mut store, 5)?); 993 func.call_async(&mut store, (5, r)).await?; 994 995 Ok(()) 996 } 997 998 #[tokio::test] 999 async fn async_gc_with_func_new_and_func_wrap() -> Result<()> { 1000 let _ = env_logger::try_init(); 1001 1002 let mut config = Config::new(); 1003 config.async_support(true); 1004 let engine = Engine::new(&config)?; 1005 1006 let module = Module::new( 1007 &engine, 1008 r#" 1009 (module $m1 1010 (import "" "a" (func $a (result externref))) 1011 (import "" "b" (func $b (result externref))) 1012 1013 (table 2 funcref) 1014 (elem (i32.const 0) func $a $b) 1015 1016 (func (export "a") 1017 (call $call (i32.const 0)) 1018 ) 1019 (func (export "b") 1020 (call $call (i32.const 1)) 1021 ) 1022 1023 (func $call (param i32) 1024 (local $cnt i32) 1025 (loop $l 1026 (drop (call_indirect (result externref) (local.get 0))) 1027 (local.set $cnt (i32.add (local.get $cnt) (i32.const 1))) 1028 1029 (if (i32.lt_u (local.get $cnt) (i32.const 1000)) 1030 (then (br $l))) 1031 ) 1032 ) 1033 ) 1034 "#, 1035 )?; 1036 1037 let mut linker = Linker::new(&engine); 1038 linker.func_wrap("", "a", |mut cx: Caller<'_, _>| { 1039 Ok(Some(ExternRef::new(&mut cx, 100)?)) 1040 })?; 1041 let ty = FuncType::new(&engine, [], [ValType::EXTERNREF]); 1042 linker.func_new("", "b", ty, |mut cx, _, results| { 1043 results[0] = ExternRef::new(&mut cx, 100)?.into(); 1044 Ok(()) 1045 })?; 1046 1047 let mut store = Store::new(&engine, ()); 1048 let instance = linker.instantiate_async(&mut store, &module).await?; 1049 let a = instance.get_typed_func::<(), ()>(&mut store, "a")?; 1050 a.call_async(&mut store, ()).await?; 1051 1052 let mut store = Store::new(&engine, ()); 1053 let instance = linker.instantiate_async(&mut store, &module).await?; 1054 let b = instance.get_typed_func::<(), ()>(&mut store, "b")?; 1055 b.call_async(&mut store, ()).await?; 1056 1057 Ok(()) 1058 } 1059