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::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 let normal = normal.clone(); 627 let overflow = overflow.clone(); 628 Box::new(async move { 629 // recursive async calls shouldn't immediately stack overflow... 630 normal.call_async(&mut caller, ()).await?; 631 632 // ... but calls that actually stack overflow should indeed stack 633 // overflow 634 let err = overflow 635 .call_async(&mut caller, ()) 636 .await 637 .unwrap_err() 638 .downcast::<Trap>()?; 639 assert_eq!(err, Trap::StackOverflow); 640 Ok(()) 641 }) 642 }); 643 f2.call_async(&mut store, &[], &mut []).await?; 644 Ok(()) 645 } 646 647 #[tokio::test] 648 async fn linker_module_command() -> Result<()> { 649 let mut store = async_store(); 650 let mut linker = Linker::new(store.engine()); 651 652 let module1 = Module::new( 653 store.engine(), 654 r#" 655 (module 656 (global $g (mut i32) (i32.const 0)) 657 658 (func (export "_start")) 659 660 (func (export "g") (result i32) 661 global.get $g 662 i32.const 1 663 global.set $g) 664 ) 665 "#, 666 )?; 667 668 let module2 = Module::new( 669 store.engine(), 670 r#" 671 (module 672 (import "" "g" (func (result i32))) 673 674 (func (export "get") (result i32) 675 call 0) 676 ) 677 "#, 678 )?; 679 680 linker.module_async(&mut store, "", &module1).await?; 681 let instance = linker.instantiate_async(&mut store, &module2).await?; 682 let f = instance.get_typed_func::<(), i32>(&mut store, "get")?; 683 assert_eq!(f.call_async(&mut store, ()).await?, 0); 684 assert_eq!(f.call_async(&mut store, ()).await?, 0); 685 686 Ok(()) 687 } 688 689 #[tokio::test] 690 async fn linker_module_reactor() -> Result<()> { 691 let mut store = async_store(); 692 let mut linker = Linker::new(store.engine()); 693 let module1 = Module::new( 694 store.engine(), 695 r#" 696 (module 697 (global $g (mut i32) (i32.const 0)) 698 699 (func (export "g") (result i32) 700 global.get $g 701 i32.const 1 702 global.set $g) 703 ) 704 "#, 705 )?; 706 let module2 = Module::new( 707 store.engine(), 708 r#" 709 (module 710 (import "" "g" (func (result i32))) 711 712 (func (export "get") (result i32) 713 call 0) 714 ) 715 "#, 716 )?; 717 718 linker.module_async(&mut store, "", &module1).await?; 719 let instance = linker.instantiate_async(&mut store, &module2).await?; 720 let f = instance.get_typed_func::<(), i32>(&mut store, "get")?; 721 assert_eq!(f.call_async(&mut store, ()).await?, 0); 722 assert_eq!(f.call_async(&mut store, ()).await?, 1); 723 724 Ok(()) 725 } 726 727 pub struct CountPending<F> { 728 future: F, 729 yields: usize, 730 } 731 732 impl<F> CountPending<F> { 733 pub fn new(future: F) -> CountPending<F> { 734 CountPending { future, yields: 0 } 735 } 736 } 737 738 impl<F> Future for CountPending<F> 739 where 740 F: Future + Unpin, 741 { 742 type Output = (F::Output, usize); 743 744 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { 745 match Pin::new(&mut self.future).poll(cx) { 746 Poll::Pending => { 747 self.yields += 1; 748 Poll::Pending 749 } 750 Poll::Ready(e) => Poll::Ready((e, self.yields)), 751 } 752 } 753 } 754 755 pub struct PollOnce<F>(Option<F>); 756 757 impl<F> PollOnce<F> { 758 pub fn new(future: F) -> PollOnce<F> { 759 PollOnce(Some(future)) 760 } 761 } 762 763 impl<F> Future for PollOnce<F> 764 where 765 F: Future + Unpin, 766 { 767 type Output = Result<F::Output, F>; 768 769 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { 770 let mut future = self.0.take().unwrap(); 771 match Pin::new(&mut future).poll(cx) { 772 Poll::Pending => Poll::Ready(Err(future)), 773 Poll::Ready(val) => Poll::Ready(Ok(val)), 774 } 775 } 776 } 777 778 fn noop_waker() -> Waker { 779 const VTABLE: RawWakerVTable = 780 RawWakerVTable::new(|ptr| RawWaker::new(ptr, &VTABLE), |_| {}, |_| {}, |_| {}); 781 const RAW: RawWaker = RawWaker::new(0 as *const (), &VTABLE); 782 unsafe { Waker::from_raw(RAW) } 783 } 784 785 #[tokio::test] 786 async fn non_stacky_async_activations() -> Result<()> { 787 let mut config = Config::new(); 788 config.async_support(true); 789 let engine = Engine::new(&config)?; 790 let mut store1: Store<Option<Pin<Box<dyn Future<Output = Result<()>> + Send>>>> = 791 Store::new(&engine, None); 792 let mut linker1 = Linker::new(&engine); 793 794 let module1 = Module::new( 795 &engine, 796 r#" 797 (module $m1 798 (import "" "host_capture_stack" (func $host_capture_stack)) 799 (import "" "start_async_instance" (func $start_async_instance)) 800 (func $capture_stack (export "capture_stack") 801 call $host_capture_stack 802 ) 803 (func $run_sync (export "run_sync") 804 call $start_async_instance 805 ) 806 ) 807 "#, 808 )?; 809 810 let module2 = Module::new( 811 &engine, 812 r#" 813 (module $m2 814 (import "" "yield" (func $yield)) 815 816 (func $run_async (export "run_async") 817 call $yield 818 ) 819 ) 820 "#, 821 )?; 822 823 let stacks = Arc::new(Mutex::new(vec![])); 824 fn capture_stack(stacks: &Arc<Mutex<Vec<WasmBacktrace>>>, store: impl AsContext) { 825 let mut stacks = stacks.lock().unwrap(); 826 stacks.push(wasmtime::WasmBacktrace::force_capture(store)); 827 } 828 829 linker1.func_wrap0_async("", "host_capture_stack", { 830 let stacks = stacks.clone(); 831 move |caller| { 832 capture_stack(&stacks, &caller); 833 Box::new(async { Ok(()) }) 834 } 835 })?; 836 837 linker1.func_wrap0_async("", "start_async_instance", { 838 let stacks = stacks.clone(); 839 move |mut caller| { 840 let stacks = stacks.clone(); 841 capture_stack(&stacks, &caller); 842 843 let module2 = module2.clone(); 844 let mut store2 = Store::new(caller.engine(), ()); 845 let mut linker2 = Linker::new(caller.engine()); 846 linker2 847 .func_wrap0_async("", "yield", { 848 let stacks = stacks.clone(); 849 move |caller| { 850 let stacks = stacks.clone(); 851 Box::new(async move { 852 capture_stack(&stacks, &caller); 853 tokio::task::yield_now().await; 854 capture_stack(&stacks, &caller); 855 Ok(()) 856 }) 857 } 858 }) 859 .unwrap(); 860 861 Box::new(async move { 862 let future = PollOnce::new(Box::pin({ 863 let stacks = stacks.clone(); 864 async move { 865 let instance2 = linker2.instantiate_async(&mut store2, &module2).await?; 866 867 instance2 868 .get_func(&mut store2, "run_async") 869 .unwrap() 870 .call_async(&mut store2, &[], &mut []) 871 .await?; 872 873 capture_stack(&stacks, &store2); 874 Ok(()) 875 } 876 }) as _) 877 .await 878 .err() 879 .unwrap(); 880 capture_stack(&stacks, &caller); 881 *caller.data_mut() = Some(future); 882 Ok(()) 883 }) 884 } 885 })?; 886 887 let instance1 = linker1.instantiate_async(&mut store1, &module1).await?; 888 instance1 889 .get_typed_func::<(), ()>(&mut store1, "run_sync")? 890 .call_async(&mut store1, ()) 891 .await?; 892 let future = store1.data_mut().take().unwrap(); 893 future.await?; 894 895 instance1 896 .get_typed_func::<(), ()>(&mut store1, "capture_stack")? 897 .call_async(&mut store1, ()) 898 .await?; 899 900 let stacks = stacks.lock().unwrap(); 901 eprintln!("stacks = {stacks:#?}"); 902 903 assert_eq!(stacks.len(), 6); 904 for (actual, expected) in stacks.iter().zip(vec![ 905 vec!["run_sync"], 906 vec!["run_async"], 907 vec!["run_sync"], 908 vec!["run_async"], 909 vec![], 910 vec!["capture_stack"], 911 ]) { 912 eprintln!("expected = {expected:?}"); 913 eprintln!("actual = {actual:?}"); 914 assert_eq!(actual.frames().len(), expected.len()); 915 for (actual, expected) in actual.frames().iter().zip(expected) { 916 assert_eq!(actual.func_name(), Some(expected)); 917 } 918 } 919 920 Ok(()) 921 } 922 923 #[tokio::test] 924 async fn gc_preserves_externref_on_historical_async_stacks() -> Result<()> { 925 let _ = env_logger::try_init(); 926 927 let mut config = Config::new(); 928 config.async_support(true); 929 let engine = Engine::new(&config)?; 930 931 let module = Module::new( 932 &engine, 933 r#" 934 (module $m1 935 (import "" "gc" (func $gc)) 936 (import "" "recurse" (func $recurse (param i32))) 937 (import "" "test" (func $test (param i32 externref))) 938 (func (export "run") (param i32 externref) 939 local.get 0 940 if 941 local.get 0 942 i32.const -1 943 i32.add 944 call $recurse 945 else 946 call $gc 947 end 948 949 local.get 0 950 local.get 1 951 call $test 952 ) 953 ) 954 "#, 955 )?; 956 957 type F = TypedFunc<(i32, Option<Rooted<ExternRef>>), ()>; 958 959 let mut store = Store::new(&engine, None); 960 let mut linker = Linker::<Option<F>>::new(&engine); 961 linker.func_wrap("", "gc", |mut cx: Caller<'_, _>| cx.gc())?; 962 linker.func_wrap( 963 "", 964 "test", 965 |cx: Caller<'_, _>, val: i32, handle: Option<Rooted<ExternRef>>| -> Result<()> { 966 assert_eq!(handle.unwrap().data(&cx)?.downcast_ref(), Some(&val)); 967 Ok(()) 968 }, 969 )?; 970 linker.func_wrap1_async("", "recurse", |mut cx: Caller<'_, _>, val: i32| { 971 let func = cx.data().clone().unwrap(); 972 let r = Some(ExternRef::new(&mut cx, val)); 973 Box::new(async move { func.call_async(&mut cx, (val, r)).await }) 974 })?; 975 let instance = linker.instantiate_async(&mut store, &module).await?; 976 let func: F = instance.get_typed_func(&mut store, "run")?; 977 *store.data_mut() = Some(func.clone()); 978 979 let r = Some(ExternRef::new(&mut store, 5)); 980 func.call_async(&mut store, (5, r)).await?; 981 982 Ok(()) 983 } 984