perf(tree): only chunk multiproof targets if needed (#19326)

Author: shekhirin

crates/engine/tree/src/tree/payload_processor/multiproof.rs

@@ -743,7 +743,12 @@ impl MultiProofTask {
     /// Handles request for proof prefetch.
     ///
     /// Returns a number of proofs that were spawned.
-    #[instrument(level = "debug", target = "engine::tree::payload_processor::multiproof", skip_all, fields(accounts = targets.len()))]
+    #[instrument(
+        level = "debug",
+        target = "engine::tree::payload_processor::multiproof",
+        skip_all,
+        fields(accounts = targets.len(), chunks = 0)
+    )]
     fn on_prefetch_proof(&mut self, targets: MultiProofTargets) -> u64 {
         let proof_targets = self.get_prefetch_proof_targets(targets);
         self.fetched_proof_targets.extend_ref(&proof_targets);
@@ -785,10 +790,16 @@ impl MultiProofTask {
             chunks += 1;
         };
 
-        if should_chunk && let Some(chunk_size) = self.chunk_size {
+        if should_chunk &&
+            let Some(chunk_size) = self.chunk_size &&
+            proof_targets.chunking_length() > chunk_size
+        {
+            let mut chunks = 0usize;
             for proof_targets_chunk in proof_targets.chunks(chunk_size) {
                 dispatch(proof_targets_chunk);
+                chunks += 1;
             }
+            tracing::Span::current().record("chunks", chunks);
         } else {
             dispatch(proof_targets);
         }
@@ -874,7 +885,12 @@ impl MultiProofTask {
     /// Handles state updates.
     ///
     /// Returns a number of proofs that were spawned.
-    #[instrument(level = "debug", target = "engine::tree::payload_processor::multiproof", skip(self, update), fields(accounts = update.len()))]
+    #[instrument(
+        level = "debug",
+        target = "engine::tree::payload_processor::multiproof",
+        skip(self, update),
+        fields(accounts = update.len(), chunks = 0)
+    )]
     fn on_state_update(&mut self, source: StateChangeSource, update: EvmState) -> u64 {
         let hashed_state_update = evm_state_to_hashed_post_state(update);
 
@@ -934,10 +950,16 @@ impl MultiProofTask {
             chunks += 1;
         };
 
-        if should_chunk && let Some(chunk_size) = self.chunk_size {
+        if should_chunk &&
+            let Some(chunk_size) = self.chunk_size &&
+            not_fetched_state_update.chunking_length() > chunk_size
+        {
+            let mut chunks = 0usize;
             for chunk in not_fetched_state_update.chunks(chunk_size) {
                 dispatch(chunk);
+                chunks += 1;
             }
+            tracing::Span::current().record("chunks", chunks);
         } else {
             dispatch(not_fetched_state_update);
         }

crates/trie/common/src/hashed_state.rs

@@ -278,6 +278,15 @@ impl HashedPostState {
         ChunkedHashedPostState::new(self, size)
     }
 
+    /// Returns the number of items that will be considered during chunking in `[Self::chunks]`.
+    pub fn chunking_length(&self) -> usize {
+        self.accounts.len() +
+            self.storages
+                .values()
+                .map(|storage| if storage.wiped { 1 } else { 0 } + storage.storage.len())
+                .sum::<usize>()
+    }
+
     /// Extend this hashed post state with contents of another.
     /// Entries in the second hashed post state take precedence.
     pub fn extend(&mut self, other: Self) {
@@ -1246,4 +1255,65 @@ mod tests {
         assert_eq!(storage3.zero_valued_slots.len(), 1);
         assert!(storage3.zero_valued_slots.contains(&B256::from([4; 32])));
     }
+
+    #[test]
+    fn test_hashed_post_state_chunking_length() {
+        let addr1 = B256::from([1; 32]);
+        let addr2 = B256::from([2; 32]);
+        let addr3 = B256::from([3; 32]);
+        let addr4 = B256::from([4; 32]);
+        let slot1 = B256::from([1; 32]);
+        let slot2 = B256::from([2; 32]);
+        let slot3 = B256::from([3; 32]);
+
+        let state = HashedPostState {
+            accounts: B256Map::from_iter([(addr1, None), (addr2, None), (addr4, None)]),
+            storages: B256Map::from_iter([
+                (
+                    addr1,
+                    HashedStorage {
+                        wiped: false,
+                        storage: B256Map::from_iter([
+                            (slot1, U256::ZERO),
+                            (slot2, U256::ZERO),
+                            (slot3, U256::ZERO),
+                        ]),
+                    },
+                ),
+                (
+                    addr2,
+                    HashedStorage {
+                        wiped: true,
+                        storage: B256Map::from_iter([
+                            (slot1, U256::ZERO),
+                            (slot2, U256::ZERO),
+                            (slot3, U256::ZERO),
+                        ]),
+                    },
+                ),
+                (
+                    addr3,
+                    HashedStorage {
+                        wiped: false,
+                        storage: B256Map::from_iter([
+                            (slot1, U256::ZERO),
+                            (slot2, U256::ZERO),
+                            (slot3, U256::ZERO),
+                        ]),
+                    },
+                ),
+            ]),
+        };
+
+        let chunking_length = state.chunking_length();
+        for size in 1..=state.clone().chunks(1).count() {
+            let chunk_count = state.clone().chunks(size).count();
+            let expected_count = chunking_length.div_ceil(size);
+            assert_eq!(
+                chunk_count, expected_count,
+                "chunking_length: {}, size: {}",
+                chunking_length, size
+            );
+        }
+    }
 }

crates/trie/common/src/proofs.rs

@@ -89,6 +89,11 @@ impl MultiProofTargets {
     pub fn chunks(self, size: usize) -> ChunkedMultiProofTargets {
         ChunkedMultiProofTargets::new(self, size)
     }
+
+    /// Returns the number of items that will be considered during chunking in `[Self::chunks]`.
+    pub fn chunking_length(&self) -> usize {
+        self.values().map(|slots| 1 + slots.len().saturating_sub(1)).sum::<usize>()
+    }
 }
 
 /// An iterator that yields chunks of the proof targets of at most `size` account and storage
@@ -1067,4 +1072,33 @@ mod tests {
         acc.storage_root = EMPTY_ROOT_HASH;
         assert_eq!(acc, inverse);
     }
+
+    #[test]
+    fn test_multiproof_targets_chunking_length() {
+        let mut targets = MultiProofTargets::default();
+        targets.insert(B256::with_last_byte(1), B256Set::default());
+        targets.insert(
+            B256::with_last_byte(2),
+            B256Set::from_iter([B256::with_last_byte(10), B256::with_last_byte(20)]),
+        );
+        targets.insert(
+            B256::with_last_byte(3),
+            B256Set::from_iter([
+                B256::with_last_byte(30),
+                B256::with_last_byte(31),
+                B256::with_last_byte(32),
+            ]),
+        );
+
+        let chunking_length = targets.chunking_length();
+        for size in 1..=targets.clone().chunks(1).count() {
+            let chunk_count = targets.clone().chunks(size).count();
+            let expected_count = chunking_length.div_ceil(size);
+            assert_eq!(
+                chunk_count, expected_count,
+                "chunking_length: {}, size: {}",
+                chunking_length, size
+            );
+        }
+    }
 }