How to create longer-lived objects in Wasm that would outlive Wasm function scope, but not completely static ? In JS heap

[artem-tim]

Hello,
Here's my problem. I'm using a Wasm function to process HTTP Response body in a web worker, which is a ReadableStream.
For that, I mirror the JS way to do so : create a Transformer, define a transform function Closure, and pipe the stream. Something like that :

pub async fn my_wasm_response_processor(resp: Response) -> Response {
  let body: ReadableStream = response.body().unwrap();
  let mut transformer = Transformer::new();
  // My custom struct which has a method .process(chunk) which takes in a Uint8Array and returns another Uint8Array.
  // It is a custom function which defines how each response's body needs to be processed based on response's headers and so on
  let my_processor: MyProcessor = MyProcessor::new(resp.headers());
  let transform_closure = Closure::new(move |chunk: Uint8Array, controller: TransformStreamDefaultController| {
    let new_chunk: Uint8Array = my_processor.process(chunk);
    controller.enqueue_with_chunk(&new_chunk).unwrap();
  });
  let transform_fn: &Function = transform_closure.as_ref().unchecked_ref::<Function>();
  let ts: TransformStream = TransformStream::new_with_transformer(&transformer).unwrap();
  #[allow(unused_must_use)] {
      // pipe_to is async, but we don't want to await anything here
      body.pipe_to(&ts.writable());
  }

  let mut resp_init = ResponseInit::new();
  resp_init.headers(&resp.headers());
  Response::new_with_opt_readable_stream_and_init(Some(&ts.readable()), &resp_init).unwrap()
}

And the above code compiles, but generates the following error at runtime :

closure invoked after being dropped

I've read the docs and I understand that there is manual deallocation of JS's Function as soon as Rust's Closure goes out of scope (which is at the end of the function), however I don't want that ! The whole point is to return the response (send it back to the user), and keep transforming the stream.

Now, I've actually managed to "kind of" solve it by using thread_local and putting the Closure inside a statically lived RefCell (along with the associated MyProcessor), having the transformer borrow it every time it is needed, however this is a no good solution in the real world - thread_local static variables are shared between all requests, where I need a separate Closure & MyProcessor for each request. I don't want them to live forever, just enough to finish stream processing of the current request, but without conflicting with other requests.

There must be a way to make these objects live in the JS heap, and have JS's GC take care of them once no longer needed ?

[Liamolucko]

The easy solution to this is to use Closure::forget (or Closure::into_js_value) and set WASM_BINDGEN_WEAKREF=1 when running wasm-bindgen, which does exactly what you said: making the closure live in the JS heap and get deallocated by the GC. If you don't set WASM_BINDGEN_WEAKREF though, Closure::forget just leaks memory, since doing this requires the 'weak references' proposal (specifically FinalizationRegistry). It's been supported by all browsers for a while now, but wasm-bindgen still has it disabled by default because apparently some serverless runtimes don't (see #3384).

Alternatively, you can manually deallocate the Closure inside flush like this:

pub async fn my_wasm_response_processor(resp: Response) -> Response {
    let body: ReadableStream = resp.body().unwrap();
    let mut transformer = Transformer::new();
    // My custom struct which has a method .process(chunk) which takes in a Uint8Array and returns another Uint8Array.
    // It is a custom function which defines how each response's body needs to be processed based on response's headers and so on
    let my_processor: MyProcessor = MyProcessor::new(resp.headers());
    let transform_closure = Closure::<dyn Fn(_, _)>::new(
        move |chunk: Uint8Array, controller: TransformStreamDefaultController| {
            let new_chunk: Uint8Array = my_processor.process(chunk);
            controller.enqueue_with_chunk(&new_chunk).unwrap();
        },
    );
    transformer.transform(transform_closure.as_ref().unchecked_ref::<Function>());
    // NEW STUFF
    let closure_handles = Rc::new(RefCell::new(None));
    let flush_closure = Closure::<dyn Fn()>::new({
        let closure_handles = Rc::clone(&closure_handles);
        move || {
            // The stream is finished and so we can deallocate both closures.
            *closure_handles.borrow_mut() = None;
        }
    });
    transformer.flush(flush_closure.as_ref().unchecked_ref::<Function>());
    *closure_handles.borrow_mut() = Some((transform_closure, flush_closure));
    // END NEW STUFF
    let ts: TransformStream = TransformStream::new_with_transformer(&transformer).unwrap();
    #[allow(unused_must_use)]
    {
        // pipe_to is async, but we don't want to await anything here
        body.pipe_to(&ts.writable());
    }

    let mut resp_init = ResponseInit::new();
    resp_init.headers(&resp.headers());
    Response::new_with_opt_readable_stream_and_init(Some(&ts.readable()), &resp_init).unwrap()
}

This is basically just a tweaked version of the requestAnimationFrame example.

[artem-tim]

Thank you for pointing out Closure::forget / into_js_value which as you said do exactly what I need - somehow I've missed them. It would be the perfect solution if the workers platform supported FinalizationRegistry, unfortunately at least Cloudflare's web workers don't support it, so thank you also for the alternative solution.

My only concern is that in case of a network (or any other) error with the stream, flush() might not be called (thus leaking memory), so the same manual deallocation (*closure_handles.borrow_mut() = None;) needs to also be added to stream error handlers.