1 use super::util::test_run;
2 use crate::scenario::util::{config, make_component};
3 use component_async_tests::util;
4 use component_async_tests::{Ctx, yield_};
5 use std::future;
6 use std::pin::pin;
7 use std::task::Poll;
8 use wasmtime::Result;
9 use wasmtime::component::{Linker, ResourceTable};
10 use wasmtime::{AsContextMut, Engine, Store, StoreContextMut};
11 use wasmtime_wasi::WasiCtxBuilder;
12
13 mod yield_post_return {
14 wasmtime::component::bindgen!({
15 path: "wit",
16 world: "yield-post-return-callee",
17 });
18 }
19
20 // No-op function; we only test this by composing it in `async_post_return_caller`
21 #[allow(
22 dead_code,
23 reason = "here only to make the `assert_test_exists` macro happy"
24 )]
async_post_return_callee()25 pub fn async_post_return_callee() {}
26
27 #[tokio::test]
async_post_return_caller() -> Result<()>28 pub async fn async_post_return_caller() -> Result<()> {
29 test_run(&[
30 test_programs_artifacts::ASYNC_POST_RETURN_CALLER_COMPONENT,
31 test_programs_artifacts::ASYNC_POST_RETURN_CALLEE_COMPONENT,
32 ])
33 .await
34 }
35
36 #[tokio::test]
async_yield_post_return_caller() -> Result<()>37 pub async fn async_yield_post_return_caller() -> Result<()> {
38 test_yield_post_return(&[
39 test_programs_artifacts::ASYNC_YIELD_POST_RETURN_CALLER_COMPONENT,
40 test_programs_artifacts::ASYNC_YIELD_POST_RETURN_CALLEE_COMPONENT,
41 ])
42 .await
43 }
44
45 #[tokio::test]
async_yield_post_return_callee() -> Result<()>46 pub async fn async_yield_post_return_callee() -> Result<()> {
47 test_yield_post_return(&[test_programs_artifacts::ASYNC_YIELD_POST_RETURN_CALLEE_COMPONENT])
48 .await
49 }
50
test_yield_post_return(components: &[&str]) -> Result<()>51 async fn test_yield_post_return(components: &[&str]) -> Result<()> {
52 let engine = Engine::new(&config())?;
53
54 let component = make_component(&engine, components).await?;
55
56 let mut linker = Linker::new(&engine);
57
58 wasmtime_wasi::p2::add_to_linker_async(&mut linker)?;
59 yield_::local::local::yield_::add_to_linker::<_, Ctx>(&mut linker, |ctx| ctx)?;
60
61 let mut store = Store::new(
62 &engine,
63 Ctx {
64 wasi: WasiCtxBuilder::new().inherit_stdio().build(),
65 table: ResourceTable::default(),
66 continue_: false,
67 },
68 );
69
70 let guest = yield_post_return::YieldPostReturnCallee::instantiate_async(
71 &mut store, &component, &linker,
72 )
73 .await?;
74
75 async fn run(
76 store: StoreContextMut<'_, Ctx>,
77 guest: &yield_post_return::YieldPostReturnCallee,
78 ) -> Result<()> {
79 store
80 .run_concurrent(async |accessor| {
81 guest
82 .local_local_yield_post_return()
83 .call_run(accessor, 5)
84 .await?;
85
86 // Go idle for a bit before doing it again. This tests that
87 // `StoreContextMut::run_concurrent` is okay with having no
88 // outstanding guest or host tasks to poll for a while, trusting
89 // that we'll resolve the future independently, with or without
90 // giving it more work to do.
91 util::yield_times(100).await;
92
93 guest
94 .local_local_yield_post_return()
95 .call_run(accessor, 5)
96 .await?;
97
98 // Wait again for everything to be flushed to ensure that the
99 // assertion the store is empty below is correct.
100 util::yield_times(100).await;
101
102 wasmtime::error::Ok(())
103 })
104 .await?
105 }
106
107 run(store.as_context_mut(), &guest).await?;
108 // At this point, all subtasks should have exited, meaning no waitables,
109 // tasks, or other concurrent state should remain present in the store.
110 store.assert_concurrent_state_empty();
111
112 // Do it again, but this time cancel the event loop before it exits:
113 assert!(
114 future::poll_fn(|cx| Poll::Ready(pin!(run(store.as_context_mut(), &guest)).poll(cx)))
115 .await
116 .is_pending()
117 );
118
119 // Assuming the event loop is cancel-safe, this should complete without
120 // errors or panics:
121 run(store.as_context_mut(), &guest).await?;
122
123 Ok(())
124 }
125