address comments

near · Dec 19, 2024 · 7f8bb8f · 7f8bb8f
1 parent bf64b82
commit 7f8bb8f
Show file tree

Hide file tree

Showing 6 changed files with 35 additions and 33 deletions.
diff --git a/chain/client/src/metrics.rs b/chain/client/src/metrics.rs
@@ -638,15 +638,16 @@ pub(crate) static BLOCK_PRODUCER_MISSING_ENDORSEMENT_COUNT: LazyLock<HistogramVe
         .unwrap()
     });
 
-pub(crate) static PARTIAL_WITNESS_ENCODE_TIME: LazyLock<HistogramVec> = LazyLock::new(|| {
-    try_create_histogram_vec(
-        "near_partial_witness_encode_time",
-        "Partial state witness generation from encoded state witness time in seconds",
-        &["shard_id"],
-        Some(linear_buckets(0.0, 0.005, 20).unwrap()),
-    )
-    .unwrap()
-});
+pub(crate) static PARTIAL_WITNESS_ENCODE_AND_SEND_TIME: LazyLock<HistogramVec> =
+    LazyLock::new(|| {
+        try_create_histogram_vec(
+            "near_partial_witness_encode_and_send_time",
+            "Partial state witness generation from encoded state witness time in seconds",
+            &["shard_id"],
+            Some(linear_buckets(0.0, 0.005, 20).unwrap()),
+        )
+        .unwrap()
+    });
 
 pub(crate) static PARTIAL_WITNESS_TIME_TO_LAST_PART: LazyLock<HistogramVec> = LazyLock::new(|| {
     try_create_histogram_vec(

diff --git a/chain/client/src/stateless_validation/partial_witness/partial_witness_actor.rs b/chain/client/src/stateless_validation/partial_witness/partial_witness_actor.rs
@@ -268,26 +268,27 @@ impl PartialWitnessActor {
         let encoder = self.witness_encoders.entry(chunk_validators.len());
         let (parts, encoded_length) = encoder.encode(&witness_bytes);
 
-        chunk_validators.iter().zip_eq(parts).enumerate().for_each(
-            |(part_ord, (chunk_validator, part))| {
-                // It's fine to unwrap part here as we just constructed the parts above and we expect
-                // all of them to be present.
-                let partial_witness = PartialEncodedStateWitness::new(
-                    epoch_id,
-                    chunk_header.clone(),
-                    part_ord,
-                    part.unwrap().to_vec(),
-                    encoded_length,
-                    signer,
-                );
-                let validator_witness_tuple = (chunk_validator.clone(), partial_witness);
+        for (part_ord, (chunk_validator, part)) in chunk_validators.iter().zip_eq(parts).enumerate()
+        {
+            // It's fine to unwrap part here as we just constructed the parts above and we expect
+            // all of them to be present.
+            let partial_witness = PartialEncodedStateWitness::new(
+                epoch_id,
+                chunk_header.clone(),
+                part_ord,
+                part.unwrap().to_vec(),
+                encoded_length,
+                signer,
+            );
 
-                // Send the parts to the corresponding chunk validator owners.
-                self.network_adapter.send(PeerManagerMessageRequest::NetworkRequests(
-                    NetworkRequests::PartialEncodedStateWitness(validator_witness_tuple),
-                ));
-            },
-        );
+            // Send the parts to the corresponding chunk validator owners.
+            self.network_adapter.send(PeerManagerMessageRequest::NetworkRequests(
+                NetworkRequests::PartialEncodedStateWitness(
+                    chunk_validator.clone(),
+                    partial_witness,
+                ),
+            ));
+        }
 
         Ok(())
     }
@@ -344,7 +345,7 @@ impl PartialWitnessActor {
 
         // Record time taken to encode the state witness parts.
         let shard_id_label = chunk_header.shard_id().to_string();
-        let encode_timer = metrics::PARTIAL_WITNESS_ENCODE_TIME
+        let encode_timer = metrics::PARTIAL_WITNESS_ENCODE_AND_SEND_TIME
             .with_label_values(&[shard_id_label.as_str()])
             .start_timer();
         self.generate_and_send_state_witness_parts(

diff --git a/chain/client/src/test_utils/setup.rs b/chain/client/src/test_utils/setup.rs
@@ -756,7 +756,7 @@ fn process_peer_manager_message_default(
                 }
             }
         }
-        NetworkRequests::PartialEncodedStateWitness((account, partial_witness)) => {
+        NetworkRequests::PartialEncodedStateWitness(account, partial_witness) => {
             for (i, name) in validators.iter().enumerate() {
                 if name == account {
                     connectors[i]

diff --git a/chain/network/src/peer_manager/peer_manager_actor.rs b/chain/network/src/peer_manager/peer_manager_actor.rs
@@ -1069,7 +1069,7 @@ impl PeerManagerActor {
                 );
                 NetworkResponses::NoResponse
             }
-            NetworkRequests::PartialEncodedStateWitness((chunk_validator, partial_witness)) => {
+            NetworkRequests::PartialEncodedStateWitness(chunk_validator, partial_witness) => {
                 self.state.send_message_to_account(
                     &self.clock,
                     &chunk_validator,

diff --git a/chain/network/src/test_loop.rs b/chain/network/src/test_loop.rs
@@ -348,7 +348,7 @@ fn network_message_to_partial_witness_handler(
             None
         }
 
-        NetworkRequests::PartialEncodedStateWitness((target, partial_witness)) => {
+        NetworkRequests::PartialEncodedStateWitness(target, partial_witness) => {
             shared_state
                 .senders_for_account(&target)
                 .partial_witness_sender

diff --git a/chain/network/src/types.rs b/chain/network/src/types.rs
@@ -288,7 +288,7 @@ pub enum NetworkRequests {
     /// Message for a chunk endorsement, sent by a chunk validator to the block producer.
     ChunkEndorsement(AccountId, ChunkEndorsement),
     /// Message from chunk producer to set of chunk validators to send state witness part.
-    PartialEncodedStateWitness((AccountId, PartialEncodedStateWitness)),
+    PartialEncodedStateWitness(AccountId, PartialEncodedStateWitness),
     /// Message from chunk validator to all other chunk validators to forward state witness part.
     PartialEncodedStateWitnessForward(Vec<AccountId>, PartialEncodedStateWitness),
     /// Requests an epoch sync