Possible Wasmtime Improvement for Async Wasm Component execution

[benbrandt]

Hi all 👋🏻

I was having a look at your Wasmtime usage in the extension_host crate and think there is a possible improvement. However, I will freely admit I haven't fully understood how you handle async executors or if this is an issue so feel free to point out what I may have missed here 😄 I wanted to check before submitting a PR as it might have unforeseen effects if I missed anything.

Background

Since you have turned on async_support on the Wasmtime Engine, the Wasmtime docs mention the potential problem this introduces, because technically the wasm execution happens within the Future::poll method:

The Future::poll method is the main driving force behind Rust’s futures. That method’s own documentation states “an implementation of poll should strive to return quickly, and should not block”. This, however, can be at odds with executing WebAssembly code as part of the poll method itself. If your WebAssembly is untrusted then this could allow the poll method to take arbitrarily long in the worst case, likely blocking all other asynchronous tasks.

So, there is a potential that extension code with a bug (or being malicious) could waste lots of CPU cycles and block an async thread.

Looking through your codebase, you seem to have a pattern where you use a select to either wait to receive a message or timeout, which should avoid issues in the main thread, which is great. So this would just be an optimization in terms of extensions sharing resources amongst themselves, and also to make sure a current poll still doesn't block up one of the threads that the extension is running on.

Proposal

This assumes I have correctly diagnosed the current situation, which I may not have. But assuming this isn't just a theoretical concern, the docs mention there are two options, Config::epoch_interruption and Config::consume_fuel. My hunch is epoch interruption would be good enough to help solve the theoretical issue, to make sure that the extensions yield at some point back to the exector at regular intervals.

To integrate this would require a few changes:

1. Turn on Epoch Interruption

This part is fairly straightforward, updating the creation of the Engine:

fn wasm_engine() -> wasmtime::Engine {
    static WASM_ENGINE: OnceLock<wasmtime::Engine> = OnceLock::new();

    WASM_ENGINE
        .get_or_init(|| {
            let mut config = wasmtime::Config::new();
            config.wasm_component_model(true);
            config.async_support(true);
            config.epoch_interruption(true);
            wasmtime::Engine::new(&config).unwrap()
        })
        .clone()
}

2. Increment the Epoch

This requires something outside of the engine in the embedder. One option is to do this in a thread that wakes up at specific intervals and increments the epoch:

fn wasm_engine() -> wasmtime::Engine {
    static WASM_ENGINE: OnceLock<wasmtime::Engine> = OnceLock::new();

    WASM_ENGINE
        .get_or_init(|| {
            let mut config = wasmtime::Config::new();
            config.wasm_component_model(true);
            config.async_support(true);
            config.epoch_interruption(true);
            
            let engine = wasmtime::Engine::new(&config).unwrap();
            
            let engine_ref = engine.weak();
    
            std::thread::spawn(move || {
                loop {
                    std::thread::sleep(EPOCH_INTERVAL);
                    let Some(engine) = engine_ref.upgrade() else {
                        break;
                    };
                    engine.increment_epoch();
                }
            });
            
            engine
        })
        .clone()
}

There are some points worth highlighting here. According to the docs for increment_epoch, because the Engine is an Arc, we want an EngineWeak in the thread so that it doesn't keep the engine alive and leak, and it will terminate once the Engine is actually dropped.

I submitted a PR around this for Fermyon Spin if you want to see the discussion there on this pattern.

TODO: Decide on how long EPOCH_INTERVAL is. Basically, how often should an extension check to see if they need to yield.

3. Configure the Store to yield

Once this is set, then we have to also configure the store for how often it should yield.

            let mut store = wasmtime::Store::new(
                &this.engine,
                WasmState {
                    ctx: this.build_wasi_ctx(&manifest).await?,
                    manifest: manifest.clone(),
                    table: ResourceTable::new(),
                    host: this.clone(),
                },
            );
            store.set_epoch_deadline(1);
            store.epoch_deadline_async_yield_and_update(1);

This will yield after every a deadline of number of ticks (configured to some interval) and once it has yielded will be given a new deadline of some amount of ticks in the future. These values could also be adjusted to some other interval.

If you want other behavior, like an upper-bound on total execution time, you can also use epoch_deadline_callback which allows you to return a result, with an Ok of either continuing to process for a number of ticks, yielding with a new deadline, or Err and stop execution.

This is where there is a lot of flexibility, and likely needs testing and benchmarking.

Next Steps

1. I'd like to get feedback on if I missed something in your architecture where this is accounted for in another way, and all of this isn't necessary, or needs to be handled differently. I'm also not sure what your current plans for handling extensions are in terms of making sure they don't hog resources or maximum execution time, etc. So any feedback in this area would be helpful. If this doesn't seem like something you want, that is fine too 😄

2. I'd be happy to push a PR or pair on this with some of you if that is desirable.

3. If I do, I'd be very happy to get insights on what sort of tests/benchmarks you'd want to make sure this does indeed make extension behavior better, and not worse. I love Zed's performance and would hate to introduce something that takes away from that!

Thanks, and let me know if and how I can help with this!