Improve robustness of cache loading/storing

crates/api/src/frame_info.rs

@@ -69,7 +69,7 @@ impl GlobalFrameInfo {
             if end > max {
                 max = end;
             }
-            let func_index = module.module().func_index(i);
+            let func_index = module.module().local.func_index(i);
             assert!(functions.insert(end, (start, func_index)).is_none());
         }
         if functions.len() == 0 {

crates/api/src/trampoline/create_handle.rs

@@ -27,6 +27,7 @@ pub(crate) fn create_handle(
 
     // Compute indices into the shared signature table.
     let signatures = module
+        .local
         .signatures
         .values()
         .map(|sig| store.compiler().signatures().register(sig))

crates/api/src/trampoline/func.rs

@@ -81,7 +81,8 @@ unsafe extern "C" fn stub_fn(
 
         let (args, returns_len) = {
             let module = instance.module_ref();
-            let signature = &module.signatures[module.functions[FuncIndex::new(call_id as usize)]];
+            let signature =
+                &module.local.signatures[module.local.functions[FuncIndex::new(call_id as usize)]];
 
             let mut args = Vec::new();
             for i in 2..signature.params.len() {
@@ -101,7 +102,8 @@ unsafe extern "C" fn stub_fn(
         state.func.call(&args, &mut returns)?;
 
         let module = instance.module_ref();
-        let signature = &module.signatures[module.functions[FuncIndex::new(call_id as usize)]];
+        let signature =
+            &module.local.signatures[module.local.functions[FuncIndex::new(call_id as usize)]];
         for (i, ret) in returns.iter_mut().enumerate() {
             if ret.ty().get_wasmtime_type() != Some(signature.returns[i].value_type) {
                 return Err(Trap::new(
@@ -266,8 +268,8 @@ pub fn create_handle_with_function(
         PrimaryMap::new();
     let mut code_memory = CodeMemory::new();
 
-    let sig_id = module.signatures.push(sig.clone());
-    let func_id = module.functions.push(sig_id);
+    let sig_id = module.local.signatures.push(sig.clone());
+    let func_id = module.local.functions.push(sig_id);
     module
         .exports
         .insert("trampoline".to_string(), Export::Function(func_id));
@@ -313,8 +315,8 @@ pub unsafe fn create_handle_with_raw_function(
     let mut module = Module::new();
     let mut finished_functions = PrimaryMap::new();
 
-    let sig_id = module.signatures.push(sig.clone());
-    let func_id = module.functions.push(sig_id);
+    let sig_id = module.local.signatures.push(sig.clone());
+    let func_id = module.local.functions.push(sig_id);
     module
         .exports
         .insert("trampoline".to_string(), Export::Function(func_id));

crates/api/src/trampoline/global.rs

@@ -25,7 +25,7 @@ pub fn create_global(store: &Store, gt: &GlobalType, val: Val) -> Result<Instanc
         },
     };
     let mut module = Module::new();
-    let global_id = module.globals.push(global);
+    let global_id = module.local.globals.push(global);
     module.exports.insert(
         "global".to_string(),
         wasmtime_environ::Export::Global(global_id),

crates/api/src/trampoline/memory.rs

@@ -17,7 +17,7 @@ pub fn create_handle_with_memory(store: &Store, memory: &MemoryType) -> Result<I
     let tunable = Default::default();
 
     let memory_plan = wasmtime_environ::MemoryPlan::for_memory(memory, &tunable);
-    let memory_id = module.memory_plans.push(memory_plan);
+    let memory_id = module.local.memory_plans.push(memory_plan);
     module.exports.insert(
         "memory".to_string(),
         wasmtime_environ::Export::Memory(memory_id),

crates/api/src/trampoline/table.rs

@@ -20,7 +20,7 @@ pub fn create_handle_with_table(store: &Store, table: &TableType) -> Result<Inst
     let tunable = Default::default();
 
     let table_plan = wasmtime_environ::TablePlan::for_table(table, &tunable);
-    let table_id = module.table_plans.push(table_plan);
+    let table_id = module.local.table_plans.push(table_plan);
     module.exports.insert(
         "table".to_string(),
         wasmtime_environ::Export::Table(table_id),

crates/environ/build.rs

@@ -1,9 +1,10 @@
 use std::process::Command;
+use std::str;
 
 fn main() {
     let git_rev = match Command::new("git").args(&["rev-parse", "HEAD"]).output() {
-        Ok(output) => String::from_utf8(output.stdout).unwrap(),
-        Err(_) => String::from("git-not-found"),
+        Ok(output) => str::from_utf8(&output.stdout).unwrap().trim().to_string(),
+        Err(_) => env!("CARGO_PKG_VERSION").to_string(),
     };
     println!("cargo:rustc-env=GIT_REV={}", git_rev);
 }

crates/environ/src/cache.rs

@@ -1,14 +1,13 @@
 use crate::address_map::{ModuleAddressMap, ValueLabelsRanges};
 use crate::compilation::{Compilation, Relocations, Traps};
-use crate::module::Module;
-use crate::module_environ::FunctionBodyData;
-use cranelift_codegen::{ir, isa};
+use cranelift_codegen::ir;
 use cranelift_entity::PrimaryMap;
 use cranelift_wasm::DefinedFuncIndex;
 use log::{debug, trace, warn};
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
 use std::fs;
+use std::hash::Hash;
 use std::hash::Hasher;
 use std::io::Write;
 use std::path::{Path, PathBuf};
@@ -23,7 +22,7 @@ use worker::Worker;
 pub struct ModuleCacheEntry<'config>(Option<ModuleCacheEntryInner<'config>>);
 
 struct ModuleCacheEntryInner<'config> {
-    mod_cache_path: PathBuf,
+    root_path: PathBuf,
     cache_config: &'config CacheConfig,
 }
 
@@ -51,21 +50,10 @@ pub type ModuleCacheDataTupleType = (
 struct Sha256Hasher(Sha256);
 
 impl<'config> ModuleCacheEntry<'config> {
-    pub fn new<'data>(
-        module: &Module,
-        function_body_inputs: &PrimaryMap<DefinedFuncIndex, FunctionBodyData<'data>>,
-        isa: &dyn isa::TargetIsa,
-        compiler_name: &str,
-        generate_debug_info: bool,
-        cache_config: &'config CacheConfig,
-    ) -> Self {
+    pub fn new<'data>(compiler_name: &str, cache_config: &'config CacheConfig) -> Self {
         if cache_config.enabled() {
             Self(Some(ModuleCacheEntryInner::new(
-                module,
-                function_body_inputs,
-                isa,
                 compiler_name,
-                generate_debug_info,
                 cache_config,
             )))
         } else {
@@ -78,42 +66,53 @@ impl<'config> ModuleCacheEntry<'config> {
         Self(Some(inner))
     }
 
-    pub fn get_data(&self) -> Option<ModuleCacheData> {
-        if let Some(inner) = &self.0 {
-            inner.get_data().map(|val| {
-                inner
-                    .cache_config
-                    .worker()
-                    .on_cache_get_async(&inner.mod_cache_path); // call on success
-                val
-            })
-        } else {
-            None
-        }
-    }
+    pub fn get_data<T: Hash, E>(
+        &self,
+        state: T,
+        compute: fn(T) -> Result<ModuleCacheDataTupleType, E>,
+    ) -> Result<ModuleCacheData, E> {
+        let mut hasher = Sha256Hasher(Sha256::new());
+        state.hash(&mut hasher);
+        let hash: [u8; 32] = hasher.0.result().into();
+        // standard encoding uses '/' which can't be used for filename
+        let hash = base64::encode_config(&hash, base64::URL_SAFE_NO_PAD);
 
-    pub fn update_data(&self, data: &ModuleCacheData) {
-        if let Some(inner) = &self.0 {
-            if inner.update_data(data).is_some() {
-                inner
-                    .cache_config
-                    .worker()
-                    .on_cache_update_async(&inner.mod_cache_path); // call on success
-            }
+        let inner = match &self.0 {
+            Some(inner) => inner,
+            None => return compute(state).map(ModuleCacheData::from_tuple),
+        };
+
+        if let Some(cached_val) = inner.get_data(&hash) {
+            let mod_cache_path = inner.root_path.join(&hash);
+            inner
+                .cache_config
+                .worker()
+                .on_cache_get_async(&mod_cache_path); // call on success
+            return Ok(cached_val);
+        }
+        let val_to_cache = ModuleCacheData::from_tuple(compute(state)?);
+        if inner.update_data(&hash, &val_to_cache).is_some() {
+            let mod_cache_path = inner.root_path.join(&hash);
+            inner
+                .cache_config
+                .worker()
+                .on_cache_update_async(&mod_cache_path); // call on success
         }
+        Ok(val_to_cache)
     }
 }
 
 impl<'config> ModuleCacheEntryInner<'config> {
-    fn new<'data>(
-        module: &Module,
-        function_body_inputs: &PrimaryMap<DefinedFuncIndex, FunctionBodyData<'data>>,
-        isa: &dyn isa::TargetIsa,
-        compiler_name: &str,
-        generate_debug_info: bool,
-        cache_config: &'config CacheConfig,
-    ) -> Self {
-        let hash = Sha256Hasher::digest(module, function_body_inputs);
+    fn new<'data>(compiler_name: &str, cache_config: &'config CacheConfig) -> Self {
+        // If debug assertions are enabled then assume that we're some sort of
+        // local build. We don't want local builds to stomp over caches between
+        // builds, so just use a separate cache directory based on the mtime of
+        // our executable, which should roughly correlate with "you changed the
+        // source code so you get a different directory".
+        //
+        // Otherwise if this is a release build we use the `GIT_REV` env var
+        // which is either the git rev if installed from git or the crate
+        // version if installed from crates.io.
         let compiler_dir = if cfg!(debug_assertions) {
             fn self_mtime() -> Option<String> {
                 let path = std::env::current_exe().ok()?;
@@ -138,26 +137,18 @@ impl<'config> ModuleCacheEntryInner<'config> {
                 comp_ver = env!("GIT_REV"),
             )
         };
-        let mod_filename = format!(
-            "mod-{mod_hash}{mod_dbg}",
-            mod_hash = base64::encode_config(&hash, base64::URL_SAFE_NO_PAD), // standard encoding uses '/' which can't be used for filename
-            mod_dbg = if generate_debug_info { ".d" } else { "" },
-        );
-        let mod_cache_path = cache_config
-            .directory()
-            .join(isa.triple().to_string())
-            .join(compiler_dir)
-            .join(mod_filename);
+        let root_path = cache_config.directory().join(compiler_dir);
 
         Self {
-            mod_cache_path,
+            root_path,
             cache_config,
         }
     }
 
-    fn get_data(&self) -> Option<ModuleCacheData> {
-        trace!("get_data() for path: {}", self.mod_cache_path.display());
-        let compressed_cache_bytes = fs::read(&self.mod_cache_path).ok()?;
+    fn get_data(&self, hash: &str) -> Option<ModuleCacheData> {
+        let mod_cache_path = self.root_path.join(hash);
+        trace!("get_data() for path: {}", mod_cache_path.display());
+        let compressed_cache_bytes = fs::read(&mod_cache_path).ok()?;
         let cache_bytes = zstd::decode_all(&compressed_cache_bytes[..])
             .map_err(|err| warn!("Failed to decompress cached code: {}", err))
             .ok()?;
@@ -166,8 +157,9 @@ impl<'config> ModuleCacheEntryInner<'config> {
             .ok()
     }
 
-    fn update_data(&self, data: &ModuleCacheData) -> Option<()> {
-        trace!("update_data() for path: {}", self.mod_cache_path.display());
+    fn update_data(&self, hash: &str, data: &ModuleCacheData) -> Option<()> {
+        let mod_cache_path = self.root_path.join(hash);
+        trace!("update_data() for path: {}", mod_cache_path.display());
         let serialized_data = bincode::serialize(&data)
             .map_err(|err| warn!("Failed to serialize cached code: {}", err))
             .ok()?;
@@ -180,17 +172,17 @@ impl<'config> ModuleCacheEntryInner<'config> {
 
         // Optimize syscalls: first, try writing to disk. It should succeed in most cases.
         // Otherwise, try creating the cache directory and retry writing to the file.
-        if fs_write_atomic(&self.mod_cache_path, "mod", &compressed_data) {
+        if fs_write_atomic(&mod_cache_path, "mod", &compressed_data) {
             return Some(());
         }
 
         debug!(
             "Attempting to create the cache directory, because \
              failed to write cached code to disk, path: {}",
-            self.mod_cache_path.display(),
+            mod_cache_path.display(),
         );
 
-        let cache_dir = self.mod_cache_path.parent().unwrap();
+        let cache_dir = mod_cache_path.parent().unwrap();
         fs::create_dir_all(cache_dir)
             .map_err(|err| {
                 warn!(
@@ -201,7 +193,7 @@ impl<'config> ModuleCacheEntryInner<'config> {
             })
             .ok()?;
 
-        if fs_write_atomic(&self.mod_cache_path, "mod", &compressed_data) {
+        if fs_write_atomic(&mod_cache_path, "mod", &compressed_data) {
             Some(())
         } else {
             None
@@ -233,17 +225,6 @@ impl ModuleCacheData {
     }
 }
 
-impl Sha256Hasher {
-    pub fn digest<'data>(
-        module: &Module,
-        function_body_inputs: &PrimaryMap<DefinedFuncIndex, FunctionBodyData<'data>>,
-    ) -> [u8; 32] {
-        let mut hasher = Self(Sha256::new());
-        module.hash_for_cache(function_body_inputs, &mut hasher);
-        hasher.0.result().into()
-    }
-}
-
 impl Hasher for Sha256Hasher {
     fn finish(&self) -> u64 {
         panic!("Sha256Hasher doesn't support finish!");