diff --git a/packages/ic-error-types/src/lib.rs b/packages/ic-error-types/src/lib.rs index 788924458f65..985271893394 100644 --- a/packages/ic-error-types/src/lib.rs +++ b/packages/ic-error-types/src/lib.rs @@ -334,8 +334,8 @@ impl UserError { } } - pub fn count_bytes(&self) -> usize { - std::mem::size_of_val(self) + self.description.len() + pub fn count_heap_bytes(&self) -> usize { + self.description.len() } /// Panics if the error doesn't have the expected code and description. diff --git a/rs/embedders/src/compilation_cache.rs b/rs/embedders/src/compilation_cache.rs index e4ffc8ac13fb..ed9cd14068ca 100644 --- a/rs/embedders/src/compilation_cache.rs +++ b/rs/embedders/src/compilation_cache.rs @@ -42,8 +42,8 @@ pub struct CompilationCache { } impl MemoryDiskBytes for CompilationCache { - fn memory_bytes(&self) -> usize { - self.cache.lock().unwrap().memory_bytes() + fn heap_bytes(&self) -> usize { + self.cache.lock().unwrap().heap_bytes() } fn disk_bytes(&self) -> usize { diff --git a/rs/embedders/src/serialized_module.rs b/rs/embedders/src/serialized_module.rs index 4a8f68c56e15..dab1968652f8 100644 --- a/rs/embedders/src/serialized_module.rs +++ b/rs/embedders/src/serialized_module.rs @@ -82,16 +82,6 @@ pub struct SerializedModule { pub is_wasm64: bool, } -impl MemoryDiskBytes for SerializedModule { - fn memory_bytes(&self) -> usize { - self.bytes.0.len() - } - - fn disk_bytes(&self) -> usize { - 0 - } -} - impl SerializedModule { pub(crate) fn new( module: &Module, @@ -164,10 +154,6 @@ pub struct OnDiskSerializedModule { } impl MemoryDiskBytes for OnDiskSerializedModule { - fn memory_bytes(&self) -> usize { - std::mem::size_of::() - } - fn disk_bytes(&self) -> usize { (self.bytes.metadata().unwrap().len() + self.initial_state_data.metadata().unwrap().len()) as usize diff --git a/rs/embedders/src/wasm_utils.rs b/rs/embedders/src/wasm_utils.rs index 83dc75e31a06..a17fec84a4ad 100644 --- a/rs/embedders/src/wasm_utils.rs +++ b/rs/embedders/src/wasm_utils.rs @@ -10,7 +10,7 @@ use ic_replicated_state::{ EmbedderCache, NumWasmPages, PageIndex, }; use ic_sys::{PageBytes, PAGE_SIZE}; -use ic_types::{methods::WasmMethod, NumBytes, NumInstructions}; +use ic_types::{methods::WasmMethod, MemoryDiskBytes, NumBytes, NumInstructions}; use ic_wasm_types::{BinaryEncodedWasm, WasmInstrumentationError}; use serde::{Deserialize, Serialize}; @@ -58,10 +58,23 @@ struct Segment { bytes: Vec, } +impl MemoryDiskBytes for Segment { + fn heap_bytes(&self) -> usize { + self.bytes.heap_bytes() + } +} + /// Vector of heap data chunks with their offsets. #[derive(Clone, Eq, PartialEq, Debug, Default, Deserialize, Serialize)] pub struct Segments(Vec); +impl MemoryDiskBytes for Segments { + fn heap_bytes(&self) -> usize { + // Deep vector size. + self.0.iter().map(|seg| seg.memory_bytes()).sum() + } +} + impl FromIterator<(usize, Vec)> for Segments { fn from_iter)>>(iter: T) -> Self { Segments( diff --git a/rs/execution_environment/src/query_handler/query_cache.rs b/rs/execution_environment/src/query_handler/query_cache.rs index 1edd3aaf643e..498ad61b9560 100644 --- a/rs/execution_environment/src/query_handler/query_cache.rs +++ b/rs/execution_environment/src/query_handler/query_cache.rs @@ -9,7 +9,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; @@ -138,12 +138,8 @@ pub(crate) struct EntryKey { } impl MemoryDiskBytes for EntryKey { - fn memory_bytes(&self) -> usize { - size_of_val(self) + self.method_name.len() + self.method_payload.len() - } - - fn disk_bytes(&self) -> usize { - 0 + fn heap_bytes(&self) -> usize { + self.method_name.heap_bytes() + self.method_payload.heap_bytes() } } @@ -171,6 +167,12 @@ pub(crate) struct EntryEnv { pub canisters_versions_balances_stats: Vec<(CanisterId, u64, Cycles, QueryStats)>, } +impl MemoryDiskBytes for EntryEnv { + fn heap_bytes(&self) -> usize { + self.canisters_versions_balances_stats.heap_bytes() + } +} + impl EntryEnv { // Capture a state (canister version and balance) of the evaluated canisters. fn try_new( @@ -210,12 +212,8 @@ pub(crate) struct EntryValue { } impl MemoryDiskBytes for EntryValue { - fn memory_bytes(&self) -> usize { - size_of_val(self) + self.result.memory_bytes() - } - - fn disk_bytes(&self) -> usize { - 0 + fn heap_bytes(&self) -> usize { + self.env.heap_bytes() + self.result.heap_bytes() } } @@ -380,12 +378,8 @@ pub(crate) struct QueryCache { } 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 + fn heap_bytes(&self) -> usize { + self.cache.lock().unwrap().heap_bytes() } } diff --git a/rs/execution_environment/src/query_handler/query_cache/tests.rs b/rs/execution_environment/src/query_handler/query_cache/tests.rs index 001e4a0494fc..05f1267d3e73 100644 --- a/rs/execution_environment/src/query_handler/query_cache/tests.rs +++ b/rs/execution_environment/src/query_handler/query_cache/tests.rs @@ -20,7 +20,11 @@ use ic_types::{ }; use ic_types_test_utils::ids::subnet_test_id; use ic_universal_canister::call_args; -use std::{collections::BTreeMap, sync::Arc, time::Duration}; +use std::{ + collections::{BTreeMap, BTreeSet}, + sync::Arc, + time::Duration, +}; const MAX_EXPIRY_TIME: Duration = Duration::from_secs(10); const MORE_THAN_MAX_EXPIRY_TIME: Duration = Duration::from_secs(11); @@ -771,7 +775,7 @@ fn query_cache_frees_memory_after_invalidated_entries() { let res = test .non_replicated_query(id, "canister_balance_sized_reply", vec![]) .unwrap(); - assert_eq!(BIG_RESPONSE_SIZE, res.memory_bytes()); + assert_eq!(size_of_val(&res) + BIG_RESPONSE_SIZE, res.memory_bytes()); let memory_bytes = query_cache(&test).memory_bytes(); // After the first reply, the cache should have more than 1MB of data. assert!(memory_bytes > BIG_RESPONSE_SIZE); @@ -786,7 +790,7 @@ fn query_cache_frees_memory_after_invalidated_entries() { let res = test .non_replicated_query(id, "canister_balance_sized_reply", vec![]) .unwrap(); - assert_eq!(SMALL_RESPONSE_SIZE, res.memory_bytes()); + assert_eq!(size_of_val(&res) + SMALL_RESPONSE_SIZE, res.memory_bytes()); let memory_bytes = query_cache(&test).memory_bytes(); // The second 42 reply should invalidate and replace the first 1MB reply in the cache. assert!(memory_bytes > SMALL_RESPONSE_SIZE); @@ -1563,3 +1567,148 @@ fn query_cache_future_proof_test() { } } } + +#[test] +fn memory_bytes_future_proof_guard() { + const HEAP_BYTES: usize = 2; + const MORE_HEAP_BYTES: usize = 3; + + // Key with no heap data. + let key = EntryKey { + source: user_test_id(1), + receiver: CanisterId::from_u64(1), + method_name: String::new(), + method_payload: vec![], + }; + assert_eq!(size_of_val(&key), 112); + assert_eq!(key.memory_bytes(), size_of_val(&key)); + + // Key with some heap data. + let key = EntryKey { + source: user_test_id(1), + receiver: CanisterId::from_u64(1), + method_name: " ".repeat(HEAP_BYTES), + method_payload: vec![42; HEAP_BYTES], + }; + assert_eq!(size_of_val(&key), 112); + assert_eq!(key.memory_bytes(), size_of_val(&key) + HEAP_BYTES * 2); + + // Key with more heap data. + let key = EntryKey { + source: user_test_id(1), + receiver: CanisterId::from_u64(1), + method_name: "0".repeat(MORE_HEAP_BYTES), + method_payload: vec![42; MORE_HEAP_BYTES], + }; + assert_eq!(size_of_val(&key), 112); + assert_eq!(key.memory_bytes(), size_of_val(&key) + MORE_HEAP_BYTES * 2); + + // Value with no heap data. + let env = EntryEnv { + batch_time: time::GENESIS, + canisters_versions_balances_stats: vec![], + }; + let value = EntryValue::new( + env, + Result::Ok(WasmResult::Reply(vec![])), + &SystemApiCallCounters::default(), + ); + assert_eq!(size_of_val(&value), 80); + assert_eq!(value.memory_bytes(), size_of_val(&value)); + + // Value with some heap data. + let env = EntryEnv { + batch_time: time::GENESIS, + canisters_versions_balances_stats: vec![ + ( + CanisterId::from_u64(1), + 0, + 0_u64.into(), + QueryStats::default(), + ); + HEAP_BYTES + ], + }; + let env_vec_size = size_of_val(&*env.canisters_versions_balances_stats); + let value = EntryValue::new( + env, + Result::Ok(WasmResult::Reply(vec![42; HEAP_BYTES])), + &SystemApiCallCounters::default(), + ); + assert_eq!(size_of_val(&value), 80); + assert_eq!( + value.memory_bytes(), + size_of_val(&value) + env_vec_size + HEAP_BYTES + ); + + // Value with more heap data. + let env = EntryEnv { + batch_time: time::GENESIS, + canisters_versions_balances_stats: vec![ + ( + CanisterId::from_u64(1), + 0, + 0_u64.into(), + QueryStats::default(), + ); + MORE_HEAP_BYTES + ], + }; + let env_vec_size = size_of_val(&*env.canisters_versions_balances_stats); + let value = EntryValue::new( + env, + Result::Ok(WasmResult::Reply(vec![42; MORE_HEAP_BYTES])), + &SystemApiCallCounters::default(), + ); + assert_eq!(size_of_val(&value), 80); + assert_eq!( + value.memory_bytes(), + size_of_val(&value) + env_vec_size + MORE_HEAP_BYTES + ); +} + +#[test] +fn btree_set_memory_bytes_future_proof_guard() { + let mut set = BTreeSet::new(); + assert_eq!(set.heap_bytes(), 0); + assert_eq!(set.memory_bytes(), size_of::>()); + set.insert(42_u64); + assert_eq!(set.heap_bytes(), size_of::()); + assert_eq!( + set.memory_bytes(), + size_of::>() + size_of::() + ); + + let mut set = BTreeSet::new(); + set.insert(" ".repeat(100)); + set.insert(" ".repeat(1000)); + // By default, the `heap_bytes` returns a constant time estimation. + assert_eq!(set.heap_bytes(), size_of::() * 2); + assert_eq!( + set.memory_bytes(), + size_of::>() + size_of::() * 2 + ); +} + +#[test] +fn vec_memory_bytes_future_proof_guard() { + let mut vec = vec![]; + assert_eq!(vec.heap_bytes(), 0); + assert_eq!(vec.memory_bytes(), size_of::>()); + vec.push(42_u64); + assert_eq!(vec.heap_bytes(), size_of::()); + assert_eq!( + vec.memory_bytes(), + size_of::>() + size_of::() + ); + + let mut vec = BTreeSet::new(); + vec.insert(" ".repeat(100)); + vec.insert(" ".repeat(1000)); + // By default, the `heap_bytes` returns a constant time estimation. + assert_eq!(vec.heap_bytes(), size_of::() * 2); + assert_eq!( + vec.memory_bytes(), + size_of::>() + size_of::() * 2 + ); +} diff --git a/rs/interfaces/src/execution_environment/errors.rs b/rs/interfaces/src/execution_environment/errors.rs index 5c4e4efbd2ca..9abe917da891 100644 --- a/rs/interfaces/src/execution_environment/errors.rs +++ b/rs/interfaces/src/execution_environment/errors.rs @@ -55,6 +55,9 @@ pub struct CanisterOutOfCyclesError { pub reveal_top_up: bool, } +// No heap allocations. +impl MemoryDiskBytes for CanisterOutOfCyclesError {} + impl std::error::Error for CanisterOutOfCyclesError {} impl std::fmt::Display for CanisterOutOfCyclesError { @@ -74,6 +77,13 @@ impl std::fmt::Display for CanisterOutOfCyclesError { #[derive(Clone, Eq, PartialEq, Debug, Deserialize, Serialize)] pub struct CanisterBacktrace(pub Vec<(u32, Option)>); +impl MemoryDiskBytes for CanisterBacktrace { + fn heap_bytes(&self) -> usize { + // Deep vector size. + self.0.iter().map(|(_, o)| o.memory_bytes()).sum() + } +} + impl std::fmt::Display for CanisterBacktrace { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { writeln!(f, "Canister Backtrace:")?; @@ -197,6 +207,57 @@ pub enum HypervisorError { }, } +impl MemoryDiskBytes for HypervisorError { + fn heap_bytes(&self) -> usize { + match self { + HypervisorError::FunctionNotFound(..) => 0, + HypervisorError::MethodNotFound(wasm_method) => wasm_method.heap_bytes(), + HypervisorError::ToolchainContractViolation { error } => error.heap_bytes(), + HypervisorError::UserContractViolation { + error, + suggestion, + doc_link, + } => error.heap_bytes() + suggestion.heap_bytes() + doc_link.heap_bytes(), + HypervisorError::InstructionLimitExceeded(_) => 0, + HypervisorError::InvalidWasm(wasm_validation_error) => { + wasm_validation_error.heap_bytes() + } + HypervisorError::InstrumentationFailed(wasm_instrumentation_error) => { + wasm_instrumentation_error.heap_bytes() + } + HypervisorError::Trapped { backtrace, .. } => backtrace.heap_bytes(), + HypervisorError::CalledTrap { message, backtrace } => { + message.heap_bytes() + backtrace.heap_bytes() + } + HypervisorError::WasmModuleNotFound => 0, + HypervisorError::OutOfMemory => 0, + HypervisorError::InvalidPrincipalId(_principal_id_error) => 0, + HypervisorError::MessageRejected => 0, + HypervisorError::InsufficientCyclesBalance(canister_out_of_cycles_error) => { + canister_out_of_cycles_error.heap_bytes() + } + HypervisorError::Cleanup { + callback_err, + cleanup_err, + } => { + // Deep size for `Box`ed errors. + callback_err.memory_bytes() + cleanup_err.memory_bytes() + } + HypervisorError::WasmEngineError(wasm_engine_error) => wasm_engine_error.heap_bytes(), + HypervisorError::ReservedPagesForOldMotoko => 0, + HypervisorError::Aborted => 0, + HypervisorError::SliceOverrun { .. } => 0, + HypervisorError::MemoryAccessLimitExceeded(s) => s.heap_bytes(), + HypervisorError::InsufficientCyclesInMemoryGrow { .. } => 0, + HypervisorError::ReservedCyclesLimitExceededInMemoryGrow { .. } => 0, + HypervisorError::InsufficientCyclesInMessageMemoryGrow { .. } => 0, + HypervisorError::WasmMemoryLimitExceeded { .. } => 0, + HypervisorError::EnvironmentVariableIndexOutOfBounds { .. } => 0, + HypervisorError::EnvironmentVariableNotFound { name } => name.heap_bytes(), + } + } +} + impl From for HypervisorError { fn from(err: WasmInstrumentationError) -> Self { Self::InstrumentationFailed(err) @@ -404,16 +465,6 @@ impl std::fmt::Display for HypervisorError { } } -impl MemoryDiskBytes for HypervisorError { - fn memory_bytes(&self) -> usize { - std::mem::size_of::() - } - - fn disk_bytes(&self) -> usize { - 0 - } -} - impl AsErrorHelp for HypervisorError { fn error_help(&self) -> ErrorHelp { match self { diff --git a/rs/types/types/src/ingress.rs b/rs/types/types/src/ingress.rs index 02b88bd81cfc..9410c32205c0 100644 --- a/rs/types/types/src/ingress.rs +++ b/rs/types/types/src/ingress.rs @@ -108,7 +108,7 @@ impl IngressStatus { match self { IngressStatus::Known { state, .. } => match state { IngressState::Completed(result) => result.memory_bytes(), - IngressState::Failed(error) => error.description().len(), + IngressState::Failed(error) => error.memory_bytes(), _ => 0, }, IngressStatus::Unknown => 0, @@ -177,16 +177,12 @@ pub enum WasmResult { } impl MemoryDiskBytes for WasmResult { - fn memory_bytes(&self) -> usize { + fn heap_bytes(&self) -> usize { match self { - WasmResult::Reply(bytes) => bytes.len(), - WasmResult::Reject(string) => string.len(), + WasmResult::Reply(bytes) => bytes.heap_bytes(), + WasmResult::Reject(string) => string.heap_bytes(), } } - - fn disk_bytes(&self) -> usize { - 0 - } } impl WasmResult { diff --git a/rs/types/types/src/lib.rs b/rs/types/types/src/lib.rs index bc097b17d095..b58d77031d03 100644 --- a/rs/types/types/src/lib.rs +++ b/rs/types/types/src/lib.rs @@ -106,6 +106,7 @@ use ic_protobuf::state::canister_state_bits::v1 as pb_state_bits; use ic_protobuf::types::v1 as pb; use phantom_newtype::{AmountOf, DisplayerOf, Id}; use serde::{Deserialize, Serialize}; +use std::collections::BTreeSet; use std::convert::TryFrom; use std::fmt; use std::sync::Arc; @@ -584,25 +585,65 @@ pub trait CountBytes { /// Allow an object to report its own byte size on disk and in memory. Not /// necessarily exact. pub trait MemoryDiskBytes { - fn memory_bytes(&self) -> usize; - fn disk_bytes(&self) -> usize; -} + /// Returns the number of bytes that this object occupies in heap memory. + fn heap_bytes(&self) -> usize { + // No heap allocations by default. + 0 + } -impl MemoryDiskBytes for Time { + /// Returns the total number of bytes that this object occupies + /// in stack and heap memory. fn memory_bytes(&self) -> usize { - 8 + size_of_val(self) + self.heap_bytes() } + /// Returns the number of bytes that this object occupies on disk. fn disk_bytes(&self) -> usize { 0 } } +impl MemoryDiskBytes for BTreeSet { + /// Provides a shallow constant time size estimation by default. + fn heap_bytes(&self) -> usize { + self.len() * size_of::() + } +} + +impl MemoryDiskBytes for Vec { + /// Provides a shallow constant time size estimation by default. + fn heap_bytes(&self) -> usize { + size_of_val(&**self) + } +} + +impl MemoryDiskBytes for String { + fn heap_bytes(&self) -> usize { + self.len() + } +} + +impl MemoryDiskBytes for Option { + fn heap_bytes(&self) -> usize { + match self { + Some(t) => t.heap_bytes(), + None => 0, + } + } + + fn disk_bytes(&self) -> usize { + match self { + Some(t) => t.disk_bytes(), + None => 0, + } + } +} + impl MemoryDiskBytes for Result { - fn memory_bytes(&self) -> usize { + fn heap_bytes(&self) -> usize { match self { - Ok(result) => result.memory_bytes(), - Err(err) => err.memory_bytes(), + Ok(result) => result.heap_bytes(), + Err(err) => err.heap_bytes(), } } @@ -615,8 +656,8 @@ impl MemoryDiskBytes for Result { } impl MemoryDiskBytes for Arc { - fn memory_bytes(&self) -> usize { - self.as_ref().memory_bytes() + fn heap_bytes(&self) -> usize { + self.as_ref().heap_bytes() } fn disk_bytes(&self) -> usize { @@ -626,11 +667,7 @@ impl MemoryDiskBytes for Arc { // Implementing `MemoryDiskBytes` in `ic_error_types` introduces a circular dependency. impl MemoryDiskBytes for ic_error_types::UserError { - fn memory_bytes(&self) -> usize { - self.count_bytes() - } - - fn disk_bytes(&self) -> usize { - 0 + fn heap_bytes(&self) -> usize { + self.count_heap_bytes() } } diff --git a/rs/types/types/src/methods.rs b/rs/types/types/src/methods.rs index 2983e356a832..e471e2b591f5 100644 --- a/rs/types/types/src/methods.rs +++ b/rs/types/types/src/methods.rs @@ -1,6 +1,7 @@ //! This module contains a collection of types and structs that define the //! various types of methods in the IC. +use crate::MemoryDiskBytes; use crate::{messages::CallContextId, time::CoarseTime, Cycles}; use ic_base_types::{CanisterId, PrincipalId}; #[cfg(test)] @@ -40,6 +41,17 @@ pub enum WasmMethod { System(SystemMethod), } +impl MemoryDiskBytes for WasmMethod { + fn heap_bytes(&self) -> usize { + match self { + WasmMethod::Update(s) | WasmMethod::Query(s) | WasmMethod::CompositeQuery(s) => { + s.heap_bytes() + } + WasmMethod::System(_system_method) => 0, + } + } +} + impl WasmMethod { pub fn name(&self) -> String { match self { diff --git a/rs/types/wasm_types/src/errors.rs b/rs/types/wasm_types/src/errors.rs index bdf423055726..f179e3ed546a 100644 --- a/rs/types/wasm_types/src/errors.rs +++ b/rs/types/wasm_types/src/errors.rs @@ -1,3 +1,4 @@ +use ic_types::MemoryDiskBytes; use serde::{Deserialize, Serialize}; /// Create a link to this section of the Execution Errors documentation. @@ -46,6 +47,12 @@ pub trait AsErrorHelp { #[derive(Clone, Eq, PartialEq, Debug, Deserialize, Serialize)] pub struct WasmError(String); +impl MemoryDiskBytes for WasmError { + fn heap_bytes(&self) -> usize { + self.0.heap_bytes() + } +} + impl WasmError { /// Creates a new `WasmError` out of an error message. pub fn new(error_message: String) -> Self { @@ -115,6 +122,35 @@ pub enum WasmValidationError { }, } +impl MemoryDiskBytes for WasmValidationError { + fn heap_bytes(&self) -> usize { + match self { + WasmValidationError::WasmtimeValidation(s) + | WasmValidationError::DecodingError(s) + | WasmValidationError::InvalidFunctionSignature(s) + | WasmValidationError::InvalidImportSection(s) + | WasmValidationError::InvalidExportSection(s) => s.heap_bytes(), + WasmValidationError::DuplicateExport { name } => name.heap_bytes(), + WasmValidationError::TooManyExports { .. } => 0, + WasmValidationError::ExportedNamesTooLong { .. } => 0, + WasmValidationError::InvalidDataSection(s) + | WasmValidationError::InvalidCustomSection(s) + | WasmValidationError::InvalidGlobalSection(s) => s.heap_bytes(), + WasmValidationError::TooManyGlobals { .. } => 0, + WasmValidationError::TooManyFunctions { .. } => 0, + WasmValidationError::TooManyCustomSections { .. } => 0, + WasmValidationError::FunctionComplexityTooHigh { .. } => 0, + WasmValidationError::UnsupportedWasmInstruction { instruction, .. } => { + instruction.heap_bytes() + } + WasmValidationError::FunctionTooLarge { .. } => 0, + WasmValidationError::CodeSectionTooLarge { .. } => 0, + WasmValidationError::ModuleTooLarge { .. } => 0, + WasmValidationError::InitialWasm64MemoryTooLarge { .. } => 0, + } + } +} + impl std::fmt::Display for WasmValidationError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { @@ -319,6 +355,18 @@ pub enum WasmInstrumentationError { InvalidFunctionType(String), } +impl MemoryDiskBytes for WasmInstrumentationError { + fn heap_bytes(&self) -> usize { + match self { + WasmInstrumentationError::WasmDeserializeError(wasm_error) + | WasmInstrumentationError::WasmSerializeError(wasm_error) => wasm_error.heap_bytes(), + WasmInstrumentationError::IncorrectNumberMemorySections { .. } => 0, + WasmInstrumentationError::InvalidDataSegment { .. } => 0, + WasmInstrumentationError::InvalidFunctionType(s) => s.heap_bytes(), + } + } +} + impl std::fmt::Display for WasmInstrumentationError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { @@ -369,6 +417,21 @@ pub enum WasmEngineError { Unexpected(String), } +impl MemoryDiskBytes for WasmEngineError { + fn heap_bytes(&self) -> usize { + match self { + WasmEngineError::FailedToInitializeEngine => 0, + WasmEngineError::FailedToInstantiateModule(s) => s.heap_bytes(), + WasmEngineError::FailedToSetAsyncStack | WasmEngineError::FailedToSetWasmStack => 0, + WasmEngineError::FailedToSerializeModule(s) + | WasmEngineError::FailedToDeserializeModule(s) + | WasmEngineError::FailedToApplySystemChanges(s) + | WasmEngineError::Other(s) + | WasmEngineError::Unexpected(s) => s.heap_bytes(), + } + } +} + impl std::fmt::Display for WasmEngineError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { diff --git a/rs/types/wasm_types/src/lib.rs b/rs/types/wasm_types/src/lib.rs index 85bed649897a..2971914f3208 100644 --- a/rs/types/wasm_types/src/lib.rs +++ b/rs/types/wasm_types/src/lib.rs @@ -188,6 +188,9 @@ impl std::hash::Hash for CanisterModule { #[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug)] pub struct WasmHash([u8; WASM_HASH_LENGTH]); +// No heap allocations. +impl MemoryDiskBytes for WasmHash {} + impl WasmHash { pub fn to_slice(&self) -> [u8; WASM_HASH_LENGTH] { self.0 @@ -219,16 +222,6 @@ impl TryFrom> for WasmHash { } } -impl MemoryDiskBytes for WasmHash { - fn memory_bytes(&self) -> usize { - self.0.len() - } - - fn disk_bytes(&self) -> usize { - 0 - } -} - impl std::fmt::Display for WasmHash { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> { for byte in self.0 { diff --git a/rs/utils/lru_cache/src/lib.rs b/rs/utils/lru_cache/src/lib.rs index 5056164d3b9e..421e2780f204 100644 --- a/rs/utils/lru_cache/src/lib.rs +++ b/rs/utils/lru_cache/src/lib.rs @@ -26,7 +26,7 @@ where K: MemoryDiskBytes + Eq + Hash, V: MemoryDiskBytes, { - fn memory_bytes(&self) -> usize { + fn heap_bytes(&self) -> usize { self.memory_size } @@ -442,27 +442,28 @@ mod tests { #[test] fn lru_cache_count_bytes_and_len() { let mut lru = LruCache::::new(NumBytes::new(10), NumBytes::new(20)); - assert_eq!(0, lru.memory_bytes()); + assert_eq!(size_of_val(&lru), lru.memory_bytes()); + assert_eq!(0, lru.heap_bytes()); assert_eq!(0, lru.disk_bytes()); assert_eq!(0, lru.len()); assert!(lru.is_empty()); lru.push(Key(0), MemoryDiskValue(0, 4, 10)); - assert_eq!(4, lru.memory_bytes()); + assert_eq!(4, lru.heap_bytes()); assert_eq!(10, lru.disk_bytes()); assert_eq!(1, lru.len()); assert!(!lru.is_empty()); lru.push(Key(1), MemoryDiskValue(1, 6, 2)); - assert_eq!(10, lru.memory_bytes()); + assert_eq!(10, lru.heap_bytes()); assert_eq!(12, lru.disk_bytes()); assert_eq!(2, lru.len()); assert!(!lru.is_empty()); lru.pop(&Key(0)); - assert_eq!(6, lru.memory_bytes()); + assert_eq!(6, lru.heap_bytes()); assert_eq!(2, lru.disk_bytes()); assert_eq!(1, lru.len()); assert!(!lru.is_empty()); lru.pop(&Key(1)); - assert_eq!(0, lru.memory_bytes()); + assert_eq!(0, lru.heap_bytes()); assert_eq!(0, lru.disk_bytes()); assert_eq!(0, lru.len()); assert!(lru.is_empty()); @@ -611,7 +612,7 @@ mod tests { update(&mut evicted_memory, &mut evicted_disk, &k, &v); } - assert_eq!(total_memory, evicted_memory + lru.memory_bytes()); + assert_eq!(size_of_val(&lru) + total_memory, evicted_memory + lru.memory_bytes()); assert_eq!(total_disk, evicted_disk + lru.disk_bytes()); } }