feat: EXC-2087: Use DeterministicHeapBytes for compilation cache

packages/ic-deterministic-heap-bytes/BUILD.bazel

@@ -14,6 +14,7 @@ rust_library(
     deps = [
         # Keep sorted.
         "@crate_index//:candid",
+        "@crate_index//:tempfile",
     ],
 )
 

packages/ic-deterministic-heap-bytes/Cargo.toml

@@ -13,3 +13,4 @@ path = "src/lib.rs"
 candid = { workspace = true }
 ic-deterministic-heap-bytes-derive = { path = "../ic-deterministic-heap-bytes-derive" }
 paste = { workspace = true }
+tempfile = { workspace = true }

packages/ic-deterministic-heap-bytes/src/lib.rs

@@ -29,6 +29,9 @@ pub trait DeterministicHeapBytes {
     }
 }
 
+////////////////////////////////////////////////////////////////////////
+// Scalar types.
+
 impl DeterministicHeapBytes for u8 {}
 impl DeterministicHeapBytes for u16 {}
 impl DeterministicHeapBytes for u32 {}
@@ -46,6 +49,11 @@ impl DeterministicHeapBytes for f64 {}
 impl DeterministicHeapBytes for bool {}
 impl DeterministicHeapBytes for char {}
 
+////////////////////////////////////////////////////////////////////////
+// Standard library types.
+
+impl DeterministicHeapBytes for std::sync::atomic::AtomicU64 {}
+
 impl DeterministicHeapBytes for std::time::Duration {}
 
 impl DeterministicHeapBytes for String {
@@ -139,7 +147,16 @@ impl<T: DeterministicHeapBytes, E: DeterministicHeapBytes> DeterministicHeapByte
     }
 }
 
+////////////////////////////////////////////////////////////////////////
+// External types.
+
 impl DeterministicHeapBytes for candid::Principal {}
+impl DeterministicHeapBytes for tempfile::TempDir {
+    fn deterministic_heap_bytes(&self) -> usize {
+        // TempDir allocates a string for the path.
+        self.path().as_os_str().len()
+    }
+}
 
 #[cfg(test)]
 mod tests;

rs/embedders/BUILD.bazel

@@ -7,6 +7,7 @@ package(default_visibility = ["//visibility:public"])
 
 DEPENDENCIES = [
     # Keep sorted.
+    "//packages/ic-deterministic-heap-bytes",
     "//packages/ic-error-types",
     "//rs/config",
     "//rs/cycles_account_manager",
@@ -50,7 +51,10 @@ DEPENDENCIES = [
     "@crate_index//:wasmtime-environ",
 ]
 
-MACRO_DEPENDENCIES = []
+MACRO_DEPENDENCIES = [
+    # Keep sorted.
+    "//packages/ic-deterministic-heap-bytes-derive",
+]
 
 DEV_DEPENDENCIES = [
     # Keep sorted.

rs/embedders/Cargo.toml

@@ -13,6 +13,7 @@ clap = { workspace = true }
 ic-btc-interface = { workspace = true }
 ic-config = { path = "../config" }
 ic-cycles-account-manager = { path = "../cycles_account_manager" }
+ic-deterministic-heap-bytes = { path = "../../packages/ic-deterministic-heap-bytes" }
 ic-error-types = { path = "../../packages/ic-error-types" }
 ic-interfaces = { path = "../interfaces" }
 ic-logger = { path = "../monitoring/logger" }

rs/embedders/src/compilation_cache.rs

@@ -9,6 +9,7 @@ use std::{
 use tempfile::TempDir;
 
 use crate::{OnDiskSerializedModule, SerializedModule};
+use ic_deterministic_heap_bytes::DeterministicHeapBytes;
 use ic_interfaces::execution_environment::{HypervisorError, HypervisorResult};
 use ic_types::{MemoryDiskBytes, NumBytes};
 use ic_utils_lru_cache::LruCache;
@@ -29,11 +30,14 @@ const DEFAULT_MEMORY_CAPACITY: NumBytes = NumBytes::new(10 * GB);
 
 /// Stores the serialized modules of wasm code that has already been compiled so
 /// that it can be used again without recompiling.
+#[derive(DeterministicHeapBytes)]
 pub struct CompilationCache {
     /// Directory holding all the temporary files. It will be deleted on
     /// drop.
     dir: TempDir,
     /// Map from wasm hash to an open fd with the serialized Module result.
+    #[deterministic_heap_bytes(with = |c: &Mutex<LruCache<WasmHash, HypervisorResult<Arc<OnDiskSerializedModule>>>>|
+        c.lock().unwrap().memory_bytes())]
     cache: Mutex<LruCache<WasmHash, HypervisorResult<Arc<OnDiskSerializedModule>>>>,
     /// Atomic counter to deduplicate files in the case of concurrent compilations of the same module.
     counter: AtomicU64,
@@ -43,7 +47,7 @@ pub struct CompilationCache {
 
 impl MemoryDiskBytes for CompilationCache {
     fn memory_bytes(&self) -> usize {
-        self.cache.lock().unwrap().memory_bytes()
+        unreachable!("To be removed in favor of DeterministicHeapBytes")
     }
 
     fn disk_bytes(&self) -> usize {

rs/execution_environment/src/hypervisor.rs

@@ -2,6 +2,7 @@ use ic_canister_sandbox_backend_lib::replica_controller::sandboxed_execution_con
 use ic_config::execution_environment::{Config, MAX_COMPILATION_CACHE_SIZE};
 use ic_config::flag_status::FlagStatus;
 use ic_cycles_account_manager::CyclesAccountManager;
+use ic_deterministic_heap_bytes::DeterministicHeapBytes;
 use ic_embedders::{
     wasm_executor::{WasmExecutionResult, WasmExecutor, WasmExecutorImpl},
     wasm_utils::decoding::decoded_wasm_size,
@@ -166,7 +167,7 @@ impl Hypervisor {
                 if let Some(compilation_result) = compilation_result {
                     self.metrics.observe_compilation_metrics(
                         &compilation_result,
-                        self.compilation_cache.memory_bytes(),
+                        self.compilation_cache.deterministic_heap_bytes(),
                         self.compilation_cache.disk_bytes(),
                     );
                 }
@@ -409,7 +410,7 @@ impl Hypervisor {
         if let Some(compilation_result) = compilation_result {
             self.metrics.observe_compilation_metrics(
                 &compilation_result,
-                self.compilation_cache.memory_bytes(),
+                self.compilation_cache.deterministic_heap_bytes(),
                 self.compilation_cache.disk_bytes(),
             );
         }