experiment: faster wasm streaming in Rust

ext/fetch/26_fetch.js

@@ -78,6 +78,8 @@ function opFetchSend(rid) {
  */
 function createResponseBodyStream(responseBodyRid, terminator) {
   const readable = readableStreamForRid(responseBodyRid);
+  // internal, used by wasm streaming
+  readable.rid = responseBodyRid;
 
   function onAbort() {
     errorReadableStream(readable, terminator.reason);

ext/fetch/lib.rs

@@ -1,6 +1,7 @@
 // Copyright 2018-2024 the Deno authors. All rights reserved. MIT license.
 
 mod fs_fetch_handler;
+mod wasm_streaming;
 
 use std::borrow::Cow;
 use std::cell::RefCell;
@@ -69,6 +70,7 @@ pub use data_url;
 pub use reqwest;
 
 pub use fs_fetch_handler::FsFetchHandler;
+pub use wasm_streaming::handle_wasm_streaming;
 
 #[derive(Clone)]
 pub struct Options {

ext/fetch/wasm_streaming.rs

@@ -0,0 +1,187 @@
+// Copyright 2018-2023 the Deno authors. All rights reserved. MIT license.
+
+use deno_core::error::type_error;
+use deno_core::error::AnyError;
+use deno_core::v8;
+use deno_core::OpState;
+use deno_core::ResourceId;
+use std::cell::RefCell;
+use std::rc::Rc;
+
+/// The Wasm streaming compilation pipeline.
+pub fn handle_wasm_streaming(
+  state: Rc<RefCell<OpState>>,
+  scope: &mut v8::HandleScope,
+  value: v8::Local<v8::Value>,
+  mut wasm_streaming: v8::WasmStreaming,
+) {
+  let (url, rid) = match compile_response(scope, value) {
+    Ok(Some((url, rid))) => (url, rid),
+    Ok(None) => {
+      // 2.7
+      wasm_streaming.finish();
+      return;
+    }
+    Err(e) => {
+      // 2.8
+      let err = v8::String::new(scope, &e.to_string()).unwrap();
+      wasm_streaming.abort(Some(err.into()));
+      return;
+    }
+  };
+
+  wasm_streaming.set_url(&url);
+
+  deno_core::unsync::spawn(async move {
+    loop {
+      let resource = state.borrow().resource_table.get_any(rid);
+      let resource = match resource {
+        Ok(r) => r,
+        Err(_) => {
+          state.borrow().borrow::<deno_core::V8TaskSpawner>().spawn(
+            move |scope| {
+              wasm_streaming.abort(Some(
+                v8::String::new(scope, "Failed to get resource.")
+                  .unwrap()
+                  .into(),
+              ))
+            },
+          );
+          return;
+        }
+      };
+
+      let view = deno_core::BufMutView::new(65536);
+      let (bytes, view) = match resource.read_byob(view).await {
+        Ok(bytes) => bytes,
+        Err(e) => {
+          state.borrow().borrow::<deno_core::V8TaskSpawner>().spawn(
+            move |scope| {
+              wasm_streaming.abort(Some(
+                v8::String::new(
+                  scope,
+                  &format!("Error reading wasm resource: {}", e),
+                )
+                .unwrap()
+                .into(),
+              ))
+            },
+          );
+
+          return;
+        }
+      };
+      /* EOF */
+      if bytes == 0 {
+        break;
+      }
+
+      wasm_streaming.on_bytes_received(&view[..bytes]);
+    }
+
+    /* Spawn a task on JS loop to finish the streaming compilation */
+    state
+      .borrow()
+      .borrow::<deno_core::V8TaskSpawner>()
+      .spawn(move |_| wasm_streaming.finish());
+  });
+}
+
+// Partially implements https://webassembly.github.io/spec/web-api/#compile-a-potential-webassembly-response
+pub fn compile_response(
+  scope: &mut v8::HandleScope,
+  value: v8::Local<v8::Value>,
+) -> Result<Option<(String, ResourceId)>, AnyError> {
+  let object = value
+    .to_object(scope)
+    .ok_or_else(|| type_error("Response is not an object."))?;
+  let url = get_string(scope, object, "url")?;
+
+  // 2.3.
+  // The spec is ambiguous here, see
+  // https://github.com/WebAssembly/spec/issues/1138. The WPT tests expect
+  // the raw value of the Content-Type attribute lowercased. We ignore this
+  // for file:// because file fetches don't have a Content-Type.
+  if !url.starts_with("file://") {
+    let headers = get_value(scope, object, "headers")?;
+    let content_type = call_method(scope, headers, "get", "Content-Type")?;
+
+    if content_type.to_lowercase() != "application/wasm" {
+      return Err(type_error("Response is not a wasm file."));
+    }
+  }
+
+  // 2.5
+  let ok = get_value(scope, object, "ok")?;
+  if !ok.is_true() {
+    return Err(type_error("Response is not ok."));
+  }
+
+  let body = get_value(scope, object, "body")?;
+
+  if body.is_null() {
+    return Ok(None);
+  }
+  let body = body
+    .to_object(scope)
+    .ok_or_else(|| type_error("Failed to get body object."))?;
+  let rid = get_value(scope, body, "rid")?
+    .to_uint32(scope)
+    .ok_or_else(|| type_error("Failed to get rid."))?
+    .value() as ResourceId;
+
+  Ok(Some((url, rid)))
+}
+
+fn get_value<'a, 'b>(
+  scope: &'b mut v8::HandleScope<'a>,
+  obj: v8::Local<'a, v8::Object>,
+  key: &'static str,
+) -> Result<v8::Local<'a, v8::Value>, AnyError> {
+  let key = v8::String::new(scope, key)
+    .ok_or_else(|| type_error("Failed to create key."))?;
+  Ok(
+    obj
+      .get(scope, key.into())
+      .ok_or_else(|| type_error("Failed to get value."))?,
+  )
+}
+
+fn get_string(
+  scope: &mut v8::HandleScope,
+  obj: v8::Local<v8::Object>,
+  key: &'static str,
+) -> Result<String, AnyError> {
+  let key = v8::String::new(scope, key)
+    .ok_or_else(|| type_error("Failed to create key."))?;
+  let value = obj
+    .get(scope, key.into())
+    .ok_or_else(|| type_error("Failed to get value."))?;
+
+  Ok(value.to_rust_string_lossy(scope))
+}
+
+fn call_method<'a>(
+  scope: &mut v8::HandleScope<'a>,
+  obj: v8::Local<'a, v8::Value>,
+  method: &'static str,
+  arg: &'static str,
+) -> Result<String, AnyError> {
+  let key = v8::String::new(scope, method)
+    .ok_or_else(|| type_error("Failed to create key."))?;
+  let function = obj
+    .to_object(scope)
+    .ok_or_else(|| type_error("Failed to create object."))?;
+  let function = function
+    .get(scope, key.into())
+    .ok_or_else(|| type_error("Failed to get value."))?;
+  let function: v8::Local<v8::Function> = function.try_into()?;
+  let arg = v8::String::new(scope, arg)
+    .ok_or_else(|| type_error("Failed to create arg."))?;
+  Ok(
+    function
+      .call(scope, obj, &[arg.into()])
+      .ok_or_else(|| type_error("Failed to call."))?
+      .to_rust_string_lossy(scope),
+  )
+}

runtime/js/99_main.js

@@ -457,7 +457,8 @@ function runtimeStart(
   tsVersion,
   target,
 ) {
-  core.setWasmStreamingCallback(fetch.handleWasmStreaming);
+  core.setMacrotaskCallback(timers.handleTimerMacrotask);
+  core.setMacrotaskCallback(promiseRejectMacrotaskCallback);
   core.setReportExceptionCallback(event.reportException);
   op_set_format_exception_callback(formatException);
   version.setVersions(

runtime/worker.rs

@@ -557,6 +557,12 @@ impl MainWorker {
     let bootstrap_fn = v8::Local::new(scope, bootstrap_fn);
     let undefined = v8::undefined(scope);
     bootstrap_fn.call(scope, undefined.into(), &[args]).unwrap();
+
+    // Set Wasm streaming callback
+    deno_core::set_wasm_streaming_callback(
+      scope,
+      deno_fetch::handle_wasm_streaming,
+    );
   }
 
   /// See [JsRuntime::execute_script](deno_core::JsRuntime::execute_script)