impl

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 =

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<DistributeStateWitnessRequest>,
+    pub update_chain_head: Sender<UpdateChainHeadRequest>,
 }
 
 impl Handler<DistributeStateWitnessRequest> for PartialWitnessActor {
@@ -67,6 +75,11 @@ impl Handler<DistributeStateWitnessRequest> for PartialWitnessActor {
     }
 }
 
+impl Handler<UpdateChainHeadRequest> for PartialWitnessActor {
+    #[perf]
+    fn handle(&mut self, _msg: UpdateChainHeadRequest) {}
+}
+
 impl Handler<ChunkStateWitnessAckMessage> for PartialWitnessActor {
     fn handle(&mut self, msg: ChunkStateWitnessAckMessage) {
         self.handle_chunk_state_witness_ack(msg.0);

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<DistributeStateWitnessRequest> for MockPartialWitnessAdapter {
     }
 }
 
+impl CanSend<UpdateChainHeadRequest> for MockPartialWitnessAdapter {
+    fn send(&self, _msg: UpdateChainHeadRequest) {}
+}
+
 impl MockPartialWitnessAdapter {
     pub fn pop_distribution_request(&self) -> Option<DistributeStateWitnessRequest> {
         self.distribution_request.write().unwrap().pop_front()

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);
+        }
     })
 }