Made ModuleCache sync again

wasmerio · May 3, 2023 · d7a7080 · d7a7080
1 parent 7b71bec
commit d7a7080
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 21 deletions.
diff --git a/lib/wasi/src/bin_factory/exec.rs b/lib/wasi/src/bin_factory/exec.rs
@@ -26,14 +26,7 @@ pub fn spawn_exec(
     // The deterministic id for this engine
     let compiler = store.engine().deterministic_id();
 
-    let tasks = runtime.task_manager();
-
-    let module = tasks.block_on(compiled_modules.get_compiled_module(
-        &store,
-        binary.hash().as_str(),
-        compiler,
-        &**tasks,
-    ));
+    let module = compiled_modules.get_compiled_module(&**runtime, binary.hash().as_str(), compiler);
 
     let module = match (module, binary.entry.as_ref()) {
         (Some(a), _) => a,
@@ -52,12 +45,12 @@ pub fn spawn_exec(
             }
             let module = module?;
 
-            tasks.block_on(compiled_modules.set_compiled_module(
+            compiled_modules.set_compiled_module(
+                &**runtime,
                 binary.hash().as_str(),
                 compiler,
                 &module,
-                &**tasks,
-            ));
+            );
             module
         }
         (None, None) => {

diff --git a/lib/wasi/src/bin_factory/module_cache.rs b/lib/wasi/src/bin_factory/module_cache.rs
@@ -4,7 +4,7 @@ use anyhow::Context;
 use wasmer::Module;
 
 use super::BinaryPackage;
-use crate::{runtime::module_cache::CompiledModuleCache, VirtualTaskManager, WasiRuntime};
+use crate::{runtime::module_cache::CompiledModuleCache, WasiRuntime};
 
 pub const DEFAULT_COMPILED_PATH: &str = "~/.wasmer/compiled";
 
@@ -59,31 +59,34 @@ impl ModuleCache {
         })
     }
 
-    pub async fn get_compiled_module(
+    pub fn get_compiled_module(
         &self,
-        engine: &impl wasmer::AsEngineRef,
+        runtime: &dyn WasiRuntime,
         data_hash: &str,
         compiler: &str,
-        task_manager: &dyn VirtualTaskManager,
     ) -> Option<Module> {
         let key = format!("{}-{}", data_hash, compiler);
+        let engine = runtime.engine()?;
 
-        self.0
-            .load(&key, engine.as_engine_ref().engine(), task_manager)
-            .await
+        let tasks = runtime.task_manager();
+        tasks
+            .block_on(async { self.0.load(&key, &engine, tasks).await })
             .ok()
     }
 
-    pub async fn set_compiled_module(
+    pub fn set_compiled_module(
         &self,
+        runtime: &dyn WasiRuntime,
         data_hash: &str,
         compiler: &str,
         module: &Module,
-        task_manager: &dyn VirtualTaskManager,
     ) {
         let key = format!("{}-{}", data_hash, compiler);
 
-        if let Err(e) = self.0.save(&key, module, task_manager).await {
+        let tasks = runtime.task_manager();
+        let result = tasks.block_on(async { self.0.save(&key, module, tasks).await });
+
+        if let Err(e) = result {
             tracing::warn!(
                 data_hash,
                 compiler,