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