chore: send partial witnesses faster

chain/client/src/stateless_validation/partial_witness/partial_witness_actor.rs

@@ -249,14 +249,14 @@ impl PartialWitnessActor {
     }
 
     // Function to generate the parts of the state witness and return them as a tuple of chunk_validator and part.
-    fn generate_state_witness_parts(
+    fn generate_and_send_state_witness_parts(
         &mut self,
         epoch_id: EpochId,
         chunk_header: &ShardChunkHeader,
         witness_bytes: EncodedChunkStateWitness,
         chunk_validators: &[AccountId],
         signer: &ValidatorSigner,
-    ) -> Result<Vec<(AccountId, PartialEncodedStateWitness)>, Error> {
+    ) -> Result<(), Error> {
         tracing::debug!(
             target: "client",
             chunk_hash=?chunk_header.chunk_hash(),
@@ -268,11 +268,8 @@ impl PartialWitnessActor {
         let encoder = self.witness_encoders.entry(chunk_validators.len());
         let (parts, encoded_length) = encoder.encode(&witness_bytes);
 
-        Ok(chunk_validators
-            .iter()
-            .zip_eq(parts)
-            .enumerate()
-            .map(|(part_ord, (chunk_validator, part))| {
+        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(
@@ -283,9 +280,16 @@ impl PartialWitnessActor {
                     encoded_length,
                     signer,
                 );
-                (chunk_validator.clone(), partial_witness)
-            })
-            .collect_vec())
+                let validator_witness_tuple = (chunk_validator.clone(), partial_witness);
+
+                // Send the parts to the corresponding chunk validator owners.
+                self.network_adapter.send(PeerManagerMessageRequest::NetworkRequests(
+                    NetworkRequests::PartialEncodedStateWitness(validator_witness_tuple),
+                ));
+            },
+        );
+
+        Ok(())
     }
 
     fn generate_contract_deploys_parts(
@@ -343,7 +347,7 @@ impl PartialWitnessActor {
         let encode_timer = metrics::PARTIAL_WITNESS_ENCODE_TIME
             .with_label_values(&[shard_id_label.as_str()])
             .start_timer();
-        let validator_witness_tuple = self.generate_state_witness_parts(
+        self.generate_and_send_state_witness_parts(
             epoch_id,
             chunk_header,
             witness_bytes,
@@ -357,13 +361,9 @@ impl PartialWitnessActor {
         self.state_witness_tracker.record_witness_sent(
             chunk_hash,
             witness_size_in_bytes,
-            validator_witness_tuple.len(),
+            chunk_validators.len(),
         );
 
-        // Send the parts to the corresponding chunk validator owners.
-        self.network_adapter.send(PeerManagerMessageRequest::NetworkRequests(
-            NetworkRequests::PartialEncodedStateWitness(validator_witness_tuple),
-        ));
         Ok(())
     }
 
@@ -598,7 +598,7 @@ impl PartialWitnessActor {
 
     /// Sends the contract accesses to the same chunk validators
     /// (except for the chunk producers that track the same shard),
-    /// which will receive the state witness for the new chunk.  
+    /// which will receive the state witness for the new chunk.
     fn send_contract_accesses_to_chunk_validators(
         &self,
         key: ChunkProductionKey,

chain/client/src/test_utils/setup.rs

@@ -756,14 +756,12 @@ fn process_peer_manager_message_default(
                 }
             }
         }
-        NetworkRequests::PartialEncodedStateWitness(partial_witnesses) => {
-            for (account, partial_witness) in partial_witnesses {
-                for (i, name) in validators.iter().enumerate() {
-                    if name == account {
-                        connectors[i]
-                            .partial_witness_sender
-                            .send(PartialEncodedStateWitnessMessage(partial_witness.clone()));
-                    }
+        NetworkRequests::PartialEncodedStateWitness((account, partial_witness)) => {
+            for (i, name) in validators.iter().enumerate() {
+                if name == account {
+                    connectors[i]
+                        .partial_witness_sender
+                        .send(PartialEncodedStateWitnessMessage(partial_witness.clone()));
                 }
             }
         }

chain/network/src/peer_manager/peer_manager_actor.rs

@@ -1069,14 +1069,12 @@ impl PeerManagerActor {
                 );
                 NetworkResponses::NoResponse
             }
-            NetworkRequests::PartialEncodedStateWitness(validator_witness_tuple) => {
-                for (chunk_validator, partial_witness) in validator_witness_tuple {
-                    self.state.send_message_to_account(
-                        &self.clock,
-                        &chunk_validator,
-                        RoutedMessageBody::PartialEncodedStateWitness(partial_witness),
-                    );
-                }
+            NetworkRequests::PartialEncodedStateWitness((chunk_validator, partial_witness)) => {
+                self.state.send_message_to_account(
+                    &self.clock,
+                    &chunk_validator,
+                    RoutedMessageBody::PartialEncodedStateWitness(partial_witness),
+                );
                 NetworkResponses::NoResponse
             }
             NetworkRequests::PartialEncodedStateWitnessForward(

chain/network/src/test_loop.rs

@@ -348,13 +348,11 @@ fn network_message_to_partial_witness_handler(
             None
         }
 
-        NetworkRequests::PartialEncodedStateWitness(validator_witness_tuple) => {
-            for (target, partial_witness) in validator_witness_tuple.into_iter() {
-                shared_state
-                    .senders_for_account(&target)
-                    .partial_witness_sender
-                    .send(PartialEncodedStateWitnessMessage(partial_witness));
-            }
+        NetworkRequests::PartialEncodedStateWitness((target, partial_witness)) => {
+            shared_state
+                .senders_for_account(&target)
+                .partial_witness_sender
+                .send(PartialEncodedStateWitnessMessage(partial_witness));
             None
         }
         NetworkRequests::PartialEncodedStateWitnessForward(chunk_validators, partial_witness) => {

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(Vec<(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