handle prefetch proofs correctly, more docs

crates/engine/tree/src/tree/root.rs

@@ -515,6 +515,16 @@ struct StateRootTaskMetrics {
     /// Histogram of sparse trie final update durations.
     pub sparse_trie_final_update_duration_histogram: Histogram,
 
+    /// Histogram of proof targets chunk sizes that were calculated from state updates.
+    pub state_update_chunks_histogram: Histogram,
+    /// Histogram of number of proof targets chunks that were calculated from state updates.
+    pub state_update_chunk_size_histogram: Histogram,
+
+    /// Histogram of proof targets chunk sizes that were calculated from prefetch targets.
+    pub prefetch_chunks_histogram: Histogram,
+    /// Histogram of number of proof targets chunks that were calculated from prefetch targets.
+    pub prefetch_chunk_size_histogram: Histogram,
+
     /// Histogram of state updates received.
     pub state_updates_received_histogram: Histogram,
     /// Histogram of proofs processed.
@@ -706,17 +716,23 @@ where
     }
 
     /// Handles request for proof prefetch.
-    fn on_prefetch_proof(&mut self, targets: MultiProofTargets) {
+    ///
+    /// Returns number of proofs generated by state update.
+    fn on_prefetch_proof(&mut self, targets: MultiProofTargets) -> u64 {
         let proof_targets = self.get_prefetch_proof_targets(targets);
         extend_multi_proof_targets_ref(&mut self.fetched_proof_targets, &proof_targets);
 
+        let mut total_proofs = 0;
         for chunk in ChunkedProofTargets::new(
             proof_targets,
             MULTIPROOF_ACCOUNTS_CHUNK_SIZE,
             MULTIPROOF_STORAGES_CHUNK_SIZE,
         )
         .flatten()
         {
+            total_proofs += 1;
+            self.metrics.prefetch_chunk_size_histogram.record(chunk.len() as f64);
+
             self.multiproof_manager.spawn_or_queue(MultiproofInput {
                 config: self.config.clone(),
                 source: None,
@@ -726,6 +742,9 @@ where
                 state_root_message_sender: self.tx.clone(),
             });
         }
+
+        self.metrics.prefetch_chunks_histogram.record(total_proofs as f64);
+        total_proofs as u64
     }
 
     /// Calls `get_proof_targets` with existing proof targets for prefetching.
@@ -786,15 +805,12 @@ where
     ///
     /// After chunking, [`MultiproofManager::spawn_or_queue`] is called for each state update.
     ///
-    /// Returns number of new updates generated by one state update.
+    /// Returns number of proofs generated by state update.
     fn on_state_update(&mut self, source: StateChangeSource, update: EvmState) -> u64 {
         let mut state_update = evm_state_to_hashed_post_state(update);
         let proof_targets = hashed_post_state_to_targets(&state_update);
 
-        debug!(target: "engine::root", ?state_update, "State update before");
-        debug!(target: "engine::root", ?proof_targets, "Proof targets");
-
-        let mut total_updates = 0;
+        let mut total_proofs = 0;
 
         for chunk in ChunkedProofTargets::new(
             proof_targets,
@@ -803,9 +819,8 @@ where
         )
         .flatten()
         {
-            debug!(target: "engine::root", ?chunk, "Chunk");
-
-            total_updates += 1;
+            total_proofs += 1;
+            self.metrics.state_update_chunk_size_histogram.record(chunk.len() as f64);
 
             let mut chunked_state_update = HashedPostState {
                 accounts: B256Map::with_capacity_and_hasher(chunk.len(), Default::default()),
@@ -885,7 +900,6 @@ where
                 chunked_state_update.storages.insert(address, chuncked_hashed_storage);
             }
 
-            debug!(target: "engine::tree", ?chunked_state_update, ?chunked_proof_targets, "Spawning multiproof");
             self.multiproof_manager.spawn_or_queue(MultiproofInput {
                 config: self.config.clone(),
                 source: Some(source),
@@ -896,9 +910,8 @@ where
             });
         }
 
-        debug!(target: "engine::tree", ?state_update, "State update after");
-
-        total_updates as u64
+        self.metrics.state_update_chunks_histogram.record(total_proofs as f64);
+        total_proofs as u64
     }
 
     /// Handler for new proof calculated, aggregates all the existing sequential proofs.
@@ -983,15 +996,20 @@ where
                 Ok(message) => match message {
                     StateRootMessage::PrefetchProofs(targets) => {
                         trace!(target: "engine::root", "processing StateRootMessage::PrefetchProofs");
-                        prefetch_proofs_received += 1;
+
+                        let targets_len = targets.len();
+                        let storage_targets_len =
+                            targets.values().map(|slots| slots.len()).sum::<usize>();
+                        let new_proofs = self.on_prefetch_proof(targets);
+                        prefetch_proofs_received += new_proofs;
+
                         debug!(
                             target: "engine::root",
-                            targets = targets.len(),
-                            storage_targets = targets.values().map(|slots| slots.len()).sum::<usize>(),
+                            targets_len,
+                            storage_targets_len,
                             total_prefetches = prefetch_proofs_received,
-                            "Prefetching proofs"
+                            "Received new prefetching request"
                         );
-                        self.on_prefetch_proof(targets);
                     }
                     StateRootMessage::StateUpdate(source, update) => {
                         trace!(target: "engine::root", "processing StateRootMessage::StateUpdate");
@@ -1002,12 +1020,13 @@ where
                         last_update_time = Some(Instant::now());
 
                         let update_size = update.len();
-                        let new_updates = self.on_state_update(source, update);
-                        updates_received += new_updates;
+                        let new_proofs = self.on_state_update(source, update);
+                        updates_received += new_proofs;
+
                         debug!(
                             target: "engine::root",
                             update_size,
-                            new_updates,
+                            new_proofs,
                             total_updates = updates_received,
                             "Received new state update"
                         );