How to call async exported function properly with tokio? · general

Stream: general

Topic: How to call async exported function properly with tokio?

raindust (Jun 24 2022 at 07:44):

I wrote a simple example to export an async function like the following:

#[no_mangle]
pub async extern "C" fn answer() -> i32 {
    println!("call inner answer");
    42
}

After use cargo wasi build to get the wasm target, I create an example with tokio:

[package]
name = "lib-async"
version = "0.1.0"
edition = "2021"

[dependencies]
anyhow = "1.0"
tokio = { version = "1.8.0", features = ["full"] }
wasmtime = "0.38"
wasmtime-wasi = { version = "0.38", features = ["tokio"] }

and the source code is:

use wasmtime::{Config, Engine, Linker, Module, Store};
use wasmtime_wasi::{tokio::WasiCtxBuilder, WasiCtx};

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let mut config = Config::new();
    config.async_support(true);
    config.consume_fuel(true);

    let engine = Engine::new(&config)?;
    let module = Module::from_file(&engine, "target/wasm32-wasi/debug/lib_async_example.wasm")?;

    let mut linker = Linker::new(&engine);
    wasmtime_wasi::tokio::add_to_linker(&mut linker, |s| s)?;

    let join = tokio::task::spawn(async move { get_answer(&mut linker, &engine, &module).await });

    let result = join.await??;
    println!("answer is {}", result);
    Ok(())
}

async fn get_answer(
    linker: &mut Linker<WasiCtx>,
    engine: &Engine,
    module: &Module,
) -> anyhow::Result<i32> {
    let wasi = WasiCtxBuilder::new().inherit_stdout().build();
    let mut store = Store::new(engine, wasi);
    store.out_of_fuel_async_yield(u64::MAX, 10000);

    let instance = linker.instantiate_async(&mut store, module).await?;

    let answer = instance.get_typed_func::<(), i32, _>(&mut store, "answer")?;
    let result = answer.call_async(&mut store, ()).await?;

    Ok(result)
}

I expect the print answer is 42, but the actual answer is 0.

bjorn3 (Jun 24 2022 at 09:06):

You can't directly call an async function from the host. I am surprised rustc even allows extern "C" on async functions. Async functions return a future you need to poll from an executor. You can't poll them from the host. Wasmtime's call_async requires a sync wasm function. The difference with non-async call I believe is that functions from the host called by the wasm module can be async. If such a function is called and it isn't ready yet, wasmtime will skip past the wasm functions entirely and return a not-ready result from call_async, if you poll it again, it will immediately poll the host async function that was originally called. Only when this function is ready will it return back to the wasm caller. In other words wasmtime async support is transparent to the wasm module. If you want async functions on the wasm module side you will need to write an executor for this or use an existing one like tokio. https://github.com/tokio-rs/tokio/pull/4716 will add wasi support to tokio.

Wasm32-wasi target support by rjzak · Pull Request #4716 · tokio-rs/tokio

Motivation Getting Tokio working on WebAssembly-WASI so that web applications can be executed by a Wasi runtime, such as Wasmtime and others. Solution Working on supporting the wasm32-wasi target. ...

raindust (Jun 24 2022 at 10:51):

@bjorn3 thanks for reply.
I thought wasm (or wasi) already support async functions because wasm-bindgen can wrap a rust Future into JS Promise, and we compile the extern "C" on async functions with wasm32-wasi, is it means can translate a rust Future to the wasm grammar but wasmtime still can't support it yet?

https://github.com/tokio-rs/tokio/pull/4716

The issue above I think it is used to support tokio when cargo build with wasm32-wasi. But I think what we really need is let wasmtime can recognize (or give use the choice to identify) there is an async function, so we can call it in the async way.

raindust (Jun 24 2022 at 10:55):

And if I add dependencies in https://github.com/bytecodealliance/wasmtime/blob/main/examples/tokio/wasm/Cargo.toml like the following:

[dependencies]
futures = "0.3"

and modify https://github.com/bytecodealliance/wasmtime/blob/main/examples/tokio/wasm/tokio-wasi.rs like this:

use futures::executor::block_on;

async fn hello_world() {
    println!("hello world from async!");
}

fn main() {
    let name = std::env::var("NAME").unwrap();
    println!("Hello, world! My name is {}", name);
    std::thread::sleep(std::time::Duration::from_secs(1));
    println!("Goodbye from {}", name);

    let future = hello_world(); // Nothing is printed
    block_on(future); // `future` is run and "hello, world!" is printed
}

It works properly.

wasmtime/Cargo.toml at main · bytecodealliance/wasmtime

Standalone JIT-style runtime for WebAssembly, using Cranelift - wasmtime/Cargo.toml at main · bytecodealliance/wasmtime

wasmtime/tokio-wasi.rs at main · bytecodealliance/wasmtime

Standalone JIT-style runtime for WebAssembly, using Cranelift - wasmtime/tokio-wasi.rs at main · bytecodealliance/wasmtime

bjorn3 (Jun 24 2022 at 11:00):

Wasm itself doesn't have the concept of async functions and as such wasmtime can't recognize them. In rust async functions are lowered to a function that returns a value implementing the Future trait. This trait has a method called poll that needs to be called to drive the future to completion. This method uses the "Rust" calling convention and as such isn't callable outside of the rust code from which you compile the wasm module. The wasm_bindgen_futures crate has an async executor that can handle javascript Promise's. There is no equivalent to this for wasi yet. Note that futures::block_on only works in this case as the future is always immediately ready. If it isn't, futures::block_on will attempt to park the current thread until the waker is called. I think libstd implements thread parking as a panic on wasm/wasi as they are single threaded platforms at the moment and thus parking the current thread doesn't make sense as nobody could unpark it again. In addition neither wasm nor wasi has an api to park a thread in the first place.

raindust (Jun 24 2022 at 13:32):

Ah, I forgot wasm/wasi are single threaded platforms. So if I wanna get rid of the "callback hell" in a wasi library, write a single export function like the following still works, am I Right?

#[no_mangle]
pub extern "C" fn message_loop(op_len: i32, req_len: i32) -> i32 {
    task::block_on(future::poll_fn(move |cx: &mut Context<'_>| {
        loop {
                    // message handling related logic
                }
         }
}

bjorn3 (Jun 24 2022 at 14:08):

What exactly are you using async for?

raindust (Jun 24 2022 at 15:05):

bjorn3 said:

What exactly are you using async for?

I have a handler map in a wasi library used to store callbacks like this:

lazy_static! {
    pub static ref MAP_HANDLER: Mutex<HashMap<u64, Vec<CallbackItem>>> = Mutex::new(HashMap::new());
}

pub type Callback = dyn FnOnce(TeaResult<Vec<u8>>) -> TeaResult<()> + Sync + Send + 'static;
pub struct CallbackItem {
    pub callback: Box<Callback>,
    pub timeout: Duration,
    pub start_at: SystemTime,
}

pub fn add_callback<F>(seq_number: u64, callback: F) -> TeaResult<()>
where
    F: FnOnce(TeaResult<Vec<u8>>) -> TeaResult<()> + Sync + Send + 'static,
{
    let mut map = MAP_HANDLER.lock().unwrap();
    let callback_item = CallbackItem::with_default_duration(Box::new(callback))?;
    match map.get_mut(&seq_number) {
        Some(callback_list) => callback_list.push(callback_item),
        None => {
            map.insert(seq_number, vec![callback_item]);
        }
    }
    Ok(())
}

And when the callback reply reached, I use the following function to handle it:

pub fn result_handler(msg: &[u8], seq_number: u64) -> TeaResult<()> {
    let callback = {
        match MAP_HANDLER.lock() {
            Ok(mut hash_map) => match hash_map.get_mut(&seq_number) {
                Some(callback_list) => {
                    let callback = callback_list.pop();
                    if callback_list.is_empty() {
                        hash_map.remove(&seq_number);
                    }
                    callback
                }
                None => None,
            },
            Err(e) => {
                error!("Result handler lock failed, details: {:?}", e);
                None
            }
        }
    };

    match deserialize::<TeaResult<Vec<u8>>, _>(&msg) {
        Ok(content) => {
            match callback {
                Some(callback) => (callback.callback)(content),
                None => {
                    error!("Cannot find callback function (seq_number: {}) from hashmap. Cannot callback", seq_number);
                    Ok(())
                }
            }
        }
        Err(e) => {
            warn!(
                "failed to deserialize reply message (seq_number: {}): {:?}",
                seq_number, e
            );
            Ok(())
        }
    }
}

Using the handler map I wrote above, I can do something asynchronously using callback like this:

pub fn call_async<F>(id: &str, operation: &str, callback: F) -> TeaResult<()>
where
    F: FnOnce(TeaResult<Vec<u8>>) -> TeaResult<()> + Sync + Send + 'static,
{
    let seq_number = get_seq_number()?;
    add_callback(seq_number, callback)?;

    // do something synchronously

    Ok(())
}

As you can see, using callback is hard to read and debug, so I want replace these callback functions with async.

bjorn3 (Jun 24 2022 at 15:26):

I see. I did recommend writing your own executor. https://rust-lang.github.io/async-book/02_execution/04_executor.html has an explanation on how to do this. https://os.phil-opp.com/async-await/ is also a really nice explanation of how async in rust internally works. It is written in the context of implementing an OS in rust, but a lot is applicable outside of this context too.

raindust (Jun 25 2022 at 01:51):

Thanks a lot, I will read these tutorials carefully and write my own executor

Notification Bot (Jun 25 2022 at 01:52):

raindust has marked this topic as resolved.

Notification Bot (Jun 28 2022 at 12:46):

raindust has marked this topic as unresolved.