1 use std::future::Future;
2 use std::pin::Pin;
3 use std::task::{self, Context, Poll};
4
5 use super::util::{config, make_component, test_run, test_run_with_count};
6 use cancel::exports::local::local::cancel::Mode;
7 use component_async_tests::transmit::bindings::exports::local::local::transmit::Control;
8 use component_async_tests::util::{OneshotConsumer, OneshotProducer, PipeConsumer, PipeProducer};
9 use component_async_tests::{Ctx, transmit, yield_};
10 use futures::{
11 FutureExt, SinkExt, StreamExt, TryStreamExt,
12 channel::{mpsc, oneshot},
13 stream::FuturesUnordered,
14 };
15 use wasmtime::component::{
16 Accessor, Component, Destination, FutureConsumer, FutureProducer, FutureReader, HasSelf,
17 Instance, Linker, ResourceTable, Source, StreamConsumer, StreamProducer, StreamReader,
18 StreamResult, Val,
19 };
20 use wasmtime::{AsContextMut, Engine, Result, Store, StoreContextMut, format_err};
21 use wasmtime_wasi::WasiCtxBuilder;
22
23 struct BufferStreamProducer {
24 buffer: Vec<u8>,
25 }
26
27 impl<D> StreamProducer<D> for BufferStreamProducer {
28 type Item = u8;
29 type Buffer = Option<u8>;
30
poll_produce<'a>( self: Pin<&mut Self>, _: &mut Context<'_>, mut store: StoreContextMut<'a, D>, destination: Destination<'a, Self::Item, Self::Buffer>, _: bool, ) -> Poll<Result<StreamResult>>31 fn poll_produce<'a>(
32 self: Pin<&mut Self>,
33 _: &mut Context<'_>,
34 mut store: StoreContextMut<'a, D>,
35 destination: Destination<'a, Self::Item, Self::Buffer>,
36 _: bool,
37 ) -> Poll<Result<StreamResult>> {
38 let me = self.get_mut();
39 let capacity = destination.remaining(store.as_context_mut());
40 if capacity == Some(0) {
41 Poll::Ready(Ok(StreamResult::Completed))
42 } else {
43 assert_eq!(capacity, Some(me.buffer.len()));
44 let mut destination = destination.as_direct(store, me.buffer.len());
45 destination.remaining().copy_from_slice(&me.buffer);
46 destination.mark_written(me.buffer.len());
47
48 Poll::Ready(Ok(StreamResult::Dropped))
49 }
50 }
51 }
52
53 struct BufferStreamConsumer {
54 expected: Vec<u8>,
55 }
56
57 impl<D> StreamConsumer<D> for BufferStreamConsumer {
58 type Item = u8;
59
poll_consume( self: Pin<&mut Self>, _: &mut Context<'_>, mut store: StoreContextMut<D>, source: Source<Self::Item>, _: bool, ) -> Poll<Result<StreamResult>>60 fn poll_consume(
61 self: Pin<&mut Self>,
62 _: &mut Context<'_>,
63 mut store: StoreContextMut<D>,
64 source: Source<Self::Item>,
65 _: bool,
66 ) -> Poll<Result<StreamResult>> {
67 let me = self.get_mut();
68 let available = source.remaining(store.as_context_mut());
69 if available == 0 {
70 Poll::Ready(Ok(StreamResult::Completed))
71 } else {
72 assert_eq!(available, me.expected.len());
73 let mut source = source.as_direct(store);
74 assert_eq!(&me.expected, source.remaining());
75 source.mark_read(me.expected.len());
76
77 Poll::Ready(Ok(StreamResult::Dropped))
78 }
79 }
80 }
81
82 struct ValueFutureProducer {
83 value: u8,
84 }
85
86 impl<D> FutureProducer<D> for ValueFutureProducer {
87 type Item = u8;
88
poll_produce<'a>( self: Pin<&mut Self>, _: &mut Context<'_>, _: StoreContextMut<'a, D>, _: bool, ) -> Poll<Result<Option<Self::Item>>>89 fn poll_produce<'a>(
90 self: Pin<&mut Self>,
91 _: &mut Context<'_>,
92 _: StoreContextMut<'a, D>,
93 _: bool,
94 ) -> Poll<Result<Option<Self::Item>>> {
95 Poll::Ready(Ok(Some(self.value)))
96 }
97 }
98
99 struct ValueFutureConsumer {
100 expected: u8,
101 }
102
103 impl<D> FutureConsumer<D> for ValueFutureConsumer {
104 type Item = u8;
105
poll_consume( self: Pin<&mut Self>, _: &mut Context<'_>, store: StoreContextMut<D>, mut source: Source<'_, Self::Item>, _: bool, ) -> Poll<Result<()>>106 fn poll_consume(
107 self: Pin<&mut Self>,
108 _: &mut Context<'_>,
109 store: StoreContextMut<D>,
110 mut source: Source<'_, Self::Item>,
111 _: bool,
112 ) -> Poll<Result<()>> {
113 let value = &mut None;
114 source.read(store, value)?;
115 assert_eq!(value.take(), Some(self.expected));
116 Poll::Ready(Ok(()))
117 }
118 }
119
120 struct DelayedStreamProducer<P> {
121 inner: P,
122 maybe_yield: Pin<Box<dyn Future<Output = ()> + Send>>,
123 }
124
125 impl<D, P: StreamProducer<D>> StreamProducer<D> for DelayedStreamProducer<P> {
126 type Item = P::Item;
127 type Buffer = P::Buffer;
128
poll_produce<'a>( mut self: Pin<&mut Self>, cx: &mut Context<'_>, store: StoreContextMut<'a, D>, destination: Destination<'a, Self::Item, Self::Buffer>, finish: bool, ) -> Poll<Result<StreamResult>>129 fn poll_produce<'a>(
130 mut self: Pin<&mut Self>,
131 cx: &mut Context<'_>,
132 store: StoreContextMut<'a, D>,
133 destination: Destination<'a, Self::Item, Self::Buffer>,
134 finish: bool,
135 ) -> Poll<Result<StreamResult>> {
136 // SAFETY: We never move out of `self`.
137 let maybe_yield = unsafe { &mut self.as_mut().get_unchecked_mut().maybe_yield };
138 task::ready!(maybe_yield.as_mut().poll(cx));
139 *maybe_yield = async {}.boxed();
140
141 // SAFETY: This is a standard pin-projection, and we never move out
142 // of `self`.
143 let inner = unsafe { self.map_unchecked_mut(|v| &mut v.inner) };
144 inner.poll_produce(cx, store, destination, finish)
145 }
146 }
147
148 struct DelayedStreamConsumer<C> {
149 inner: C,
150 maybe_yield: Pin<Box<dyn Future<Output = ()> + Send>>,
151 }
152
153 impl<D, C: StreamConsumer<D>> StreamConsumer<D> for DelayedStreamConsumer<C> {
154 type Item = C::Item;
155
poll_consume( mut self: Pin<&mut Self>, cx: &mut Context<'_>, store: StoreContextMut<D>, source: Source<Self::Item>, finish: bool, ) -> Poll<Result<StreamResult>>156 fn poll_consume(
157 mut self: Pin<&mut Self>,
158 cx: &mut Context<'_>,
159 store: StoreContextMut<D>,
160 source: Source<Self::Item>,
161 finish: bool,
162 ) -> Poll<Result<StreamResult>> {
163 // SAFETY: We never move out of `self`.
164 let maybe_yield = unsafe { &mut self.as_mut().get_unchecked_mut().maybe_yield };
165 task::ready!(maybe_yield.as_mut().poll(cx));
166 *maybe_yield = async {}.boxed();
167
168 // SAFETY: This is a standard pin-projection, and we never move out
169 // of `self`.
170 let inner = unsafe { self.map_unchecked_mut(|v| &mut v.inner) };
171 inner.poll_consume(cx, store, source, finish)
172 }
173 }
174
175 struct DelayedFutureProducer<P> {
176 inner: P,
177 maybe_yield: Pin<Box<dyn Future<Output = ()> + Send>>,
178 }
179
180 impl<D, P: FutureProducer<D>> FutureProducer<D> for DelayedFutureProducer<P> {
181 type Item = P::Item;
182
poll_produce<'a>( mut self: Pin<&mut Self>, cx: &mut Context<'_>, store: StoreContextMut<'a, D>, finish: bool, ) -> Poll<Result<Option<Self::Item>>>183 fn poll_produce<'a>(
184 mut self: Pin<&mut Self>,
185 cx: &mut Context<'_>,
186 store: StoreContextMut<'a, D>,
187 finish: bool,
188 ) -> Poll<Result<Option<Self::Item>>> {
189 // SAFETY: We never move out of `self`.
190 let maybe_yield = unsafe { &mut self.as_mut().get_unchecked_mut().maybe_yield };
191 task::ready!(maybe_yield.as_mut().poll(cx));
192 *maybe_yield = async {}.boxed();
193
194 // SAFETY: This is a standard pin-projection, and we never move out
195 // of `self`.
196 let inner = unsafe { self.map_unchecked_mut(|v| &mut v.inner) };
197 inner.poll_produce(cx, store, finish)
198 }
199 }
200
201 struct DelayedFutureConsumer<C> {
202 inner: C,
203 maybe_yield: Pin<Box<dyn Future<Output = ()> + Send>>,
204 }
205
206 impl<D, C: FutureConsumer<D>> FutureConsumer<D> for DelayedFutureConsumer<C> {
207 type Item = C::Item;
208
poll_consume( mut self: Pin<&mut Self>, cx: &mut Context<'_>, store: StoreContextMut<D>, source: Source<'_, Self::Item>, finish: bool, ) -> Poll<Result<()>>209 fn poll_consume(
210 mut self: Pin<&mut Self>,
211 cx: &mut Context<'_>,
212 store: StoreContextMut<D>,
213 source: Source<'_, Self::Item>,
214 finish: bool,
215 ) -> Poll<Result<()>> {
216 // SAFETY: We never move out of `self`.
217 let maybe_yield = unsafe { &mut self.as_mut().get_unchecked_mut().maybe_yield };
218 task::ready!(maybe_yield.as_mut().poll(cx));
219 *maybe_yield = async {}.boxed();
220
221 // SAFETY: This is a standard pin-projection, and we never move out
222 // of `self`.
223 let inner = unsafe { self.map_unchecked_mut(|v| &mut v.inner) };
224 inner.poll_consume(cx, store, source, finish)
225 }
226 }
227
228 struct ProcrastinatingStreamProducer<P>(P);
229
230 impl<D, P: StreamProducer<D>> StreamProducer<D> for ProcrastinatingStreamProducer<P> {
231 type Item = P::Item;
232 type Buffer = P::Buffer;
233
poll_produce<'a>( self: Pin<&mut Self>, cx: &mut Context<'_>, store: StoreContextMut<'a, D>, destination: Destination<'a, Self::Item, Self::Buffer>, finish: bool, ) -> Poll<Result<StreamResult>>234 fn poll_produce<'a>(
235 self: Pin<&mut Self>,
236 cx: &mut Context<'_>,
237 store: StoreContextMut<'a, D>,
238 destination: Destination<'a, Self::Item, Self::Buffer>,
239 finish: bool,
240 ) -> Poll<Result<StreamResult>> {
241 if finish {
242 // SAFETY: This is a standard pin-projection, and we never move out
243 // of `self`.
244 let producer = unsafe { self.map_unchecked_mut(|v| &mut v.0) };
245 producer.poll_produce(cx, store, destination, true)
246 } else {
247 Poll::Pending
248 }
249 }
250 }
251
252 struct ProcrastinatingStreamConsumer<C>(C);
253
254 impl<D, C: StreamConsumer<D>> StreamConsumer<D> for ProcrastinatingStreamConsumer<C> {
255 type Item = C::Item;
256
poll_consume( self: Pin<&mut Self>, cx: &mut Context<'_>, store: StoreContextMut<D>, source: Source<'_, Self::Item>, finish: bool, ) -> Poll<Result<StreamResult>>257 fn poll_consume(
258 self: Pin<&mut Self>,
259 cx: &mut Context<'_>,
260 store: StoreContextMut<D>,
261 source: Source<'_, Self::Item>,
262 finish: bool,
263 ) -> Poll<Result<StreamResult>> {
264 if finish {
265 // SAFETY: This is a standard pin-projection, and we never move out
266 // of `self`.
267 let consumer = unsafe { self.map_unchecked_mut(|v| &mut v.0) };
268 consumer.poll_consume(cx, store, source, true)
269 } else {
270 Poll::Pending
271 }
272 }
273 }
274
275 struct ProcrastinatingFutureProducer<P>(P);
276
277 impl<D, P: FutureProducer<D>> FutureProducer<D> for ProcrastinatingFutureProducer<P> {
278 type Item = P::Item;
279
poll_produce<'a>( self: Pin<&mut Self>, cx: &mut Context<'_>, store: StoreContextMut<'a, D>, finish: bool, ) -> Poll<Result<Option<Self::Item>>>280 fn poll_produce<'a>(
281 self: Pin<&mut Self>,
282 cx: &mut Context<'_>,
283 store: StoreContextMut<'a, D>,
284 finish: bool,
285 ) -> Poll<Result<Option<Self::Item>>> {
286 if finish {
287 // SAFETY: This is a standard pin-projection, and we never move out
288 // of `self`.
289 let producer = unsafe { self.map_unchecked_mut(|v| &mut v.0) };
290 producer.poll_produce(cx, store, true)
291 } else {
292 Poll::Pending
293 }
294 }
295 }
296
297 struct ProcrastinatingFutureConsumer<C>(C);
298
299 impl<D, C: FutureConsumer<D>> FutureConsumer<D> for ProcrastinatingFutureConsumer<C> {
300 type Item = C::Item;
301
poll_consume( self: Pin<&mut Self>, cx: &mut Context<'_>, store: StoreContextMut<D>, source: Source<'_, Self::Item>, finish: bool, ) -> Poll<Result<()>>302 fn poll_consume(
303 self: Pin<&mut Self>,
304 cx: &mut Context<'_>,
305 store: StoreContextMut<D>,
306 source: Source<'_, Self::Item>,
307 finish: bool,
308 ) -> Poll<Result<()>> {
309 if finish {
310 // SAFETY: This is a standard pin-projection, and we never move out
311 // of `self`.
312 let consumer = unsafe { self.map_unchecked_mut(|v| &mut v.0) };
313 consumer.poll_consume(cx, store, source, true)
314 } else {
315 Poll::Pending
316 }
317 }
318 }
319
yield_times(n: usize)320 async fn yield_times(n: usize) {
321 for _ in 0..n {
322 tokio::task::yield_now().await;
323 }
324 }
325
326 mod readiness {
327 wasmtime::component::bindgen!({
328 path: "wit",
329 world: "readiness-guest",
330 });
331 }
332
333 #[tokio::test]
async_readiness() -> Result<()>334 pub async fn async_readiness() -> Result<()> {
335 let component = test_programs_artifacts::ASYNC_READINESS_COMPONENT;
336
337 let engine = Engine::new(&config())?;
338
339 let component = make_component(&engine, &[component]).await?;
340
341 let mut linker = Linker::new(&engine);
342
343 wasmtime_wasi::p2::add_to_linker_async(&mut linker)?;
344
345 let mut store = Store::new(
346 &engine,
347 Ctx {
348 wasi: WasiCtxBuilder::new().inherit_stdio().build(),
349 table: ResourceTable::default(),
350 continue_: false,
351 },
352 );
353
354 let readiness_guest =
355 readiness::ReadinessGuest::instantiate_async(&mut store, &component, &linker).await?;
356 let expected = vec![2u8, 4, 6, 8, 9];
357 let rx = StreamReader::new(
358 &mut store,
359 DelayedStreamProducer {
360 inner: BufferStreamProducer {
361 buffer: expected.clone(),
362 },
363 maybe_yield: yield_times(10).boxed(),
364 },
365 )?;
366 store
367 .run_concurrent(async move |accessor| {
368 let (rx, expected) = readiness_guest
369 .local_local_readiness()
370 .call_start(accessor, rx, expected)
371 .await?;
372
373 accessor.with(|access| {
374 rx.pipe(
375 access,
376 DelayedStreamConsumer {
377 inner: BufferStreamConsumer { expected },
378 maybe_yield: yield_times(10).boxed(),
379 },
380 )
381 })?;
382
383 Ok(())
384 })
385 .await?
386 }
387
388 #[tokio::test]
async_poll_synchronous() -> Result<()>389 pub async fn async_poll_synchronous() -> Result<()> {
390 test_run(&[test_programs_artifacts::ASYNC_POLL_SYNCHRONOUS_COMPONENT]).await
391 }
392
393 #[tokio::test]
async_poll_stackless() -> Result<()>394 pub async fn async_poll_stackless() -> Result<()> {
395 test_run(&[test_programs_artifacts::ASYNC_POLL_STACKLESS_COMPONENT]).await
396 }
397
398 mod cancel {
399 wasmtime::component::bindgen!({
400 path: "wit",
401 world: "cancel-host",
402 exports: { default: async | store },
403 });
404 }
405
406 // No-op function; we only test this by composing it in `async_cancel_caller`
407 #[allow(
408 dead_code,
409 reason = "here only to make the `assert_test_exists` macro happy"
410 )]
async_cancel_callee()411 pub fn async_cancel_callee() {}
412
413 #[tokio::test]
async_cancel_caller() -> Result<()>414 pub async fn async_cancel_caller() -> Result<()> {
415 test_cancel(Mode::Normal).await
416 }
417
418 #[tokio::test]
async_cancel_caller_leak_task_after_cancel() -> Result<()>419 pub async fn async_cancel_caller_leak_task_after_cancel() -> Result<()> {
420 test_cancel(Mode::LeakTaskAfterCancel).await
421 }
422
423 #[tokio::test]
async_trap_cancel_guest_after_start_cancelled() -> Result<()>424 pub async fn async_trap_cancel_guest_after_start_cancelled() -> Result<()> {
425 test_cancel_trap(Mode::TrapCancelGuestAfterStartCancelled).await
426 }
427
428 #[tokio::test]
async_trap_cancel_guest_after_return_cancelled() -> Result<()>429 pub async fn async_trap_cancel_guest_after_return_cancelled() -> Result<()> {
430 test_cancel_trap(Mode::TrapCancelGuestAfterReturnCancelled).await
431 }
432
433 #[tokio::test]
async_trap_cancel_guest_after_return() -> Result<()>434 pub async fn async_trap_cancel_guest_after_return() -> Result<()> {
435 test_cancel_trap(Mode::TrapCancelGuestAfterReturn).await
436 }
437
438 #[tokio::test]
async_trap_cancel_host_after_return_cancelled() -> Result<()>439 pub async fn async_trap_cancel_host_after_return_cancelled() -> Result<()> {
440 test_cancel_trap(Mode::TrapCancelHostAfterReturnCancelled).await
441 }
442
443 #[tokio::test]
async_trap_cancel_host_after_return() -> Result<()>444 pub async fn async_trap_cancel_host_after_return() -> Result<()> {
445 test_cancel_trap(Mode::TrapCancelHostAfterReturn).await
446 }
447
test_cancel_trap(mode: Mode) -> Result<()>448 async fn test_cancel_trap(mode: Mode) -> Result<()> {
449 let message = "`subtask.cancel` called after terminal status delivered";
450 let trap = test_cancel(mode).await.unwrap_err();
451 assert!(
452 format!("{trap:?}").contains(message),
453 "expected `{message}`; got `{trap:?}`",
454 );
455 Ok(())
456 }
457
test_cancel(mode: Mode) -> Result<()>458 async fn test_cancel(mode: Mode) -> Result<()> {
459 let engine = Engine::new(&config())?;
460
461 let component = make_component(
462 &engine,
463 &[
464 test_programs_artifacts::ASYNC_CANCEL_CALLER_COMPONENT,
465 test_programs_artifacts::ASYNC_CANCEL_CALLEE_COMPONENT,
466 ],
467 )
468 .await?;
469
470 let mut linker = Linker::new(&engine);
471
472 wasmtime_wasi::p2::add_to_linker_async(&mut linker)?;
473 yield_::local::local::yield_::add_to_linker::<_, Ctx>(&mut linker, |ctx| ctx)?;
474
475 let mut store = Store::new(
476 &engine,
477 Ctx {
478 wasi: WasiCtxBuilder::new().inherit_stdio().build(),
479 table: ResourceTable::default(),
480 continue_: false,
481 },
482 );
483
484 let cancel_host =
485 cancel::CancelHost::instantiate_async(&mut store, &component, &linker).await?;
486 store
487 .run_concurrent(async move |accessor| {
488 cancel_host
489 .local_local_cancel()
490 .call_run(accessor, mode, 100)
491 .await?;
492 Ok::<_, wasmtime::Error>(())
493 })
494 .await??;
495
496 Ok(())
497 }
498
499 #[tokio::test]
async_intertask_communication() -> Result<()>500 pub async fn async_intertask_communication() -> Result<()> {
501 test_run_with_count(
502 &[test_programs_artifacts::ASYNC_INTERTASK_COMMUNICATION_COMPONENT],
503 2,
504 )
505 .await
506 }
507
508 #[tokio::test]
async_transmit_caller() -> Result<()>509 pub async fn async_transmit_caller() -> Result<()> {
510 test_run(&[
511 test_programs_artifacts::ASYNC_TRANSMIT_CALLER_COMPONENT,
512 test_programs_artifacts::ASYNC_TRANSMIT_CALLEE_COMPONENT,
513 ])
514 .await
515 }
516
517 #[tokio::test]
async_transmit_callee() -> Result<()>518 pub async fn async_transmit_callee() -> Result<()> {
519 test_transmit(test_programs_artifacts::ASYNC_TRANSMIT_CALLEE_COMPONENT).await
520 }
521
522 pub trait TransmitTest {
523 type Instance: Send + Sync;
524 type Params;
525 type Result: Send + Sync + 'static;
526
instantiate( store: impl AsContextMut<Data = Ctx>, component: &Component, linker: &Linker<Ctx>, ) -> impl Future<Output = Result<Self::Instance>>527 fn instantiate(
528 store: impl AsContextMut<Data = Ctx>,
529 component: &Component,
530 linker: &Linker<Ctx>,
531 ) -> impl Future<Output = Result<Self::Instance>>;
532
call<'a>( accessor: &'a Accessor<Ctx, HasSelf<Ctx>>, instance: &'a Self::Instance, params: Self::Params, ) -> impl Future<Output = Result<Self::Result>> + Send + 'a533 fn call<'a>(
534 accessor: &'a Accessor<Ctx, HasSelf<Ctx>>,
535 instance: &'a Self::Instance,
536 params: Self::Params,
537 ) -> impl Future<Output = Result<Self::Result>> + Send + 'a;
538
into_params( store: impl AsContextMut<Data = Ctx>, control: StreamReader<Control>, caller_stream: StreamReader<String>, caller_future1: FutureReader<String>, caller_future2: FutureReader<String>, ) -> Self::Params539 fn into_params(
540 store: impl AsContextMut<Data = Ctx>,
541 control: StreamReader<Control>,
542 caller_stream: StreamReader<String>,
543 caller_future1: FutureReader<String>,
544 caller_future2: FutureReader<String>,
545 ) -> Self::Params;
546
from_result( store: impl AsContextMut<Data = Ctx>, result: Self::Result, ) -> Result<( StreamReader<String>, FutureReader<String>, FutureReader<String>, )>547 fn from_result(
548 store: impl AsContextMut<Data = Ctx>,
549 result: Self::Result,
550 ) -> Result<(
551 StreamReader<String>,
552 FutureReader<String>,
553 FutureReader<String>,
554 )>;
555 }
556
557 struct StaticTransmitTest;
558
559 impl TransmitTest for StaticTransmitTest {
560 type Instance = transmit::bindings::TransmitCallee;
561 type Params = (
562 StreamReader<Control>,
563 StreamReader<String>,
564 FutureReader<String>,
565 FutureReader<String>,
566 );
567 type Result = (
568 StreamReader<String>,
569 FutureReader<String>,
570 FutureReader<String>,
571 );
572
instantiate( store: impl AsContextMut<Data = Ctx>, component: &Component, linker: &Linker<Ctx>, ) -> Result<Self::Instance>573 async fn instantiate(
574 store: impl AsContextMut<Data = Ctx>,
575 component: &Component,
576 linker: &Linker<Ctx>,
577 ) -> Result<Self::Instance> {
578 transmit::bindings::TransmitCallee::instantiate_async(store, &component, &linker).await
579 }
580
call<'a>( accessor: &'a Accessor<Ctx, HasSelf<Ctx>>, instance: &'a Self::Instance, params: Self::Params, ) -> impl Future<Output = Result<Self::Result>> + Send + 'a581 fn call<'a>(
582 accessor: &'a Accessor<Ctx, HasSelf<Ctx>>,
583 instance: &'a Self::Instance,
584 params: Self::Params,
585 ) -> impl Future<Output = Result<Self::Result>> + Send + 'a {
586 instance
587 .local_local_transmit()
588 .call_exchange(accessor, params.0, params.1, params.2, params.3)
589 }
590
into_params( _store: impl AsContextMut<Data = Ctx>, control: StreamReader<Control>, caller_stream: StreamReader<String>, caller_future1: FutureReader<String>, caller_future2: FutureReader<String>, ) -> Self::Params591 fn into_params(
592 _store: impl AsContextMut<Data = Ctx>,
593 control: StreamReader<Control>,
594 caller_stream: StreamReader<String>,
595 caller_future1: FutureReader<String>,
596 caller_future2: FutureReader<String>,
597 ) -> Self::Params {
598 (control, caller_stream, caller_future1, caller_future2)
599 }
600
from_result( _: impl AsContextMut<Data = Ctx>, result: Self::Result, ) -> Result<( StreamReader<String>, FutureReader<String>, FutureReader<String>, )>601 fn from_result(
602 _: impl AsContextMut<Data = Ctx>,
603 result: Self::Result,
604 ) -> Result<(
605 StreamReader<String>,
606 FutureReader<String>,
607 FutureReader<String>,
608 )> {
609 Ok(result)
610 }
611 }
612
613 struct DynamicTransmitTest;
614
615 impl TransmitTest for DynamicTransmitTest {
616 type Instance = Instance;
617 type Params = Vec<Val>;
618 type Result = Val;
619
instantiate( store: impl AsContextMut<Data = Ctx>, component: &Component, linker: &Linker<Ctx>, ) -> Result<Self::Instance>620 async fn instantiate(
621 store: impl AsContextMut<Data = Ctx>,
622 component: &Component,
623 linker: &Linker<Ctx>,
624 ) -> Result<Self::Instance> {
625 linker.instantiate_async(store, component).await
626 }
627
call<'a>( accessor: &'a Accessor<Ctx, HasSelf<Ctx>>, instance: &'a Self::Instance, params: Self::Params, ) -> Result<Self::Result>628 async fn call<'a>(
629 accessor: &'a Accessor<Ctx, HasSelf<Ctx>>,
630 instance: &'a Self::Instance,
631 params: Self::Params,
632 ) -> Result<Self::Result> {
633 let exchange_function = accessor.with(|mut store| {
634 let transmit_instance = instance
635 .get_export_index(store.as_context_mut(), None, "local:local/transmit")
636 .ok_or_else(|| format_err!("can't find `local:local/transmit` in instance"))?;
637 let exchange_function = instance
638 .get_export_index(store.as_context_mut(), Some(&transmit_instance), "exchange")
639 .ok_or_else(|| format_err!("can't find `exchange` in instance"))?;
640 instance
641 .get_func(store.as_context_mut(), exchange_function)
642 .ok_or_else(|| format_err!("can't find `exchange` in instance"))
643 })?;
644
645 let mut results = vec![Val::Bool(false)];
646 exchange_function
647 .call_concurrent(accessor, ¶ms, &mut results)
648 .await?;
649 Ok(results.pop().unwrap())
650 }
651
into_params( mut store: impl AsContextMut<Data = Ctx>, control: StreamReader<Control>, caller_stream: StreamReader<String>, caller_future1: FutureReader<String>, caller_future2: FutureReader<String>, ) -> Self::Params652 fn into_params(
653 mut store: impl AsContextMut<Data = Ctx>,
654 control: StreamReader<Control>,
655 caller_stream: StreamReader<String>,
656 caller_future1: FutureReader<String>,
657 caller_future2: FutureReader<String>,
658 ) -> Self::Params {
659 vec![
660 control.try_into_stream_any(&mut store).unwrap().into(),
661 caller_stream
662 .try_into_stream_any(&mut store)
663 .unwrap()
664 .into(),
665 caller_future1
666 .try_into_future_any(&mut store)
667 .unwrap()
668 .into(),
669 caller_future2
670 .try_into_future_any(&mut store)
671 .unwrap()
672 .into(),
673 ]
674 }
675
from_result( _store: impl AsContextMut<Data = Ctx>, result: Self::Result, ) -> Result<( StreamReader<String>, FutureReader<String>, FutureReader<String>, )>676 fn from_result(
677 _store: impl AsContextMut<Data = Ctx>,
678 result: Self::Result,
679 ) -> Result<(
680 StreamReader<String>,
681 FutureReader<String>,
682 FutureReader<String>,
683 )> {
684 let Val::Tuple(fields) = result else {
685 unreachable!()
686 };
687 let mut fields = fields.into_iter();
688 let Val::Stream(stream) = fields.next().unwrap() else {
689 unreachable!()
690 };
691 let Val::Future(future1) = fields.next().unwrap() else {
692 unreachable!()
693 };
694 let Val::Future(future2) = fields.next().unwrap() else {
695 unreachable!()
696 };
697 let stream = StreamReader::try_from_stream_any(stream).unwrap();
698 let future1 = FutureReader::try_from_future_any(future1).unwrap();
699 let future2 = FutureReader::try_from_future_any(future2).unwrap();
700 Ok((stream, future1, future2))
701 }
702 }
703
test_transmit(component: &str) -> Result<()>704 async fn test_transmit(component: &str) -> Result<()> {
705 test_transmit_with::<StaticTransmitTest>(component).await?;
706 test_transmit_with::<DynamicTransmitTest>(component).await
707 }
708
test_transmit_with<Test: TransmitTest + 'static>(component: &str) -> Result<()>709 async fn test_transmit_with<Test: TransmitTest + 'static>(component: &str) -> Result<()> {
710 let engine = Engine::new(&config())?;
711
712 let component = make_component(&engine, &[component]).await?;
713
714 let mut linker = Linker::new(&engine);
715
716 wasmtime_wasi::p2::add_to_linker_async(&mut linker)?;
717
718 let mut store = Store::new(
719 &engine,
720 Ctx {
721 wasi: WasiCtxBuilder::new().inherit_stdio().build(),
722 table: ResourceTable::default(),
723 continue_: false,
724 },
725 );
726
727 let test = Test::instantiate(&mut store, &component, &linker).await?;
728
729 enum Event<Test: TransmitTest> {
730 Result(Test::Result),
731 ControlWriteA(mpsc::Sender<Control>),
732 ControlWriteB(mpsc::Sender<Control>),
733 ControlWriteC(mpsc::Sender<Control>),
734 ControlWriteD,
735 WriteA,
736 ReadC(mpsc::Receiver<String>, Option<String>),
737 ReadD(mpsc::Receiver<String>, Option<String>),
738 ReadNone(Option<String>),
739 }
740
741 let (mut control_tx, control_rx) = mpsc::channel(1);
742 let control_rx = StreamReader::new(&mut store, PipeProducer::new(control_rx))?;
743 let (mut caller_stream_tx, caller_stream_rx) = mpsc::channel(1);
744 let caller_stream_rx = StreamReader::new(&mut store, PipeProducer::new(caller_stream_rx))?;
745 let (caller_future1_tx, caller_future1_rx) = oneshot::channel();
746 let caller_future1_rx = FutureReader::new(&mut store, OneshotProducer::new(caller_future1_rx))?;
747 let (_, caller_future2_rx) = oneshot::channel();
748 let caller_future2_rx = FutureReader::new(&mut store, OneshotProducer::new(caller_future2_rx))?;
749 let (callee_future1_tx, callee_future1_rx) = oneshot::channel();
750 let (callee_stream_tx, callee_stream_rx) = mpsc::channel(1);
751 store
752 .run_concurrent(async |accessor| {
753 let mut caller_future1_tx = Some(caller_future1_tx);
754 let mut callee_future1_tx = Some(callee_future1_tx);
755 let mut callee_future1_rx = Some(callee_future1_rx);
756 let mut callee_stream_tx = Some(callee_stream_tx);
757 let mut callee_stream_rx = Some(callee_stream_rx);
758 let mut complete = false;
759 let mut futures = FuturesUnordered::<
760 Pin<Box<dyn Future<Output = Result<Event<Test>>> + Send>>,
761 >::new();
762
763 futures.push(
764 async move {
765 control_tx.send(Control::ReadStream("a".into())).await?;
766 Ok(Event::ControlWriteA(control_tx))
767 }
768 .boxed(),
769 );
770
771 futures.push(
772 async move {
773 caller_stream_tx.send(String::from("a")).await?;
774 Ok(Event::WriteA)
775 }
776 .boxed(),
777 );
778
779 let params = accessor.with(|s| {
780 Test::into_params(
781 s,
782 control_rx,
783 caller_stream_rx,
784 caller_future1_rx,
785 caller_future2_rx,
786 )
787 });
788
789 futures.push(
790 Test::call(accessor, &test, params)
791 .map(|v| v.map(Event::Result))
792 .boxed(),
793 );
794
795 while let Some(event) = futures.try_next().await? {
796 match event {
797 Event::Result(result) => {
798 accessor.with(|mut store| {
799 let (callee_stream_rx, callee_future1_rx, _) =
800 Test::from_result(&mut store, result)?;
801 callee_stream_rx.pipe(
802 &mut store,
803 PipeConsumer::new(callee_stream_tx.take().unwrap()),
804 )?;
805 callee_future1_rx.pipe(
806 &mut store,
807 OneshotConsumer::new(callee_future1_tx.take().unwrap()),
808 )?;
809 wasmtime::error::Ok(())
810 })?;
811 }
812 Event::ControlWriteA(mut control_tx) => {
813 futures.push(
814 async move {
815 control_tx.send(Control::ReadFuture("b".into())).await?;
816 Ok(Event::ControlWriteB(control_tx))
817 }
818 .boxed(),
819 );
820 }
821 Event::WriteA => {
822 _ = caller_future1_tx.take().unwrap().send("b".into());
823 let mut callee_stream_rx = callee_stream_rx.take().unwrap();
824 futures.push(
825 async move {
826 let value = callee_stream_rx.next().await;
827 Ok(Event::ReadC(callee_stream_rx, value))
828 }
829 .boxed(),
830 );
831 }
832 Event::ControlWriteB(mut control_tx) => {
833 futures.push(
834 async move {
835 control_tx.send(Control::WriteStream("c".into())).await?;
836 Ok(Event::ControlWriteC(control_tx))
837 }
838 .boxed(),
839 );
840 }
841 Event::ControlWriteC(mut control_tx) => {
842 futures.push(
843 async move {
844 control_tx.send(Control::WriteFuture("d".into())).await?;
845 Ok(Event::ControlWriteD)
846 }
847 .boxed(),
848 );
849 }
850 Event::ReadC(callee_stream_rx, mut value) => {
851 assert_eq!(value.take().as_deref(), Some("c"));
852 futures.push(
853 callee_future1_rx
854 .take()
855 .unwrap()
856 .map(|v| Event::ReadD(callee_stream_rx, v.ok()))
857 .map(Ok)
858 .boxed(),
859 );
860 }
861 Event::ControlWriteD => {}
862 Event::ReadD(_, None) => unreachable!(),
863 Event::ReadD(mut callee_stream_rx, Some(value)) => {
864 assert_eq!(&value, "d");
865 futures.push(
866 async move { Ok(Event::ReadNone(callee_stream_rx.next().await)) }
867 .boxed(),
868 );
869 }
870 Event::ReadNone(Some(_)) => unreachable!(),
871 Event::ReadNone(None) => {
872 complete = true;
873 }
874 }
875 }
876
877 assert!(complete);
878
879 wasmtime::error::Ok(())
880 })
881 .await??;
882 Ok(())
883 }
884
885 mod synchronous_transmit {
886 wasmtime::component::bindgen!({
887 path: "wit",
888 world: "synchronous-transmit-guest",
889 });
890 }
891
892 #[tokio::test]
async_cancel_transmit() -> Result<()>893 pub async fn async_cancel_transmit() -> Result<()> {
894 test_synchronous_transmit(
895 test_programs_artifacts::ASYNC_CANCEL_TRANSMIT_COMPONENT,
896 true,
897 )
898 .await
899 }
900
901 #[tokio::test]
async_synchronous_transmit() -> Result<()>902 pub async fn async_synchronous_transmit() -> Result<()> {
903 test_synchronous_transmit(
904 test_programs_artifacts::ASYNC_SYNCHRONOUS_TRANSMIT_COMPONENT,
905 false,
906 )
907 .await
908 }
909
test_synchronous_transmit(component: &str, procrastinate: bool) -> Result<()>910 async fn test_synchronous_transmit(component: &str, procrastinate: bool) -> Result<()> {
911 let engine = Engine::new(&config())?;
912
913 let component = make_component(&engine, &[component]).await?;
914
915 let mut linker = Linker::new(&engine);
916
917 wasmtime_wasi::p2::add_to_linker_async(&mut linker)?;
918
919 let mut store = Store::new(
920 &engine,
921 Ctx {
922 wasi: WasiCtxBuilder::new().inherit_stdio().build(),
923 table: ResourceTable::default(),
924 continue_: false,
925 },
926 );
927
928 let instance = linker.instantiate_async(&mut store, &component).await?;
929 let guest = synchronous_transmit::SynchronousTransmitGuest::new(&mut store, &instance)?;
930 let stream_expected = vec![2u8, 4, 6, 8, 9];
931 let producer = DelayedStreamProducer {
932 inner: BufferStreamProducer {
933 buffer: stream_expected.clone(),
934 },
935 maybe_yield: yield_times(10).boxed(),
936 };
937 let stream = if procrastinate {
938 StreamReader::new(&mut store, ProcrastinatingStreamProducer(producer))?
939 } else {
940 StreamReader::new(&mut store, producer)?
941 };
942 let future_expected = 10;
943 let producer = DelayedFutureProducer {
944 inner: ValueFutureProducer {
945 value: future_expected,
946 },
947 maybe_yield: yield_times(10).boxed(),
948 };
949 let future = if procrastinate {
950 FutureReader::new(&mut store, ProcrastinatingFutureProducer(producer))?
951 } else {
952 FutureReader::new(&mut store, producer)?
953 };
954 store
955 .run_concurrent(async move |accessor| {
956 let (stream, stream_expected, future, future_expected) = guest
957 .local_local_synchronous_transmit()
958 .call_start(accessor, stream, stream_expected, future, future_expected)
959 .await?;
960
961 accessor.with(|mut access| -> wasmtime::Result<_> {
962 let consumer = DelayedStreamConsumer {
963 inner: BufferStreamConsumer {
964 expected: stream_expected,
965 },
966 maybe_yield: yield_times(10).boxed(),
967 };
968 if procrastinate {
969 stream.pipe(&mut access, ProcrastinatingStreamConsumer(consumer))?;
970 } else {
971 stream.pipe(&mut access, consumer)?;
972 }
973 let consumer = DelayedFutureConsumer {
974 inner: ValueFutureConsumer {
975 expected: future_expected,
976 },
977 maybe_yield: yield_times(10).boxed(),
978 };
979 if procrastinate {
980 future.pipe(access, ProcrastinatingFutureConsumer(consumer))?;
981 } else {
982 future.pipe(access, consumer)?;
983 }
984 Ok(())
985 })?;
986
987 Ok(())
988 })
989 .await?
990 }
991