Stubbed out caching

wasmerio · Aug 21, 2023 · 48cbb11 · 48cbb11
1 parent c5dddc8
commit 48cbb11
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 190 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -3,7 +3,6 @@ wasm_bindgen_test::wasm_bindgen_test_configure!(run_in_browser);
 
 mod facade;
 mod instance;
-mod module_cache;
 mod net;
 mod run;
 mod runtime;

diff --git a/src/module_cache.rs b/src/module_cache.rs
diff --git a/src/tasks/pool.rs b/src/tasks/pool.rs
@@ -8,6 +8,7 @@ use std::{
 use anyhow::{Context, Error};
 use futures::{future::LocalBoxFuture, Future};
 use tokio::sync::mpsc::{self, UnboundedSender};
+use wasm_bindgen::{JsCast, JsValue};
 use wasmer_wasix::{
     runtime::{resolver::WebcHash, task_manager::TaskWasm},
     WasiThreadError,
@@ -164,7 +165,7 @@ impl Scheduler {
                 move_worker(worker_id, &mut self.busy, &mut self.idle)
             }
             Message::CacheModule { hash, module } => {
-                let module = js_sys::WebAssembly::Module::from(module);
+                let module: js_sys::WebAssembly::Module = JsValue::from(module).unchecked_into();
                 self.cached_modules.insert(hash, module.clone());
 
                 for worker in self.idle.iter().chain(self.busy.iter()) {
@@ -242,8 +243,13 @@ impl Scheduler {
         self.next_id += 1;
         let handle = WorkerHandle::spawn(id, self.mailbox.clone())?;
 
-        for (hash, module) in &self.cached_modules {
-            todo!();
+        // Prime the worker's module cache
+        for (&hash, module) in &self.cached_modules {
+            let msg = PostMessagePayload::CacheModule {
+                hash,
+                module: module.clone(),
+            };
+            handle.send(msg)?;
         }
 
         Ok(handle)

diff --git a/src/tasks/worker2.js → src/tasks/worker.js b/src/tasks/worker2.js → src/tasks/worker.js
diff --git a/src/tasks/worker.rs b/src/tasks/worker.rs
@@ -1,4 +1,4 @@
-use std::pin::Pin;
+use std::{mem::ManuallyDrop, pin::Pin};
 
 use anyhow::{Context, Error};
 use futures::Future;
@@ -9,6 +9,7 @@ use wasm_bindgen::{
     prelude::{wasm_bindgen, Closure},
     JsCast, JsValue,
 };
+use wasmer_wasix::runtime::resolver::WebcHash;
 
 use crate::tasks::pool::{Message, PostMessagePayload};
 
@@ -118,7 +119,7 @@ static WORKER_URL: Lazy<String> = Lazy::new(|| {
 
     tracing::debug!(import_url = IMPORT_META_URL.as_str());
 
-    let script = include_str!("worker2.js").replace("$IMPORT_META_URL", &IMPORT_META_URL);
+    let script = include_str!("worker.js").replace("$IMPORT_META_URL", &IMPORT_META_URL);
 
     let blob = web_sys::Blob::new_with_u8_array_sequence_and_options(
         Array::from_iter([Uint8Array::from(script.as_bytes())]).as_ref(),
@@ -136,30 +137,45 @@ static WORKER_URL: Lazy<String> = Lazy::new(|| {
 #[derive(serde::Serialize, serde::Deserialize)]
 #[serde(tag = "type", rename_all = "kebab-case")]
 pub(crate) enum PostMessagePayloadRepr {
-    SpawnAsync { ptr: usize },
-    SpawnBlocking { ptr: usize },
+    SpawnAsync {
+        ptr: usize,
+    },
+    SpawnBlocking {
+        ptr: usize,
+    },
+    #[serde(skip)]
+    CacheModule {
+        hash: WebcHash,
+        module: js_sys::WebAssembly::Module,
+    },
 }
 
 impl PostMessagePayloadRepr {
     pub(crate) unsafe fn reconstitute(self) -> PostMessagePayload {
-        match self {
+        let this = ManuallyDrop::new(self);
+
+        match &*this {
             PostMessagePayloadRepr::SpawnAsync { ptr } => {
                 let boxed = Box::from_raw(
-                    ptr as *mut Box<
+                    *ptr as *mut Box<
                         dyn FnOnce() -> Pin<Box<dyn Future<Output = ()> + 'static>>
                             + Send
                             + 'static,
                     >,
                 );
-                std::mem::forget(self);
                 PostMessagePayload::SpawnAsync(*boxed)
             }
 
             PostMessagePayloadRepr::SpawnBlocking { ptr } => {
-                let boxed = Box::from_raw(ptr as *mut Box<dyn FnOnce() + Send + 'static>);
-                std::mem::forget(self);
+                let boxed = Box::from_raw(*ptr as *mut Box<dyn FnOnce() + Send + 'static>);
                 PostMessagePayload::SpawnBlocking(*boxed)
             }
+            PostMessagePayloadRepr::CacheModule { hash, ref module } => {
+                PostMessagePayload::CacheModule {
+                    hash: std::ptr::read(hash),
+                    module: std::ptr::read(module),
+                }
+            }
         }
     }
 }
@@ -182,6 +198,9 @@ impl From<PostMessagePayload> for PostMessagePayloadRepr {
                     ptr: Box::into_raw(boxed) as usize,
                 }
             }
+            PostMessagePayload::CacheModule { hash, module } => {
+                PostMessagePayloadRepr::CacheModule { hash, module }
+            }
         }
     }
 }
@@ -291,6 +310,9 @@ pub async fn __worker_handle_message(msg: JsValue) -> Result<(), crate::utils::E
         match msg.reconstitute() {
             PostMessagePayload::SpawnAsync(thunk) => thunk().await,
             PostMessagePayload::SpawnBlocking(thunk) => thunk(),
+            PostMessagePayload::CacheModule { hash, .. } => {
+                tracing::warn!(%hash, "XXX Caching module");
+            }
         }
     }