How to stop a wasm function?

help
1

So the question I have is rather simple, but it has lead to me to some rabbit holes and I'm getting to darker places slowly but surely so I'm trying to get a fresh point of view from other rust users.

I have a command line tool, performing some iterative optimization algorithm that can take a bit of time (from few seconds to a minute or so). In the terminal, it's fairly easy to stop it via [Ctrl+C] if needed. But I have ported the algorithm to WebAssembly and I'm struggling to find a solution to stop the function from user input in the web app.

What I've got so far

The wasm module is loaded in a web worker to prevent the main UI thread to freeze. I can communicate from wasm to the web app via imported JS functions, that's what enabled my implementation of Log to pass the program logs to the main thread as they occur. Below is a sreenshot of the logs while the algorithm is running and a "stop!" button that I'd like to be able to click to stop the algorithm.

lowrr-logs
lowrr-logs1087×573 58.1 KB

I thought I could also inject a function as a parameter in my algorithm, with something looking like:

fn algorithm(should_stop: &mut dyn FnMut() -> bool) {
    loop {
        if should_stop() { return; }
        ...
    }
}

And that should_stop function would be imported from JS and would check a global value stop_order in the worker that could be changed anytime by a user interaction in the main thread. But that's where the rabbit hole starts.

Since JS is monothreaded, I can't change that value until the algorithm finishes running, which defies its purpose. UNLESS algorithm and should_stop are actually async. In such case, every call to should_stop would give control back to the JS loop, which will execute the received worker message telling it to change that stop_order global value and the next call to should_stop would actually terminate the algorithm.

But I think doing that (if possible, which seems it is) means I bring the whole async ecosystem to my crate, just because I need it in the WebAssembly port. Currently I have a workspace with <mycrate>-lib, <mycrate>-bin (a CLI app) and <mycrate>-wasm to use in the web app. The algorithm function lives in <mycrate>-lib and I'd like to avoid bringing the whole async ecosystem to it.

So I thought I could just gate this behind a feature, but then again I'm hitting an unknown. The algoritm function is roughly 200 lines and can't really get much shorter. I would hate having to duplicate this function to have another one (with some if should_stop() { .. } sprinkled) under a cfg_attr because it means the two will eventually get out of sync (pun intended ^^). Is there a simple way to do that in Rust without duplicating the function? I'm afraid not because the other function also must have a different signature starting with async fn ... right? (I've actually never used async in Rust yet)

Is there a simpler way?

I'm getting quite far from my original goal of just being able to stop that wasm function? Is there a simpler way?

PS: here is the app in question in case interested: https://lowrr.pizenberg.fr.

1 Like
2

There are only two options you have. Either you periodically check if you should stop with the function or you use a web worker and use the .terminate() method to stop the web worker. For periodic checking you don't have to duplicate the function. You can also use #[cfg(target_arch = "wasm32")] { if should_stop () { ... } }. Note that #[cfg] isn't allowed directly on if, but it is allowed on blocks.

1 Like
3

That's good to know!

That won't be usefull since in our case it means we loose the images data loaded into the wasm memory and it's roughly equivalent to just reload the page. So I guess periodic check it is!

How about the fact that one function is fn algorithm() while the other is async fn algorithm(should_stop ...). There needs to be two different functions right? Or is it something that can be solved with macro somehow? I've 0 experience with macros.

4

I think you could use a macro like

macro_rules! my_algorithm {
    ($arg1:expr, $arg2:expr, $maybe_stop:block) => {
        // part of the algorithm
        $maybe_stop
        // some other part
    }
}

And then invoked like my_algorithm!(the_first_arg, the_second_arg, {}) on native and like my_algorithm!(the_first_arg, the_second_arg, { if { should_stop() { return; } }) on wasm.

1 Like
5

I'll try to validate first the wasm-bindgen async approach with duplicated code and once that ok I'll try the macro as you suggest and let you know how things go. Thanks!

6

I'd like to avoid bringing the whole async ecosystem to it.

Having written async webassembly, the only real pain is potentially having to do a global refactor to make everything async. If you expose only async rust functions to javascript you can simply await those on the javascript side. There is no need to choose an async executor.

Since JS is monothreaded, I can't change that value until the algorithm finishes running

Try a SharedArrayBuffer here. These can be shared between main and (multiple) webworkers.

People have written hacks (using SharedArrayBuffer) to write libraries enabling multithreaded wasm. Perhaps one of these is useful for you.

7

Yeah I'd totally have gone that way if I didn't want to support Safari :slight_smile:

I'll try that! I haven't got the time yet but will try later today.

1 Like
8

setTimeout is the classic way to yield to the event loop.

1 Like
9

If you don't want to pull async into your WebAssembly, one option is to turn the loop inside out.

So instead of this:

fn algorithm(should_stop: &mut dyn FnMut() -> bool) {
    loop {
        if should_stop() { return; }
        ...
    }
}

You would break algorithm() into small chunks and write something like this:

fn poll_algorithm(state: &mut AlgorithmState) {
  // do a little bit of computation
}

Then as part of your web worker you might use setTiimeout() to poll the algorithm and re-schedule it to be polled again asap.

import { AlgorithmState, poll_algorithm } from "./wasm";

let keepGoing = true;
const state = new AlgorithmState();

const action = () => {
  if (keepGoing) {
    setTimeout(action, 0);
    poll_algorithm(state);
  }
};

setTimeout(action, 0);

If you squint a little you'll see this is effectively like writing your own Future by hand.

In the normal desktop version you'd then create a wrapper function that calls your poll_algorithm() function to completion.

2 Likes
10

Thanks for the insight, it's interesting to see it that way. Unfortunately breaking the algorithm into small pieces really isn't practical because of it's complexity. But the good news is I've managed to do the async change :slight_smile: It was far from trivial though since I am not used to all the things you have to know when dealing with async functions and wasm bindgen such as this or that.

The approach I took to make it work was to have a newtype wrapper for an Rc<RefCell<...>>. It looks like this:

// src/lib.rs

use wasm_bindgen::prelude::*;
// some other imports

#[wasm_bindgen(raw_module = "../worker.mjs")]
extern "C" {
    #[wasm_bindgen(js_name = "shouldStop")]
    async fn should_stop() -> JsValue; // bool
}

async fn should_stop_bool() -> bool {
    let js_bool = should_stop().await;
    js_bool.as_bool().unwrap()
}

// This wrapper trick is because we cannot have async functions referencing &self.
// https://github.com/rustwasm/wasm-bindgen/issues/1858
#[wasm_bindgen]
pub struct Handler(Rc<RefCell<InnerHandler>>);

#[wasm_bindgen]
impl Handler {
    pub fn init() -> Self {
        Handler(Rc::new(RefCell::new(InnerHandler::init())))
    }
    pub fn immutable_call(&self) -> Result<JsValue, JsValue> {
        self.0.borrow().immutable_call()
    }
    pub fn mutable_call(&mut self) -> Result<(), JsValue> {
        let inner = Rc::clone(&self.0);
        let result = (*inner).borrow_mut().mutable_call(); // split in two lines needed
        result
    }
    pub fn mutable_async_call(&mut self) -> js_sys::Promise {
        let inner = Rc::clone(&self.0);
        wasm_bindgen_futures::future_to_promise(async_rc(inner))
    }
}

// Intermediate function between the async method in InnerHandler
// and the exported method that returns a promise in Handler.
async fn async_rc(mutself: Rc<RefCell<InnerHandler>>) -> Result<JsValue, JsValue> {
    let mut inner = (*mutself).borrow_mut();
    let result = inner.mutable_async_call();
    result.await
}

// Struct that initially was Handler but is now wrapped in an Rc<RefCell<>>
// the Handler new type.
struct InnerHandler { ... }

impl InnerHandler {
    fn init() -> Self { ... } // unchanged
    fn immutable_call(&self) -> Result<JsValue, JsValue> { ... } // unchanged
    fn mutable_call(&mut self) -> Result<(), JsValue> { ... } // unchanged

    // The function turned async and calling should_stop() inside our algorithm
    async fn mutable_async_call(&mut self) -> Result<JsValue, JsValue> {
        ...
        let result = algorithm_async(..., should_stop_bool).await;
        ...
    }
}

With in the JS code the following bits:

// worker.mjs

// Function regularly called in the algorithm to check if it should stop.
export async function shouldStop() {
  await sleep(0); // Force to give control back.
  return stopOrder; // Global variable in the worker being set via worker messages
}

// Small utility function.
function sleep(ms) {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

For the time being, to check that it works (and it does!) I have duplicated algorithm and algorithm_async functions like so:

pub fn algorithm() -> ... { ... }
pub async fn algorithm_async<FB: Future<Output = bool>>(
    should_stop: fn() -> FB
) -> ... { ... }

Now the only thing left is to dedup those two functions with some macro magic :slight_smile:

1 Like
11

Awesome!

I'd try to avoid this if possible.

If the algorithm has any amount of size or complexity you'll take a massive hit in readability and maintainability because you've got the normal logic intermingled with "meta-logic" (i.e. abstracting over async-ness). You see C codebases make this mistake with #if all the time, typically to hack a particular feature in, and it makes the resulting code quite hard to follow.

Instead, what about using futures::executor::block_on() to create a synchronous function which wraps your async version and polls it to completion? If you pass should_stop() in using dependency injection and that is the only async function it calls (i.e. you aren't reading from a TcpStream which requires the tokio runtime) then you could probably get away with it.

pub async fn algorithm_async(should_stop: &dyn Fn() -> Future<Item=bool>) { ... }

pub fn algorithm() {
  let should_stop = || async { false };

  futures::executor::block_on(move || algorithm_async(&should_stop));
}

I'd also recommend reading What Color is your Function? which explains the incompatibilities between sync and async functions at a higher level.

2 Likes
12

That is indeed the case. I'll see how block_on() plays with my setup, thanks for that link!

Looks like an interesting read! I like the blue-red analogy starting point. Will try to read it fully when I find time :slight_smile:

1 Like
13

Turns out the macro was very clean:

// Exact same content than previous algo_async
// except for the $(...)* expansion
// that enables calling the macro without should_stop for algo_sync
macro_rules! algo_macro {
    ($arg1: expr, $arg2: expr, $($should_stop: expr),*) => {{
        ...
        $(if $should_stop().await { return Err(...); })*
        ...
    }};
}

pub fn algo_sync(arg1: ..., arg2: ...) -> ... {
    algo_macro!(arg1, arg2,)
}

pub async fn algo_async(arg1: ..., arg2: ..., should_stop: ...) -> ... {
    algo_macro!(arg1, arg2, should_stop)
}
1 Like
14

Oh nice!

I was honestly expecting it to be a mess. Especially if you are doing recursion or have multiple sub-functions which need to switch between sync or async.

15

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.