diff --git a/chain/client/src/client.rs b/chain/client/src/client.rs index 8c2edaa0818..fff1fffef3f 100644 --- a/chain/client/src/client.rs +++ b/chain/client/src/client.rs @@ -7,7 +7,9 @@ use crate::debug::BlockProductionTracker; use crate::debug::PRODUCTION_TIMES_CACHE_SIZE; use crate::stateless_validation::chunk_endorsement_tracker::ChunkEndorsementTracker; use crate::stateless_validation::chunk_validator::ChunkValidator; -use crate::stateless_validation::partial_witness::partial_witness_actor::PartialWitnessSenderForClient; +use crate::stateless_validation::partial_witness::partial_witness_actor::{ + PartialWitnessSenderForClient, UpdateChainHeadRequest, +}; use crate::sync::adapter::SyncShardInfo; use crate::sync::block::BlockSync; use crate::sync::epoch::EpochSync; @@ -1302,10 +1304,12 @@ impl Client { apply_chunks_done_sender, ); if accepted_blocks.iter().any(|accepted_block| accepted_block.status.is_new_head()) { + let head = self.chain.head().unwrap(); self.shards_manager_adapter.send(ShardsManagerRequestFromClient::UpdateChainHeads { - head: self.chain.head().unwrap(), + head: head.clone(), header_head: self.chain.header_head().unwrap(), }); + self.partial_witness_adapter.send(UpdateChainHeadRequest { head }); } self.process_block_processing_artifact(block_processing_artifacts); let accepted_blocks_hashes = 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 21c38acda5e..97433b4527a 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 @@ -4,6 +4,7 @@ use itertools::Itertools; use near_async::messaging::{CanSend, Handler, Sender}; use near_async::time::Clock; use near_async::{MultiSend, MultiSendMessage, MultiSenderFrom}; +use near_chain::types::Tip; use near_chain::Error; use near_epoch_manager::EpochManagerAdapter; use near_network::state_witness::{ @@ -52,10 +53,17 @@ pub struct DistributeStateWitnessRequest { pub state_witness: ChunkStateWitness, } +#[derive(actix::Message, Debug)] +#[rtype(result = "()")] +pub struct UpdateChainHeadRequest { + pub head: Tip, +} + #[derive(Clone, MultiSend, MultiSenderFrom, MultiSendMessage)] #[multi_send_message_derive(Debug)] pub struct PartialWitnessSenderForClient { pub distribute_chunk_state_witness: Sender, + pub update_chain_head: Sender, } impl Handler for PartialWitnessActor { @@ -67,6 +75,11 @@ impl Handler for PartialWitnessActor { } } +impl Handler for PartialWitnessActor { + #[perf] + fn handle(&mut self, _msg: UpdateChainHeadRequest) {} +} + impl Handler for PartialWitnessActor { fn handle(&mut self, msg: ChunkStateWitnessAckMessage) { self.handle_chunk_state_witness_ack(msg.0); diff --git a/chain/client/src/test_utils/mock_partial_witness_adapter.rs b/chain/client/src/test_utils/mock_partial_witness_adapter.rs index e006f0ad479..235ff495eba 100644 --- a/chain/client/src/test_utils/mock_partial_witness_adapter.rs +++ b/chain/client/src/test_utils/mock_partial_witness_adapter.rs @@ -3,7 +3,9 @@ use std::sync::{Arc, RwLock}; use near_async::messaging::CanSend; -use crate::stateless_validation::partial_witness::partial_witness_actor::DistributeStateWitnessRequest; +use crate::stateless_validation::partial_witness::partial_witness_actor::{ + DistributeStateWitnessRequest, UpdateChainHeadRequest, +}; #[derive(Clone, Default)] pub struct MockPartialWitnessAdapter { @@ -16,6 +18,10 @@ impl CanSend for MockPartialWitnessAdapter { } } +impl CanSend for MockPartialWitnessAdapter { + fn send(&self, _msg: UpdateChainHeadRequest) {} +} + impl MockPartialWitnessAdapter { pub fn pop_distribution_request(&self) -> Option { self.distribution_request.write().unwrap().pop_front() diff --git a/chain/client/src/test_utils/test_loop/partial_witness_actor.rs b/chain/client/src/test_utils/test_loop/partial_witness_actor.rs index 65fa20742ac..9abdb184c63 100644 --- a/chain/client/src/test_utils/test_loop/partial_witness_actor.rs +++ b/chain/client/src/test_utils/test_loop/partial_witness_actor.rs @@ -24,5 +24,8 @@ pub fn forward_messages_from_client_to_partial_witness_actor( PartialWitnessSenderForClientMessage::_distribute_chunk_state_witness(msg) => { partial_witness_actor.handle(msg); } + PartialWitnessSenderForClientMessage::_update_chain_head(msg) => { + partial_witness_actor.handle(msg); + } }) }