diff --git a/crates/engine/tree/src/tree/payload_processor/configured_sparse_trie.rs b/crates/engine/tree/src/tree/payload_processor/configured_sparse_trie.rs index 90e8928dba..9e8f787823 100644 --- a/crates/engine/tree/src/tree/payload_processor/configured_sparse_trie.rs +++ b/crates/engine/tree/src/tree/payload_processor/configured_sparse_trie.rs @@ -186,4 +186,18 @@ impl SparseTrieInterface for ConfiguredSparseTrie { Self::Parallel(trie) => trie.value_capacity(), } } + + fn shrink_nodes_to(&mut self, size: usize) { + match self { + Self::Serial(trie) => trie.shrink_nodes_to(size), + Self::Parallel(trie) => trie.shrink_nodes_to(size), + } + } + + fn shrink_values_to(&mut self, size: usize) { + match self { + Self::Serial(trie) => trie.shrink_values_to(size), + Self::Parallel(trie) => trie.shrink_values_to(size), + } + } } diff --git a/crates/engine/tree/src/tree/payload_processor/mod.rs b/crates/engine/tree/src/tree/payload_processor/mod.rs index 8ab186dea5..bf3d7268ea 100644 --- a/crates/engine/tree/src/tree/payload_processor/mod.rs +++ b/crates/engine/tree/src/tree/payload_processor/mod.rs @@ -66,6 +66,29 @@ use configured_sparse_trie::ConfiguredSparseTrie; pub const PARALLEL_SPARSE_TRIE_PARALLELISM_THRESHOLDS: ParallelismThresholds = ParallelismThresholds { min_revealed_nodes: 100, min_updated_nodes: 100 }; +/// Default node capacity for shrinking the sparse trie. This is used to limit the number of trie +/// nodes in allocated sparse tries. +/// +/// Node maps have a key of `Nibbles` and value of `SparseNode`. +/// The `size_of::` is 40, and `size_of::` is 80. +/// +/// If we have 1 million entries of 120 bytes each, this conservative estimate comes out at around +/// 120MB. +pub const SPARSE_TRIE_MAX_NODES_SHRINK_CAPACITY: usize = 1_000_000; + +/// Default value capacity for shrinking the sparse trie. This is used to limit the number of values +/// in allocated sparse tries. +/// +/// There are storage and account values, the largest of the two being account values, which are +/// essentially `TrieAccount`s. +/// +/// Account value maps have a key of `Nibbles` and value of `TrieAccount`. +/// The `size_of::` is 40, and `size_of::` is 104. +/// +/// If we have 1 million entries of 144 bytes each, this conservative estimate comes out at around +/// 144MB. +pub const SPARSE_TRIE_MAX_VALUES_SHRINK_CAPACITY: usize = 1_000_000; + /// Entrypoint for executing the payload. #[derive(Debug)] pub struct PayloadProcessor @@ -439,11 +462,19 @@ where // Send state root computation result let _ = state_root_tx.send(result); - // Clear the SparseStateTrie and replace it back into the mutex _after_ sending + // Clear the SparseStateTrie, shrink, and replace it back into the mutex _after_ sending // results to the next step, so that time spent clearing doesn't block the step after // this one. let _enter = debug_span!(target: "engine::tree::payload_processor", "clear").entered(); - cleared_sparse_trie.lock().replace(ClearedSparseStateTrie::from_state_trie(trie)); + let mut cleared_trie = ClearedSparseStateTrie::from_state_trie(trie); + + // Shrink the sparse trie so that we don't have ever increasing memory. + cleared_trie.shrink_to( + SPARSE_TRIE_MAX_NODES_SHRINK_CAPACITY, + SPARSE_TRIE_MAX_VALUES_SHRINK_CAPACITY, + ); + + cleared_sparse_trie.lock().replace(cleared_trie); }); } } diff --git a/crates/trie/sparse-parallel/src/lower.rs b/crates/trie/sparse-parallel/src/lower.rs index b5454dd397..bc8ae00607 100644 --- a/crates/trie/sparse-parallel/src/lower.rs +++ b/crates/trie/sparse-parallel/src/lower.rs @@ -122,4 +122,26 @@ impl LowerSparseSubtrie { Self::Blind(None) => 0, } } + + /// Shrinks the capacity of the subtrie's node storage. + /// Works for both revealed and blind tries with allocated storage. + pub(crate) fn shrink_nodes_to(&mut self, size: usize) { + match self { + Self::Revealed(trie) | Self::Blind(Some(trie)) => { + trie.shrink_nodes_to(size); + } + Self::Blind(None) => {} + } + } + + /// Shrinks the capacity of the subtrie's value storage. + /// Works for both revealed and blind tries with allocated storage. + pub(crate) fn shrink_values_to(&mut self, size: usize) { + match self { + Self::Revealed(trie) | Self::Blind(Some(trie)) => { + trie.shrink_values_to(size); + } + Self::Blind(None) => {} + } + } } diff --git a/crates/trie/sparse-parallel/src/trie.rs b/crates/trie/sparse-parallel/src/trie.rs index 5e5a838f41..34c1ff2a96 100644 --- a/crates/trie/sparse-parallel/src/trie.rs +++ b/crates/trie/sparse-parallel/src/trie.rs @@ -883,6 +883,42 @@ impl SparseTrieInterface for ParallelSparseTrie { self.upper_subtrie.value_capacity() + self.lower_subtries.iter().map(|trie| trie.value_capacity()).sum::() } + + fn shrink_nodes_to(&mut self, size: usize) { + // Distribute the capacity across upper and lower subtries + // + // Always include upper subtrie, plus any lower subtries + let total_subtries = 1 + NUM_LOWER_SUBTRIES; + let size_per_subtrie = size / total_subtries; + + // Shrink the upper subtrie + self.upper_subtrie.shrink_nodes_to(size_per_subtrie); + + // Shrink lower subtries (works for both revealed and blind with allocation) + for subtrie in &mut self.lower_subtries { + subtrie.shrink_nodes_to(size_per_subtrie); + } + + // shrink masks maps + self.branch_node_hash_masks.shrink_to(size); + self.branch_node_tree_masks.shrink_to(size); + } + + fn shrink_values_to(&mut self, size: usize) { + // Distribute the capacity across upper and lower subtries + // + // Always include upper subtrie, plus any lower subtries + let total_subtries = 1 + NUM_LOWER_SUBTRIES; + let size_per_subtrie = size / total_subtries; + + // Shrink the upper subtrie + self.upper_subtrie.shrink_values_to(size_per_subtrie); + + // Shrink lower subtries (works for both revealed and blind with allocation) + for subtrie in &mut self.lower_subtries { + subtrie.shrink_values_to(size_per_subtrie); + } + } } impl ParallelSparseTrie { @@ -2111,6 +2147,16 @@ impl SparseSubtrie { pub(crate) fn value_capacity(&self) -> usize { self.inner.value_capacity() } + + /// Shrinks the capacity of the subtrie's node storage. + pub(crate) fn shrink_nodes_to(&mut self, size: usize) { + self.nodes.shrink_to(size); + } + + /// Shrinks the capacity of the subtrie's value storage. + pub(crate) fn shrink_values_to(&mut self, size: usize) { + self.inner.values.shrink_to(size); + } } /// Helper type for [`SparseSubtrie`] to mutably access only a subset of fields from the original @@ -2571,10 +2617,19 @@ impl SparseSubtrieBuffers { /// Clears all buffers. fn clear(&mut self) { self.path_stack.clear(); + self.path_stack.shrink_to_fit(); + self.rlp_node_stack.clear(); + self.rlp_node_stack.shrink_to_fit(); + self.branch_child_buf.clear(); + self.branch_child_buf.shrink_to_fit(); + self.branch_value_stack_buf.clear(); + self.branch_value_stack_buf.shrink_to_fit(); + self.rlp_buf.clear(); + self.rlp_buf.shrink_to_fit(); } } diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs index aef552da3d..a202ebc8b2 100644 --- a/crates/trie/sparse/src/state.rs +++ b/crates/trie/sparse/src/state.rs @@ -43,6 +43,32 @@ where Self(trie) } + /// Shrink the cleared sparse trie's capacity to the given node and value size. + /// This helps reduce memory usage when the trie has excess capacity. + /// The capacity is distributed equally across the account trie and all storage tries. + pub fn shrink_to(&mut self, node_size: usize, value_size: usize) { + // Count total number of storage tries (active + cleared + default) + let storage_tries_count = self.0.storage.tries.len() + self.0.storage.cleared_tries.len(); + + // Total tries = 1 account trie + all storage tries + let total_tries = 1 + storage_tries_count; + + // Distribute capacity equally among all tries + let node_size_per_trie = node_size / total_tries; + let value_size_per_trie = value_size / total_tries; + + // Shrink the account trie + self.0.state.shrink_nodes_to(node_size_per_trie); + self.0.state.shrink_values_to(value_size_per_trie); + + // Give storage tries the remaining capacity after account trie allocation + let storage_node_size = node_size.saturating_sub(node_size_per_trie); + let storage_value_size = value_size.saturating_sub(value_size_per_trie); + + // Shrink all storage tries (they will redistribute internally) + self.0.storage.shrink_to(storage_node_size, storage_value_size); + } + /// Returns the cleared [`SparseStateTrie`], consuming this instance. pub fn into_inner(self) -> SparseStateTrie { self.0 @@ -860,6 +886,31 @@ impl StorageTries { set })); } + + /// Shrinks the capacity of all storage tries (active, cleared, and default) to the given sizes. + /// The capacity is distributed equally among all tries that have allocations. + fn shrink_to(&mut self, node_size: usize, value_size: usize) { + // Count total number of tries with capacity (active + cleared + default) + let active_count = self.tries.len(); + let cleared_count = self.cleared_tries.len(); + let total_tries = 1 + active_count + cleared_count; + + // Distribute capacity equally among all tries + let node_size_per_trie = node_size / total_tries; + let value_size_per_trie = value_size / total_tries; + + // Shrink active storage tries + for trie in self.tries.values_mut() { + trie.shrink_nodes_to(node_size_per_trie); + trie.shrink_values_to(value_size_per_trie); + } + + // Shrink cleared storage tries + for trie in &mut self.cleared_tries { + trie.shrink_nodes_to(node_size_per_trie); + trie.shrink_values_to(value_size_per_trie); + } + } } impl StorageTries { diff --git a/crates/trie/sparse/src/traits.rs b/crates/trie/sparse/src/traits.rs index 8fdbb78d87..5b7b6193f9 100644 --- a/crates/trie/sparse/src/traits.rs +++ b/crates/trie/sparse/src/traits.rs @@ -228,6 +228,14 @@ pub trait SparseTrieInterface: Sized + Debug + Send + Sync { /// This returns the capacity of any inner data structures which store leaf values. fn value_capacity(&self) -> usize; + + /// Shrink the capacity of the sparse trie's node storage to the given size. + /// This will reduce memory usage if the current capacity is higher than the given size. + fn shrink_nodes_to(&mut self, size: usize); + + /// Shrink the capacity of the sparse trie's value storage to the given size. + /// This will reduce memory usage if the current capacity is higher than the given size. + fn shrink_values_to(&mut self, size: usize); } /// Struct for passing around branch node mask information. diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index 737da84225..8500ea400b 100644 --- a/crates/trie/sparse/src/trie.rs +++ b/crates/trie/sparse/src/trie.rs @@ -275,6 +275,28 @@ impl SparseTrie { _ => 0, } } + + /// Shrinks the capacity of the sparse trie's node storage. + /// Works for both revealed and blind tries with allocated storage. + pub fn shrink_nodes_to(&mut self, size: usize) { + match self { + Self::Blind(Some(trie)) | Self::Revealed(trie) => { + trie.shrink_nodes_to(size); + } + _ => {} + } + } + + /// Shrinks the capacity of the sparse trie's value storage. + /// Works for both revealed and blind tries with allocated storage. + pub fn shrink_values_to(&mut self, size: usize) { + match self { + Self::Blind(Some(trie)) | Self::Revealed(trie) => { + trie.shrink_values_to(size); + } + _ => {} + } + } } /// The representation of revealed sparse trie. @@ -1088,6 +1110,16 @@ impl SparseTrieInterface for SerialSparseTrie { fn value_capacity(&self) -> usize { self.values.capacity() } + + fn shrink_nodes_to(&mut self, size: usize) { + self.nodes.shrink_to(size); + self.branch_node_tree_masks.shrink_to(size); + self.branch_node_hash_masks.shrink_to(size); + } + + fn shrink_values_to(&mut self, size: usize) { + self.values.shrink_to(size); + } } impl SerialSparseTrie {