Cache compression (#234)

src/bin/wasm2obj.rs

@@ -63,7 +63,7 @@ The translation is dependent on the environment chosen.
 The default is a dummy environment that produces placeholder values.
 
 Usage:
-    wasm2obj [--target TARGET] [-dg] [--cache | --cache-dir=<cache_dir>] [--enable-simd] <file> -o <output>
+    wasm2obj [--target TARGET] [-dg] [--cache] [--cache-dir=<cache_dir>] [--cache-compression-level=<compr_level>] [--enable-simd] <file> -o <output>
     wasm2obj --help | --version
 
 Options:
@@ -74,6 +74,8 @@ Options:
     -c, --cache         enable caching system, use default cache directory
     --cache-dir=<cache_dir>
                         enable caching system, use specified cache directory
+    --cache-compression-level=<compr_level>
+                        enable caching system, use custom compression level for new cache, values 1-21
     --enable-simd       enable proposed SIMD instructions
     --version           print the Cranelift version
     -d, --debug         enable debug output on stderr/stdout
@@ -88,6 +90,7 @@ struct Args {
     flag_debug: bool,
     flag_cache: bool,
     flag_cache_dir: Option<String>,
+    flag_cache_compression_level: Option<i32>,
     flag_enable_simd: bool,
 }
 
@@ -115,8 +118,11 @@ fn main() {
     }
 
     cache_conf::init(
-        args.flag_cache || args.flag_cache_dir.is_some(),
+        args.flag_cache
+            || args.flag_cache_dir.is_some()
+            || args.flag_cache_compression_level.is_some(),
         args.flag_cache_dir.as_ref(),
+        args.flag_cache_compression_level,
     );
 
     let path = Path::new(&args.arg_file);

src/bin/wasmtime.rs

@@ -64,8 +64,8 @@ including calling the start function if one is present. Additional functions
 given with --invoke are then called.
 
 Usage:
-    wasmtime [-odg] [--enable-simd] [--wasi-c] [--cache | --cache-dir=<cache_dir>] [--preload=<wasm>...] [--env=<env>...] [--dir=<dir>...] [--mapdir=<mapping>...] <file> [<arg>...]
+    wasmtime [-odg] [--enable-simd] [--wasi-c] [--cache] [--cache-dir=<cache_dir>] [--cache-compression-level=<compr_level>] [--preload=<wasm>...] [--env=<env>...] [--dir=<dir>...] [--mapdir=<mapping>...] <file> [<arg>...]
-    wasmtime [-odg] [--enable-simd] [--wasi-c] [--cache | --cache-dir=<cache_dir>] [--preload=<wasm>...] [--env=<env>...] [--dir=<dir>...] [--mapdir=<mapping>...] --invoke=<fn> <file> [<arg>...]
+    wasmtime [-odg] [--enable-simd] [--wasi-c] [--cache] [--cache-dir=<cache_dir>] [--cache-compression-level=<compr_level>] [--env=<env>...] [--dir=<dir>...] [--mapdir=<mapping>...] --invoke=<fn> <file> [<arg>...]
     wasmtime --help | --version
 
 Options:
@@ -74,6 +74,8 @@ Options:
     -c, --cache         enable caching system, use default cache directory
     --cache-dir=<cache_dir>
                         enable caching system, use specified cache directory
+    --cache-compression-level=<compr_level>
+                        enable caching system, use custom compression level for new cache, values 1-21
     -g                  generate debug information
     -d, --debug         enable debug output on stderr/stdout
     --enable-simd       enable proposed SIMD instructions
@@ -94,6 +96,7 @@ struct Args {
     flag_optimize: bool,
     flag_cache: bool,
     flag_cache_dir: Option<String>,
+    flag_cache_compression_level: Option<i32>,
     flag_debug: bool,
     flag_g: bool,
     flag_enable_simd: bool,
@@ -218,8 +221,11 @@ fn rmain() -> Result<(), Error> {
     }
 
     cache_conf::init(
-        args.flag_cache || args.flag_cache_dir.is_some(),
+        args.flag_cache
+            || args.flag_cache_dir.is_some()
+            || args.flag_cache_compression_level.is_some(),
         args.flag_cache_dir.as_ref(),
+        args.flag_cache_compression_level,
     );
 
     let mut flag_builder = settings::builder();

src/bin/wast.rs

@@ -41,7 +41,7 @@ const USAGE: &str = "
 Wast test runner.
 
 Usage:
-    wast [-do] [--enable-simd] [--cache | --cache-dir=<cache_dir>] <file>...
+    wast [-do] [--enable-simd] [--cache] [--cache-dir=<cache_dir>] [--cache-compression-level=<compr_level>] <file>...
     wast --help | --version
 
 Options:
@@ -51,6 +51,8 @@ Options:
     -c, --cache         enable caching system, use default cache directory
     --cache-dir=<cache_dir>
                         enable caching system, use specified cache directory
+    --cache-compression-level=<compr_level>
+                        enable caching system, use custom compression level for new cache, values 1-21
     -d, --debug         enable debug output on stderr/stdout
     --enable-simd       enable proposed SIMD instructions
 ";
@@ -63,6 +65,7 @@ struct Args {
     flag_optimize: bool,
     flag_cache: bool,
     flag_cache_dir: Option<String>,
+    flag_cache_compression_level: Option<i32>,
     flag_enable_simd: bool,
 }
 
@@ -83,8 +86,11 @@ fn main() {
     }
 
     cache_conf::init(
-        args.flag_cache || args.flag_cache_dir.is_some(),
+        args.flag_cache
+            || args.flag_cache_dir.is_some()
+            || args.flag_cache_compression_level.is_some(),
         args.flag_cache_dir.as_ref(),
+        args.flag_cache_compression_level,
     );
 
     let isa_builder = cranelift_native::builder().unwrap_or_else(|_| {

wasmtime-environ/Cargo.toml

@@ -28,6 +28,7 @@ bincode = "1.1.4"
 lazy_static = "1.3.0"
 spin = "0.5.0"
 log = { version = "0.4.8", default-features = false }
+zstd = "0.4"
 
 [dev-dependencies]
 tempfile = "3"

wasmtime-environ/src/cache.rs

@@ -29,11 +29,13 @@ pub mod conf {
     struct Config {
         pub cache_enabled: bool,
         pub cache_dir: PathBuf,
+        pub compression_level: i32,
     }
 
     // Private static, so only internal function can access it.
     static CONFIG: Once<Config> = Once::new();
     static INIT_CALLED: AtomicBool = AtomicBool::new(false);
+    static DEFAULT_COMPRESSION_LEVEL: i32 = 0; // 0 for zstd means "use default level"
 
     /// Returns true if and only if the cache is enabled.
     pub fn cache_enabled() -> bool {
@@ -54,9 +56,19 @@ pub mod conf {
             .cache_dir
     }
 
+    /// Returns cache compression level.
+    ///
+    /// Panics if the cache is disabled.
+    pub fn compression_level() -> i32 {
+        CONFIG
+            .r#try()
+            .expect("Cache system must be initialized")
+            .compression_level
+    }
+
     /// Initializes the cache system. Should be called exactly once,
     /// and before using the cache system. Otherwise it can panic.
-    pub fn init<P: AsRef<Path>>(enabled: bool, dir: Option<P>) {
+    pub fn init<P: AsRef<Path>>(enabled: bool, dir: Option<P>, compression_level: Option<i32>) {
         INIT_CALLED
             .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
             .expect("Cache system init must be called at most once");
@@ -64,10 +76,16 @@ pub mod conf {
             CONFIG.r#try().is_none(),
             "Cache system init must be called before using the system."
         );
-        let conf = CONFIG.call_once(|| Config::new(enabled, dir));
+        let conf = CONFIG.call_once(|| {
+            Config::new(
+                enabled,
+                dir,
+                compression_level.unwrap_or(DEFAULT_COMPRESSION_LEVEL),
+            )
+        });
         debug!(
-            "Cache init(): enabled={}, cache-dir={:?}",
+            "Cache init(): enabled={}, cache-dir={:?}, compression-level={}",
-            conf.cache_enabled, conf.cache_dir
+            conf.cache_enabled, conf.cache_dir, conf.compression_level,
         );
     }
 
@@ -76,15 +94,18 @@ pub mod conf {
             Self {
                 cache_enabled: false,
                 cache_dir: PathBuf::new(),
+                compression_level: DEFAULT_COMPRESSION_LEVEL,
             }
         }
 
-        pub fn new<P: AsRef<Path>>(enabled: bool, dir: Option<P>) -> Self {
+        pub fn new<P: AsRef<Path>>(enabled: bool, dir: Option<P>, compression_level: i32) -> Self {
             if enabled {
                 match dir {
-                    Some(dir) => Self::new_step2(dir.as_ref()),
+                    Some(dir) => Self::new_step2(dir.as_ref(), compression_level),
                     None => match ProjectDirs::from("", "CraneStation", "wasmtime") {
-                        Some(proj_dirs) => Self::new_step2(proj_dirs.cache_dir()),
+                        Some(proj_dirs) => {
+                            Self::new_step2(proj_dirs.cache_dir(), compression_level)
+                        }
                         None => {
                             warn!("Cache directory not specified and failed to find the default. Disabling cache.");
                             Self::new_cache_disabled()
@@ -96,7 +117,7 @@ pub mod conf {
             }
         }
 
-        fn new_step2(dir: &Path) -> Self {
+        fn new_step2(dir: &Path, compression_level: i32) -> Self {
             // On Windows, if we want long paths, we need '\\?\' prefix, but it doesn't work
             // with relative paths. One way to get absolute path (the only one?) is to use
             // fs::canonicalize, but it requires that given path exists. The extra advantage
@@ -106,6 +127,7 @@ pub mod conf {
                     Ok(p) => Self {
                         cache_enabled: true,
                         cache_dir: p,
+                        compression_level,
                     },
                     Err(err) => {
                         warn!(
@@ -223,6 +245,7 @@ impl ModuleCacheEntry {
     pub fn get_data(&self) -> Option<ModuleCacheData> {
         if let Some(p) = &self.mod_cache_path {
             match fs::read(p) {
+                Ok(compressed_cache_bytes) => match zstd::decode_all(&compressed_cache_bytes[..]) {
                     Ok(cache_bytes) => match bincode::deserialize(&cache_bytes[..]) {
                         Ok(data) => Some(data),
                         Err(err) => {
@@ -230,6 +253,11 @@ impl ModuleCacheEntry {
                             None
                         }
                     },
+                    Err(err) => {
+                        warn!("Failed to decompress cached code: {}", err);
+                        None
+                    }
+                },
                 Err(_) => None,
             }
         } else {
@@ -240,7 +268,13 @@ impl ModuleCacheEntry {
     pub fn update_data(&self, data: &ModuleCacheData) {
         if let Some(p) = &self.mod_cache_path {
             let cache_buf = match bincode::serialize(&data) {
+                Ok(data) => match zstd::encode_all(&data[..], conf::compression_level()) {
                     Ok(data) => data,
+                    Err(err) => {
+                        warn!("Failed to compress cached code: {}", err);
+                        return;
+                    }
+                },
                 Err(err) => {
                     warn!("Failed to serialize cached code: {}", err);
                     return;

wasmtime-environ/src/cache/tests.rs

@@ -26,13 +26,15 @@ use tempfile;
 fn test_write_read_cache() {
     pretty_env_logger::init();
     let dir = tempfile::tempdir().expect("Can't create temporary directory");
-    conf::init(true, Some(dir.path()));
+    let compression_level = 5;
+    conf::init(true, Some(dir.path()), Some(compression_level));
     assert!(conf::cache_enabled());
     // assumption: config init creates cache directory and returns canonicalized path
     assert_eq!(
         *conf::cache_directory(),
         fs::canonicalize(dir.path()).unwrap()
     );
+    assert_eq!(conf::compression_level(), compression_level);
 
     let mut rng = SmallRng::from_seed([
         0x42, 0x04, 0xF3, 0x44, 0x11, 0x22, 0x33, 0x44, 0x67, 0x68, 0xFF, 0x00, 0x44, 0x23, 0x7F,

wasmtime-environ/tests/cache_fail_calling_init_twice.rs

@@ -5,6 +5,6 @@ use wasmtime_environ::cache_conf;
 #[should_panic]
 fn test_fail_calling_init_twice() {
     let dir = tempfile::tempdir().expect("Can't create temporary directory");
-    cache_conf::init(true, Some(dir.path()));
+    cache_conf::init(true, Some(dir.path()), Some(5));
-    cache_conf::init(true, Some(dir.path()));
+    cache_conf::init(true, Some(dir.path()), Some(5));
 }