WASI Preview 2: rewrite streams and pollable implementation

[pchickey]

When we landed the WASI Preview 2 support in wasmtime-wasi we knew the streams/poll/scheduler implementation was going to get torn out and rewritten: see task 4 in #6370. This is that rewrite. Most of this was pair-programmed with @elliottt.

Features

wasmtime-wasi Preview 2 support is now built on top of tokio, and the host implementations for streams and pollable always use async Rust. It now passes all of the poll_oneoff tests, including on windows.

The new HostInputStream, HostOutputStream, and HostPollable traits are used to implement the wasi:io/streams.{input-stream,output-stream} and wasi:poll/poll.pollable resources.

Host{Input, Output}Stream design

input-stream and output-stream resources have a significant number of methods. For the host traits, we managed to boil this down to two required methods, as well as some optional methods to cover vectored IO, splice, and write-zeroes. The default implementations provide an unoptimized but correct implementation - the user may override them if an optimization is desired.

HostInputStream has just two methods: fn read(&mut self, ...) and async fn ready(&mut self). The most important invariant is that read must be non-blocking, and ready must block until the stream is ready for reading.

Nonblocking IO is actually impossible on regular files, but we hide that detail from WASI programs

There's a huge lie in the description above: internally, wasmtime_wasi uses InternalHost{Input,Output}Stream, which are enums that contain either a Host{Input,Output}Stream, or a private File{Input,Output}Stream. These file streams always perform a blocking read/write, whether or not the WASI program has called the nonblocking read or the blocking blocking_read method. And, any pollable associated with file read or write readiness always returns ready immediately.

This is required because there is no way to implement streams on files with epoll(7) (or on macos or windows using mio, afaik) based readiness - poll(2) will lie and always tell you files are ready for reading of writing, but the actual syscall may end up blocking the thread due to disk IO. And, rather than stick the work on a background task, we must wait until the OS read(2) or write(2) for a file has completed before returning from wasi stream (nonblocking) read or write, because the userlands in WASI expect (reasonably!) that any errors in regular file IO are reported synchronously.

So, we put an escape hatch into wasmtime-wasi which is just for regular file IO. It awaits for a tokio spawn_blocking to complete in both the blocking and non-blocking implementations of the stream reads and writes. Like all the other syscalls in the filesystem implementation, we need the async rust to put any blocking syscalls in a spawn_blocking, in order to not block the executor.

We spent a the better part of a week trying to square this circle and this is the best we could come up to make files work like we expect, but also expose the right interfaces from wasmtime-wasi so that all other streams have the proper non-blocking semantics.

HostPollable design

The trick to HostPollable is to create a Rust Future indicating whether a stream (or some other resource, e.g. http bodies) is ready without requiring both the stream and pollable hold some Arc<Mutex<...>> to the resource itself. @alexcrichton pointed out the right design here: HostPollable is an enum with two variants:

• TableEntry { index: u32, make_future: for<'a> fn(&'a mut dyn Any) -> Pin<Box<dyn Future = Result<()>> + Send + 'a>} means, get a mutable reference to some (still dynamically-typed) other entry in the table, and apply this fn to create a Future from that entry, which has the same lifetime as the mutable reference. For the streams, we use this variant, and make_future is just a call to the Host{Input,Output}Stream::ready method.
• Closure( Box<Fn() -> Pin<Box<dyn Future = Result<()>> + Send + 'static>> + Send + Sync + 'static> ) means, use this owned closure to create a Future with a static lifetime. This is used for creating timers from the monotonic and wall clocks, which are ambient and therefore do not have table entries.

Note that we discovered a major design problem with the TableEntry variant that exposes a crash vector if the parent resource of a pollable is destroyed before the pollable itself. This is a big enough problem that we are setting it aside to be solved in a follow-up PR.

Implementations

wasmtime-wasi provides a struct AsyncReadStream which wraps a tokio::io::AsyncRead to provide an impl of HostInputStream, AsyncWriteStream to wraps AsyncWrite to provide HostOutputStream.

AsyncReadStream and AsyncWriteStream will manage their own buffer in order for the HostInputStream::read and HostOutputStream::write methods to be non-blocking. Each spawns a helper tokio::task to perform async operations in the background. This requires some buffering.

Note, we discussed ensuring these background tasks are reaped properly - right now we believe that the implementations are correct and they will exit on their own when the foreground side gets dropped. However, in the future to implement e.g. splice and forward correctly, we may need to allow these tasks to (optionally!) live longer than the WASI program's execution.

Additionally, we have provided an implementation that correctly handles waiting on stdin readiness on both windows and unix. In Unix, this is built on top of tokio::os::unix::AsyncFd, allows tokio to register stdin with epoll(7). Because stdin is a process-wide singleton, and also because epoll will return an error if the same fd is registered multiple times, the stdin resource is a global singleton, created lazily and guarded by a mutex. wasmtime_wasi::preview2::stdio::stdin() returns a struct Stdin which takes a lock on that singleton in the HostInputStream::{read, ready} methods. This means that, if for some reason you have granted the process stdin to multiple wasi streams (in the same, or different, stores) you can get all sorts of strange behavior because the same global resource backs them all behind the scenes - but this is always the case, the implementation in wasmtime_wasi just allows you to do so without hanging or panicking.

On Windows, the implementation of stdin is broken, and inheriting stdin will panic. We have a path to fixing it, but we are leaving that out of this PR, because it has been open for far too long, and far too much work is predicated on landing it.

Sync and Async uses

Despite using async for its implementation, wasmtime-wasi will still work for synchronous Wasmtime embeddings. A new test-programs suite wasi-preview2-components-sync is identical to wasi-preview2-components in all regards except that it uses wasmtime's synchronous APIs, and the appropriate wasmtime_wasi::command::sync::add_to_linker, rather than wasmtime_wasi::command::add_to_linker which requires an async Config. The synchronous implementation creates its own Tokio runtime (if one does not already exist) and then invokes the async implementation inside of a block_on.

The wasi-common Preview 1 poll_oneoff implementation allowed for a synchronous implementation to work without any sort of dependency on tokio, by abstracting it behind a pluggable WasiSched. This PR eliminates that option for our Preview 2 users. If there is a compelling need, and labor available to implement and maintain it, we could collaborate on re-introducing a fully synchronous tokio-free option in the future. For this implementation, I decided that we just need a tokio backend in order to build a path towards supporting both wasi-sockets and wasi-http with the same modular stream and pollable host traits.

Bindings traits

After living with the bindings generated underneath crate::preview2::wasi and crate::preview2::wasi::command for a while, I decided that the interface traits belong under crate::preview2::bindings and command belongs under crate::preview2::command.

Beyond that code motion, we now generate async traits for the wasi:io/streams, wasi:poll/poll, and wasi:filesystem/filesystem interfaces, and sync traits for everything else. If you want to use the bindings from an async embedding (i.e. with an async Config), the crate::preview2::bindings::{interface}::add_to_linker will do it, and if you want to use them from a synchronous embedding, use crate::preview2::bindings::sync::{interface}::add_to_linker.

[alexcrichton]

Nice 👍

For the WIT changes, those are from the upstream proposals? Or local changes made temporarily for this PR?

[alexcrichton]

Instead of an Option here perhaps add Closed to WriteState?

[pchickey]

we had it factored that way, then we were taking advantage of Option::take() so changing the repr made sense, but now I think I can switch it back... Edit nope, we still use the take() to perform the move of the error case out of the state. it still could be changed with appropriate helper funcs but i dont think its worth it

[alexcrichton]

Had a good chuckle reading this :)

Regardless I agree we can't make this truly async and sync close here is the right thing to do IMO

[pchickey]

The wit has not been submitted upstream yet. I was going to work on that once this landed, as well as the rest of the upstream/wasmtime wit reconciliation that is pending

[pchickey]

I just made a big update to the PR description with the design changes since the last round of review, and the known problems in this PR. I am now going to make the windows inherit stdin panic, mark the tests that affects as cfg_attr(windows, should_panic), and land this PR. The follow up work includes, in rough priority order:

• HostPollable::TableEntry issues around pollables outliving their parent resource, which are a crash vector right now
• get windows stdin working, and better tests of unix asyncfd based stdin
• implementing stream splice and forward, and dealing with the lifetimes of background tasks accordingly