feat(sync): use routed state part request in sync actor

chain/client/src/test_utils/peer_manager_mock.rs

@@ -1,5 +1,6 @@
-use near_network::types::SetChainInfo;
-use near_network::types::{PeerManagerMessageRequest, PeerManagerMessageResponse};
+use near_network::types::{
+    PeerManagerMessageRequest, PeerManagerMessageResponse, SetChainInfo, StateSyncEvent,
+};
 
 pub struct PeerManagerMock {
     handle: Box<
@@ -37,3 +38,8 @@ impl actix::Handler<SetChainInfo> for PeerManagerMock {
     type Result = ();
     fn handle(&mut self, _msg: SetChainInfo, _ctx: &mut Self::Context) {}
 }
+
+impl actix::Handler<StateSyncEvent> for PeerManagerMock {
+    type Result = ();
+    fn handle(&mut self, _msg: StateSyncEvent, _ctx: &mut Self::Context) {}
+}

chain/network/src/peer_manager/peer_manager_actor.rs

@@ -19,7 +19,7 @@ use crate::tcp;
 use crate::types::{
     ConnectedPeerInfo, HighestHeightPeerInfo, KnownProducer, NetworkInfo, NetworkRequests,
     NetworkResponses, PeerInfo, PeerManagerMessageRequest, PeerManagerMessageResponse, PeerType,
-    SetChainInfo, SnapshotHostInfo, StatePartRequestBody, Tier3RequestBody,
+    SetChainInfo, SnapshotHostInfo, StatePartRequestBody, StateSyncEvent, Tier3RequestBody,
 };
 use ::time::ext::InstantExt as _;
 use actix::fut::future::wrap_future;
@@ -1260,6 +1260,25 @@ impl actix::Handler<WithSpanContext<PeerManagerMessageRequest>> for PeerManagerA
     }
 }
 
+impl actix::Handler<WithSpanContext<StateSyncEvent>> for PeerManagerActor {
+    type Result = ();
+    #[perf]
+    fn handle(
+        &mut self,
+        msg: WithSpanContext<StateSyncEvent>,
+        _ctx: &mut Self::Context,
+    ) -> Self::Result {
+        let (_span, msg) = handler_debug_span!(target: "network", msg);
+        let _timer =
+            metrics::PEER_MANAGER_MESSAGES_TIME.with_label_values(&[(&msg).into()]).start_timer();
+        match msg {
+            StateSyncEvent::StatePartReceived(shard_id, part_id) => {
+                self.state.snapshot_hosts.part_received(shard_id, part_id);
+            }
+        }
+    }
+}
+
 impl actix::Handler<GetDebugStatus> for PeerManagerActor {
     type Result = DebugStatus;
     #[perf]

chain/network/src/snapshot_hosts/mod.rs

@@ -323,7 +323,6 @@ impl SnapshotHostsCache {
     }
 
     /// Triggered by state sync actor after processing a state part.
-    #[allow(dead_code)]
     pub fn part_received(&self, shard_id: ShardId, part_id: u64) {
         let mut inner = self.0.lock();
         inner.peer_selector.remove(&(shard_id, part_id));

chain/network/src/test_loop.rs

@@ -13,7 +13,7 @@ use crate::state_witness::{
 };
 use crate::types::{
     NetworkRequests, NetworkResponses, PeerManagerMessageRequest, PeerManagerMessageResponse,
-    SetChainInfo,
+    SetChainInfo, StateSyncEvent,
 };
 use near_async::actix::ActixResult;
 use near_async::futures::{FutureSpawner, FutureSpawnerExt};
@@ -188,6 +188,10 @@ impl Handler<SetChainInfo> for TestLoopPeerManagerActor {
     fn handle(&mut self, _msg: SetChainInfo) {}
 }
 
+impl Handler<StateSyncEvent> for TestLoopPeerManagerActor {
+    fn handle(&mut self, _msg: StateSyncEvent) {}
+}
+
 impl Handler<PeerManagerMessageRequest> for TestLoopPeerManagerActor {
     fn handle(&mut self, msg: PeerManagerMessageRequest) -> PeerManagerMessageResponse {
         let PeerManagerMessageRequest::NetworkRequests(request) = msg else {

chain/network/src/test_utils.rs

@@ -1,7 +1,7 @@
 use crate::network_protocol::PeerInfo;
 use crate::types::{
     NetworkInfo, NetworkResponses, PeerManagerMessageRequest, PeerManagerMessageResponse,
-    SetChainInfo,
+    SetChainInfo, StateSyncEvent,
 };
 use crate::PeerManagerActor;
 use actix::{Actor, ActorContext, Context, Handler};
@@ -259,6 +259,10 @@ impl CanSend<SetChainInfo> for MockPeerManagerAdapter {
     fn send(&self, _msg: SetChainInfo) {}
 }
 
+impl CanSend<StateSyncEvent> for MockPeerManagerAdapter {
+    fn send(&self, _msg: StateSyncEvent) {}
+}
+
 impl MockPeerManagerAdapter {
     pub fn pop(&self) -> Option<PeerManagerMessageRequest> {
         self.requests.write().unwrap().pop_front()

chain/network/src/types.rs

@@ -293,6 +293,12 @@ pub enum NetworkRequests {
     EpochSyncResponse { route_back: CryptoHash, proof: CompressedEpochSyncProof },
 }
 
+#[derive(Debug, actix::Message, strum::IntoStaticStr)]
+#[rtype(result = "()")]
+pub enum StateSyncEvent {
+    StatePartReceived(ShardId, u64),
+}
+
 /// Combines peer address info, chain.
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct FullPeerInfo {
@@ -404,6 +410,7 @@ pub struct PeerManagerAdapter {
     pub async_request_sender: AsyncSender<PeerManagerMessageRequest, PeerManagerMessageResponse>,
     pub request_sender: Sender<PeerManagerMessageRequest>,
     pub set_chain_info_sender: Sender<SetChainInfo>,
+    pub state_sync_event_sender: Sender<StateSyncEvent>,
 }
 
 #[cfg(test)]