diff --git a/chain/client/src/metrics.rs b/chain/client/src/metrics.rs index 971225cc790..5149e2671ac 100644 --- 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 = 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 = + 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 = 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 index 34f9d0278f6..3464b6d83e3 100644 --- a/chain/client/src/stateless_validation/partial_witness/partial_witness_actor.rs +++ b/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, Error> { + ) -> Result<(), Error> { tracing::debug!( target: "client", chunk_hash=?chunk_header.chunk_hash(), @@ -268,24 +268,29 @@ 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))| { - // 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, - ); - (chunk_validator.clone(), partial_witness) - }) - .collect_vec()) + 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( + chunk_validator.clone(), + partial_witness, + ), + )); + } + + Ok(()) } fn generate_contract_deploys_parts( @@ -340,10 +345,10 @@ 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(); - let validator_witness_tuple = self.generate_state_witness_parts( + self.generate_and_send_state_witness_parts( epoch_id, chunk_header, witness_bytes, @@ -357,13 +362,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 +599,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, diff --git a/chain/client/src/test_utils/setup.rs b/chain/client/src/test_utils/setup.rs index 22ca09bde99..4525e33f737 100644 --- a/chain/client/src/test_utils/setup.rs +++ b/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())); } } } diff --git a/chain/network/src/peer_manager/peer_manager_actor.rs b/chain/network/src/peer_manager/peer_manager_actor.rs index 1f4a00cf8d8..b198752cef8 100644 --- a/chain/network/src/peer_manager/peer_manager_actor.rs +++ b/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( diff --git a/chain/network/src/test_loop.rs b/chain/network/src/test_loop.rs index 08e01385487..8cd5fb948c9 100644 --- a/chain/network/src/test_loop.rs +++ b/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) => { diff --git a/chain/network/src/types.rs b/chain/network/src/types.rs index cd83e57fbd3..18cd3ac6898 100644 --- 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(Vec<(AccountId, PartialEncodedStateWitness)>), + PartialEncodedStateWitness(AccountId, PartialEncodedStateWitness), /// Message from chunk validator to all other chunk validators to forward state witness part. PartialEncodedStateWitnessForward(Vec, PartialEncodedStateWitness), /// Requests an epoch sync