Host blocks when trying to write/read stream to/from guest

[segeljakt]

Hi, I am trying out the new features in wasip3-prototyping (https://github.com/segeljakt/wasmtime-component-async-test).

Everything works well, except one of the test cases. I have a function on the guest that is supposed to take an input stream, create a task that reads 10 elements from the stream and write a response to each element on an output stream:

pub mod bindings {
    use wit_bindgen::generate;
    generate!({
        world: "guest",
        path: "interface.wit",
        async: {
            exports: [ "pkg:component/intf#test4" ]
        }
    });
    pub struct Component;
    export!(Component);
}

impl Guest for bindings::Component {
    async fn test4(mut stream: StreamReader<String>) -> StreamReader<String> {
        let (mut tx, rx) = bindings::wit_stream::new::<String>();
        async_support::spawn(async move {
            for i in 0..10 {
                match stream.next().await {
                    Some(Ok(_items)) => {
                        tx.send(vec!["Response".to_string()]).await;
                    }
                    _ => {
                        tx.close().await.unwrap();
                        break;
                    }
                }
            }
        });
        rx
    }
}

I have this code on the host to test the function:

#[tokio::test]
async fn test4() {
    let (instance, mut store, intf_export) = init().await;
    let export = instance
        .get_export(&mut store, Some(&intf_export), "test4")
        .unwrap();

    let func3: TypedFunc<(HostStream<String>,), (HostStream<String>,)> =
        instance.get_typed_func(&mut store, export).unwrap();

    let (mut tx, rx) = instance.stream(&mut store).unwrap();

    let (result,) = func3.call_async(&mut store, (rx.into(),)).await.unwrap();

    let handle1 = tokio::task::spawn(async move {
        for i in 0..10 {
            println!("{i} Writing: Hello World! (test4)");
            let Some(new) = tx
                .write(Single(format!("Hello World! {i} (test4)")))
                .into_future()
                .await
            else {
                panic!("Error writing stream");
            };
            tx = new;
        }
    });

    func3.post_return_async(&mut store).await.unwrap();
    let mut result: StreamReader<Vec<String>> = result.into_reader(&mut store);

    let handle2 = tokio::task::spawn(async move {
        for i in 0..10 {
            println!("{i} Reading...");
            let Ok((new, item)) = result.read().into_future().await else {
                panic!("Error reading stream");
            };

            result = new;

            println!("Result: {:?}", item);
        }
    });
    tokio::try_join!(handle1, handle2).unwrap();
}

However, when I try to run the code, the spawned tasks on the host get stuck at tx.write(Single(...)).into_future().await and result.read().into_future().await. Is there any obvious mistake in my code? Everything seems to work well when just using futures.

Another slightly unrelated question. Is there an async API for WASI that guests can use to for example read files and sockets? I only found wasi::io::streams, but it appears to be synchronous.

[dicej]

Promise::into_future is tricky (and I should add more documentation to clarify where it should and shouldn't be used). The reason Promise exists is that operations like future/stream reads and writes (as well as guest tasks started using e.g. TypedFunc::call_concurrent) can only make progress when the event loop for the component is running.

Promise::get takes care of that automatically by running the event loop until the Future inside the Promise finishes, but that only makes sense to call from the top level of the host embedding. In a host function, we use Promise::into_future instead, since the Future returned by a concurrent host function is already being polled via the event loop, so no need to recursively run that same event loop. However, if you move that future into a task spawned using tokio::task::spawn, that's no longer true -- the Tokio executor polls that task directly without involving the Wasmtime event loop, so the operation never makes progress.

If you need to spawn a background task that uses streams or futures from a host function, you must use Accessor::spawn rather than tokio::task::spawn. That will spawn a task into the event loop, allowing the futures to make progress. There are examples of this in the wasmtime-wasi p3 test suite and in the component-async-tests suite, e.g.

wasip3-prototyping/crates/misc/component-async-tests/src/resource_stream.rs

Lines 37 to 71 in 796f75c

	impl bindings::local::local::resource_stream::Host for &mut Ctx {
	async fn foo<T: 'static>(
	accessor: &mut Accessor<T, Self>,
	count: u32,
	) -> wasmtime::Result<HostStream<Resource<ResourceStreamX>>> {
	struct Task {
	tx: StreamWriter<Single<Resource<ResourceStreamX>>>,
	count: u32,
	}
	impl<T, U: wasmtime_wasi::IoView> AccessorTask<T, U, Result<()>> for Task {
	async fn run(self, accessor: &mut Accessor<T, U>) -> Result<()> {
	let mut tx = Some(self.tx);
	for _ in 0..self.count {
	tx = accessor
	.with(|mut view| {
	let item = IoView::table(&mut *view).push(ResourceStreamX)?;
	Ok::<_, anyhow::Error>(
	tx.take().unwrap().write(Single(item)).into_future(),
	)
	})?
	.await;
	}
	Ok(())
	}
	}
	let (tx, rx) = accessor.with(|mut view| {
	let instance = view.instance();
	instance.stream(&mut view)
	})?;
	accessor.spawn(Task { tx, count });
	Ok(rx.into())
	}
	}

Hope that helps; let me know if it's still unclear.

[segeljakt]

Thank you for the help 👍 I will try it out!

[segeljakt]

@dicej Is there any way to get access to an Accessor in the host when you set up the wasmtime engine and call the guest's functions? In the example, it seems like the host gets access to the Accessor by implementing the bindings of the wit file and becoming a component itself. Is this something I should do in addition to setting up the engine?

[dicej]

Ah, apologies; I should have read your code more carefully. I see now that your spawn calls are at the top level or the embedding rather than within a host function as in the example I pointed to.

And yes, there is a way to do that, but it's currently marked #[doc(hidden)] since I had originally meant for it to only be used from wasmtime-wit-bindgen-generated code:

wasip3-prototyping/crates/wasmtime/src/runtime/component/instance.rs

Lines 325 to 340 in 796f75c

	#[cfg(feature = "component-model-async")]
	#[doc(hidden)]
	pub fn spawn(
	&self,
	mut store: impl AsContextMut,
	task: impl std::future::Future<Output = Result<()>> + Send + 'static,
	) {
	let instance = unsafe {
	&mut *store.as_context_mut().0[self.0]
	.as_ref()
	.unwrap()
	.instance_ptr()
	};
	instance.spawn(task)
	}
	}

. Also, it doesn't yet offer a way to join the spawned tasks.

Another way to do this kind of thing with the current API is by using PromisesUnordered plus an enum-based state machine similar to

wasip3-prototyping/crates/misc/component-async-tests/tests/scenario/transmit.rs

Lines 194 to 342 in 796f75c

	async fn test_transmit_with<Test: TransmitTest + 'static>(component: &[u8]) -> Result<()> {
	let mut config = Config::new();
	config.debug_info(true);
	config.cranelift_debug_verifier(true);
	config.wasm_component_model(true);
	config.wasm_component_model_async(true);
	config.async_support(true);
	let engine = Engine::new(&config)?;
	let make_store = || {
	Store::new(
	&engine,
	Ctx {
	wasi: WasiCtxBuilder::new().inherit_stdio().build(),
	table: ResourceTable::default(),
	continue_: false,
	wakers: Arc::new(Mutex::new(None)),
	},
	)
	};
	let component = Component::new(&engine, component)?;
	let mut linker = Linker::new(&engine);
	wasmtime_wasi::add_to_linker_async(&mut linker)?;
	let mut store = make_store();
	let (test, instance) = Test::instantiate(&mut store, &component, &linker).await?;
	#[allow(clippy::type_complexity)]
	enum Event<Test: TransmitTest> {
	Result(Test::Result),
	ControlWriteA(Option<StreamWriter<Single<Control>>>),
	ControlWriteB(Option<StreamWriter<Single<Control>>>),
	ControlWriteC(Option<StreamWriter<Single<Control>>>),
	ControlWriteD,
	WriteA,
	WriteB(bool),
	ReadC(Result<(StreamReader<Vec<String>>, Vec<String>), Option<ErrorContext>>),
	ReadD(Result<String, Option<ErrorContext>>),
	ReadNone(Result<(StreamReader<Vec<String>>, Vec<String>), Option<ErrorContext>>),
	}
	let (control_tx, control_rx) = instance.stream(&mut store)?;
	let (caller_stream_tx, caller_stream_rx) = instance.stream(&mut store)?;
	let (caller_future1_tx, caller_future1_rx) = instance.future(&mut store)?;
	let (_caller_future2_tx, caller_future2_rx) = instance.future(&mut store)?;
	let mut promises = PromisesUnordered::<Event<Test>>::new();
	let mut caller_future1_tx = Some(caller_future1_tx);
	let mut callee_stream_rx = None;
	let mut callee_future1_rx = None;
	let mut complete = false;
	promises.push(
	control_tx
	.write(Single(Control::ReadStream("a".into())))
	.map(Event::ControlWriteA),
	);
	promises.push(
	caller_stream_tx
	.write(Single("a".into()))
	.map(|_| Event::WriteA),
	);
	promises.push(
	Test::call(
	&mut store,
	&test,
	Test::into_params(
	control_rx.into(),
	caller_stream_rx.into(),
	caller_future1_rx.into(),
	caller_future2_rx.into(),
	),
	)
	.await?
	.map(Event::Result),
	);
	while let Some(event) = promises.next(&mut store).await? {
	match event {
	Event::Result(result) => {
	let results = Test::from_result(&mut store, instance, result)?;
	callee_stream_rx = Some(results.0.into_reader(&mut store));
	callee_future1_rx = Some(results.1.into_reader(&mut store));
	}
	Event::ControlWriteA(tx) => {
	promises.push(
	tx.unwrap()
	.write(Single(Control::ReadFuture("b".into())))
	.map(Event::ControlWriteB),
	);
	}
	Event::WriteA => {
	promises.push(
	caller_future1_tx
	.take()
	.unwrap()
	.write("b".into())
	.map(Event::WriteB),
	);
	}
	Event::ControlWriteB(tx) => {
	promises.push(
	tx.unwrap()
	.write(Single(Control::WriteStream("c".into())))
	.map(Event::ControlWriteC),
	);
	}
	Event::WriteB(delivered) => {
	assert!(delivered);
	promises.push(callee_stream_rx.take().unwrap().read().map(Event::ReadC));
	}
	Event::ControlWriteC(tx) => {
	promises.push(
	tx.unwrap()
	.write(Single(Control::WriteFuture("d".into())))
	.map(|_| Event::ControlWriteD),
	);
	}
	Event::ReadC(Err(_)) => unreachable!(),
	Event::ReadC(Ok((rx, values))) => {
	assert_eq!("c", &values[0]);
	promises.push(callee_future1_rx.take().unwrap().read().map(Event::ReadD));
	callee_stream_rx = Some(rx);
	}
	Event::ControlWriteD => {}
	Event::ReadD(Err(_)) => unreachable!(),
	Event::ReadD(Ok(value)) => {
	assert_eq!(value, "d");
	promises.push(callee_stream_rx.take().unwrap().read().map(Event::ReadNone));
	}
	Event::ReadNone(Ok(_)) => unreachable!(),
	Event::ReadNone(Err(e)) => {
	assert!(e.is_none());
	complete = true;
	}
	}
	}
	assert!(complete);
	Ok(())
	}

. I'll admit that's less convenient for your scenario than spawning a couple of tasks and joining them, though.

I'll see about removing the #[doc(hidden)] attribute from Instance::spawn and adding a way to join or cancel tasks.

[dicej]

BTW, if you're eager to play with Instance::spawn, you could try updating your code to use it now. Since there's no way to join the tasks yet, you could instead create a oneshot channel for each task, move the Senders into the task futures, and then call send on them at the end of each future. Then you can join the Receivers in your top-level code.

[dicej]

Hmm, actually I don't think what I recommended above will quite work; you'll end up hanging forever when joining the Receivers since there's currently no way to run the event loop except via a Promise, whereas you'll only have the Receivers. Sounds like I'll need to add new APIs for this to work at all.

[segeljakt]

Ah ok, I can try out the PromisesUnordered approach for now.

[segeljakt]

I tried it out and it works well 👍

I am wondering, when using WASI inside the guest, it seems like operations are synchronous. For example, wasi::io::streams::InputStream::read() is not async. Will wasi support async operations, or should I implement these by streaming from the host via StreamReader?

[dicej]

wasi-io will be deprecated in WASIp3, i.e. there will be no 0.3.x versions of it. Instead, all the WASI interfaces which previously used wasi-io have been updated to use the built-in stream and future types and other Component Model async features (e.g. wasi-filesystem and wasi-http).

So if you want to use StreamReader with WASI, you'll need to use the new p3 support in wasmtime-wasi (and soon wasmtime-wasi-http once #58 is merged). See the *_0_3_*.rs files in https://github.com/bytecodealliance/wasip3-prototyping/tree/main/crates/test-programs/src/bin for examples.

Note that there isn't any support for a wasm32-wasip3 target in wasi-libc, LLVM, Rust, etc. yet, so you need to use wit-bindgen to generate bindings and use those directly for the time being.

[segeljakt]

Thanks for the help, I managed to get it working 👍 It will be exciting to see wasi-http merged.