diff --git a/crates/trie/parallel/src/proof_task.rs b/crates/trie/parallel/src/proof_task.rs index 6cf75aa818..a5914f1e66 100644 --- a/crates/trie/parallel/src/proof_task.rs +++ b/crates/trie/parallel/src/proof_task.rs @@ -1810,11 +1810,24 @@ fn dispatch_storage_proofs( fn dispatch_v2_storage_proofs( storage_work_tx: &CrossbeamSender, account_targets: &Vec, - storage_targets: B256Map>, + mut storage_targets: B256Map>, ) -> Result>, ParallelStateRootError> { let mut storage_proof_receivers = B256Map::with_capacity_and_hasher(account_targets.len(), Default::default()); + // Collect hashed addresses from account targets that need their storage roots computed + let account_target_addresses: B256Set = account_targets.iter().map(|t| t.key()).collect(); + + // For storage targets with associated account proofs, ensure the first target has + // min_len(0) so the root node is returned for storage root computation + for (hashed_address, targets) in &mut storage_targets { + if account_target_addresses.contains(hashed_address) && + let Some(first) = targets.first_mut() + { + *first = first.with_min_len(0); + } + } + // Dispatch all proofs for targeted storage slots for (hashed_address, targets) in storage_targets { // Create channel for receiving StorageProofResultMessage diff --git a/crates/trie/parallel/src/proof_task_metrics.rs b/crates/trie/parallel/src/proof_task_metrics.rs index e303df287b..579a78ae6b 100644 --- a/crates/trie/parallel/src/proof_task_metrics.rs +++ b/crates/trie/parallel/src/proof_task_metrics.rs @@ -26,6 +26,8 @@ pub struct ProofTaskTrieMetrics { deferred_encoder_from_cache: Histogram, /// Histogram for `Sync` deferred encoder variant count. deferred_encoder_sync: Histogram, + /// Histogram for dispatched storage proofs that fell back to sync due to missing root. + deferred_encoder_dispatched_missing_root: Histogram, } impl ProofTaskTrieMetrics { @@ -54,6 +56,8 @@ impl ProofTaskTrieMetrics { self.deferred_encoder_dispatched.record(stats.dispatched_count as f64); self.deferred_encoder_from_cache.record(stats.from_cache_count as f64); self.deferred_encoder_sync.record(stats.sync_count as f64); + self.deferred_encoder_dispatched_missing_root + .record(stats.dispatched_missing_root_count as f64); } } diff --git a/crates/trie/parallel/src/value_encoder.rs b/crates/trie/parallel/src/value_encoder.rs index 0b082a08d7..e2d046edee 100644 --- a/crates/trie/parallel/src/value_encoder.rs +++ b/crates/trie/parallel/src/value_encoder.rs @@ -32,6 +32,9 @@ pub(crate) struct ValueEncoderStats { pub(crate) from_cache_count: u64, /// Number of times the `Sync` variant was used (synchronous computation). pub(crate) sync_count: u64, + /// Number of times a dispatched storage proof had no root node and fell back to sync + /// computation. + pub(crate) dispatched_missing_root_count: u64, } impl ValueEncoderStats { @@ -41,6 +44,7 @@ impl ValueEncoderStats { self.dispatched_count += other.dispatched_count; self.from_cache_count += other.from_cache_count; self.sync_count += other.sync_count; + self.dispatched_missing_root_count += other.dispatched_missing_root_count; } } @@ -55,6 +59,11 @@ pub(crate) enum AsyncAccountDeferredValueEncoder { storage_proof_results: Rc>>>, /// Shared stats for tracking wait time and counts. stats: Rc>, + /// Shared storage proof calculator for synchronous fallback when dispatched proof has no + /// root. + storage_calculator: Rc>>, + /// Cache to store computed storage roots for future reuse. + cached_storage_roots: Arc>, }, /// The storage root was found in cache. FromCache { account: Account, root: B256 }, @@ -82,6 +91,8 @@ where proof_result_rx, storage_proof_results, stats, + storage_calculator, + cached_storage_roots, } => { let wait_start = Instant::now(); let result = proof_result_rx? @@ -94,12 +105,35 @@ where .result?; stats.borrow_mut().storage_wait_time += wait_start.elapsed(); - let StorageProofResult::V2 { root: Some(root), proof } = result else { - panic!("StorageProofResult is not V2 with root: {result:?}") + let StorageProofResult::V2 { root, proof } = result else { + panic!("StorageProofResult is not V2: {result:?}") }; storage_proof_results.borrow_mut().insert(hashed_address, proof); + let root = match root { + Some(root) => root, + None => { + // In `compute_v2_account_multiproof` we ensure that all dispatched storage + // proofs computations for which there is also an account proof will return + // a root node, but it could happen randomly that an account which is not in + // the account proof targets, but _is_ in storage proof targets, will need + // to be encoded as part of general trie traversal, so we need to handle + // that case here. + stats.borrow_mut().dispatched_missing_root_count += 1; + + let mut calculator = storage_calculator.borrow_mut(); + let proof = + calculator.storage_proof(hashed_address, &mut [B256::ZERO.into()])?; + let storage_root = calculator + .compute_root_hash(&proof)? + .expect("storage_proof with dummy target always returns root"); + + cached_storage_roots.insert(hashed_address, storage_root); + storage_root + } + }; + (account, root) } Self::FromCache { account, root } => (account, root), @@ -235,6 +269,8 @@ where proof_result_rx: Ok(rx), storage_proof_results: self.storage_proof_results.clone(), stats: self.stats.clone(), + storage_calculator: self.storage_calculator.clone(), + cached_storage_roots: self.cached_storage_roots.clone(), } }