chore: EXC-2087: Fix EntryEnv memory bytes

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

@@ -20,9 +20,4 @@ rust_library(
 rust_test(
     name = "ic-deterministic-heap-bytes-test",
     crate = ":ic-deterministic-heap-bytes",
-    proc_macro_deps = [
-        # Keep sorted.
-        "//packages/ic-deterministic-heap-bytes-derive",
-        "@crate_index//:paste",
-    ],
 )

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

@@ -1,6 +1,5 @@
 pub use ic_deterministic_heap_bytes_derive::DeterministicHeapBytes;
 use paste::paste;
-use std::collections::BTreeMap;
 
 /// A trait to deterministically report heap memory usage.
 ///
@@ -47,6 +46,8 @@ impl DeterministicHeapBytes for f64 {}
 impl DeterministicHeapBytes for bool {}
 impl DeterministicHeapBytes for char {}
 
+impl DeterministicHeapBytes for std::time::Duration {}
+
 impl DeterministicHeapBytes for String {
     fn deterministic_heap_bytes(&self) -> usize {
         self.len()
@@ -79,7 +80,7 @@ impl<T: DeterministicHeapBytes> DeterministicHeapBytes for Vec<T> {
 }
 
 impl<K: DeterministicHeapBytes, V: DeterministicHeapBytes> DeterministicHeapBytes
-    for BTreeMap<K, V>
+    for std::collections::BTreeMap<K, V>
 {
     /// Calculates the precise heap size by summing the heap usage of all elements.
     ///

packages/ic-error-types/BUILD.bazel

@@ -4,12 +4,14 @@ package(default_visibility = ["//visibility:public"])
 
 DEPENDENCIES = [
     # Keep sorted.
+    "//packages/ic-deterministic-heap-bytes",
     "@crate_index//:serde",
     "@crate_index//:strum",
 ]
 
 MACRO_DEPENDENCIES = [
     "@crate_index//:strum_macros",
+    "//packages/ic-deterministic-heap-bytes-derive",
 ]
 
 rust_library(

packages/ic-error-types/Cargo.toml

@@ -13,6 +13,7 @@ authors.workspace = true
 edition.workspace = true
 
 [dependencies]
+ic-deterministic-heap-bytes = { path = "../ic-deterministic-heap-bytes" }
 serde = { workspace = true }
 strum = { workspace = true }
 strum_macros = { workspace = true }

packages/ic-error-types/src/lib.rs

@@ -1,6 +1,7 @@
 //! A crate that groups user-facing and internal error types and codes produced
 //! by the Internet Computer.
 
+use ic_deterministic_heap_bytes::DeterministicHeapBytes;
 use serde::{Deserialize, Serialize};
 use std::{convert::TryFrom, fmt};
 use str_traits::StrEllipsize;
@@ -141,7 +142,18 @@ impl From<ErrorCode> for RejectCode {
 /// code and the rest is just a sequentially assigned two-digit
 /// number.
 #[derive(
-    Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug, Deserialize, EnumIter, Serialize,
+    Copy,
+    Clone,
+    Eq,
+    DeterministicHeapBytes,
+    PartialEq,
+    Ord,
+    PartialOrd,
+    Hash,
+    Debug,
+    Deserialize,
+    EnumIter,
+    Serialize,
 )]
 pub enum ErrorCode {
     // 1xx -- `RejectCode::SysFatal`
@@ -217,7 +229,18 @@ const MAX_USER_ERROR_DESCRIPTION_LEN_BYTES: usize = 8 * 1024;
 /// The error that is sent back to users of IC if something goes
 /// wrong. It's designed to be copyable and serializable so that we
 /// can persist it in ingress history.
-#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug, Deserialize, Serialize)]
+#[derive(
+    Clone,
+    Eq,
+    DeterministicHeapBytes,
+    PartialEq,
+    Ord,
+    PartialOrd,
+    Hash,
+    Debug,
+    Deserialize,
+    Serialize,
+)]
 pub struct UserError {
     code: ErrorCode,
     description: String,

rs/execution_environment/BUILD.bazel

@@ -5,6 +5,7 @@ package(default_visibility = ["//visibility:public"])
 
 DEPENDENCIES = [
     # Keep sorted.
+    "//packages/ic-deterministic-heap-bytes",
     "//packages/ic-error-types",
     "//rs/canister_sandbox:backend_lib",
     "//rs/config",
@@ -60,7 +61,10 @@ DEPENDENCIES = [
     "@crate_index//:wasmparser",
 ]
 
-MACRO_DEPENDENCIES = []
+MACRO_DEPENDENCIES = [
+    # Keep sorted.
+    "//packages/ic-deterministic-heap-bytes-derive",
+]
 
 DEV_DEPENDENCIES = [
     # Keep sorted.

rs/execution_environment/Cargo.toml

@@ -18,6 +18,7 @@ ic-crypto-sha2 = { path = "../crypto/sha2" }
 ic-crypto-tree-hash = { path = "../crypto/tree_hash" }
 ic-crypto-utils-canister-threshold-sig = { path = "../crypto/utils/canister_threshold_sig" }
 ic-cycles-account-manager = { path = "../cycles_account_manager" }
+ic-deterministic-heap-bytes = { path = "../../packages/ic-deterministic-heap-bytes" }
 ic-embedders = { path = "../embedders" }
 ic-error-types = { path = "../../packages/ic-error-types" }
 ic-interfaces = { path = "../interfaces" }

rs/execution_environment/src/query_handler/query_cache.rs

@@ -1,4 +1,5 @@
 use ic_base_types::{CanisterId, NumBytes};
+use ic_deterministic_heap_bytes::DeterministicHeapBytes;
 use ic_error_types::UserError;
 use ic_interfaces::execution_environment::SystemApiCallCounters;
 use ic_metrics::MetricsRegistry;
@@ -9,7 +10,7 @@ use ic_types::{
 };
 use ic_utils_lru_cache::LruCache;
 use prometheus::{Histogram, IntCounter, IntGauge};
-use std::{collections::BTreeMap, mem::size_of_val, sync::Mutex, time::Duration};
+use std::{collections::BTreeMap, sync::Mutex, time::Duration};
 
 use crate::metrics::duration_histogram;
 
@@ -125,7 +126,7 @@ impl QueryCacheMetrics {
 ///
 /// The key is to distinguish query cache entries, i.e. entries with different
 /// keys are (almost) completely independent from each other.
-#[derive(Clone, Eq, PartialEq, Hash)]
+#[derive(Clone, Eq, DeterministicHeapBytes, PartialEq, Hash)]
 pub(crate) struct EntryKey {
     /// Query source.
     pub source: UserId,
@@ -137,9 +138,10 @@ pub(crate) struct EntryKey {
     pub method_payload: Vec<u8>,
 }
 
+// The MemoryDiskBytes trait is still required by the LruCache.
 impl MemoryDiskBytes for EntryKey {
     fn memory_bytes(&self) -> usize {
-        size_of_val(self) + self.method_name.len() + self.method_payload.len()
+        self.deterministic_heap_bytes()
     }
 
     fn disk_bytes(&self) -> usize {
@@ -163,7 +165,7 @@ impl From<&Query> for EntryKey {
 ///
 /// The cache entry is valid as long as the metadata is unchanged,
 /// or it can be proven that the query does not depend on the change.
-#[derive(PartialEq)]
+#[derive(DeterministicHeapBytes, PartialEq)]
 pub(crate) struct EntryEnv {
     /// The consensus-determined time when the query is executed.
     pub batch_time: Time,
@@ -196,6 +198,7 @@ impl EntryEnv {
 
 ////////////////////////////////////////////////////////////////////////
 /// Query Cache entry value.
+#[derive(DeterministicHeapBytes)]
 pub(crate) struct EntryValue {
     /// Query Cache entry environment metadata captured before the query execution.
     env: EntryEnv,
@@ -209,9 +212,10 @@ pub(crate) struct EntryValue {
     ignore_canister_balances: bool,
 }
 
+// The MemoryDiskBytes trait is still required by the LruCache.
 impl MemoryDiskBytes for EntryValue {
     fn memory_bytes(&self) -> usize {
-        size_of_val(self) + self.result.memory_bytes()
+        self.deterministic_heap_bytes()
     }
 
     fn disk_bytes(&self) -> usize {
@@ -367,28 +371,24 @@ impl EntryValue {
 
 ////////////////////////////////////////////////////////////////////////
 /// Replica Side Query Cache.
+#[derive(DeterministicHeapBytes)]
 pub(crate) struct QueryCache {
     // We can't use `RwLock`, as the `LruCache::get()` requires mutable reference
     // to update the LRU.
+    // For now, avoid using the derive macro on LruCache.
+    #[deterministic_heap_bytes(with =
+        |c: &Mutex<LruCache<EntryKey, EntryValue>>| c.lock().unwrap().memory_bytes())]
     cache: Mutex<LruCache<EntryKey, EntryValue>>,
     /// The upper limit on how long the cache entry stays valid in the query cache.
     max_expiry_time: Duration,
     /// The upper limit on how long the data certificate stays valid in the query cache.
     data_certificate_expiry_time: Duration,
-    /// Query cache metrics (public for tests)
+    /// Query cache metrics (public for tests).
+    // Assume the metrics do not use any heap allocations.
+    #[deterministic_heap_bytes(with = |_| 0)]
     pub(crate) metrics: QueryCacheMetrics,
 }
 
-impl MemoryDiskBytes for QueryCache {
-    fn memory_bytes(&self) -> usize {
-        size_of_val(self) + self.cache.lock().unwrap().memory_bytes()
-    }
-
-    fn disk_bytes(&self) -> usize {
-        0
-    }
-}
-
 impl QueryCache {
     /// Create a new `QueryCache` instance.
     pub(crate) fn new(