add NetworkRequests::SnapshotHostInfo

saketh-are · Nov 2, 2023 · 4b6e295 · 4b6e295
1 parent a7ddcd0
commit 4b6e295
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 5 deletions.
diff --git a/chain/client/src/test_utils/setup.rs b/chain/client/src/test_utils/setup.rs
@@ -826,6 +826,7 @@ pub fn setup_mock_all_validators(
                         NetworkRequests::ForwardTx(_, _)
                         | NetworkRequests::BanPeer { .. }
                         | NetworkRequests::TxStatus(_, _, _)
+                        | NetworkRequests::SnapshotHostInfo(_)
                         | NetworkRequests::Challenge(_) => {}
                     };
                 }

diff --git a/chain/network/src/network_protocol/mod.rs b/chain/network/src/network_protocol/mod.rs
@@ -5,7 +5,7 @@ mod borsh_conv;
 mod edge;
 mod peer;
 mod proto_conv;
-pub mod state_sync;
+mod state_sync;
 pub use edge::*;
 pub use peer::*;
 pub use state_sync::*;
@@ -24,7 +24,6 @@ pub use _proto::network as proto;
 use crate::network_protocol::proto_conv::trace_context::{
     extract_span_context, inject_trace_context,
 };
-use crate::network_protocol::state_sync::SyncSnapshotHosts;
 use borsh::BorshDeserialize as _;
 use near_async::time;
 use near_crypto::PublicKey;

diff --git a/chain/network/src/peer_manager/connection/mod.rs b/chain/network/src/peer_manager/connection/mod.rs
@@ -1,10 +1,9 @@
 use crate::concurrency::arc_mutex::ArcMutex;
 use crate::concurrency::atomic_cell::AtomicCell;
 use crate::concurrency::demux;
-use crate::network_protocol::state_sync::SyncSnapshotHosts;
 use crate::network_protocol::{
     PeerInfo, PeerMessage, RoutedMessageBody, SignedAccountData, SignedOwnedAccount,
-    SnapshotHostInfo, SyncAccountsData,
+    SnapshotHostInfo, SyncAccountsData, SyncSnapshotHosts,
 };
 use crate::peer::peer_actor;
 use crate::peer::peer_actor::PeerActor;

diff --git a/chain/network/src/peer_manager/peer_manager_actor.rs b/chain/network/src/peer_manager/peer_manager_actor.rs
@@ -1,6 +1,7 @@
 use crate::client;
 use crate::config;
 use crate::debug::{DebugStatus, GetDebugStatus};
+use crate::network_protocol::SyncSnapshotHosts;
 use crate::network_protocol::{
     Disconnect, Edge, PeerIdOrHash, PeerMessage, Ping, Pong, RawRoutedMessage, RoutedMessageBody,
 };
@@ -783,6 +784,12 @@ impl PeerManagerActor {
                     NetworkResponses::RouteNotFound
                 }
             }
+            NetworkRequests::SnapshotHostInfo(snapshot_host_info) => {
+                self.state.tier2.broadcast_message(Arc::new(PeerMessage::SyncSnapshotHosts(
+                    SyncSnapshotHosts { hosts: vec![snapshot_host_info] },
+                )));
+                NetworkResponses::NoResponse
+            }
             NetworkRequests::BanPeer { peer_id, ban_reason } => {
                 self.state.disconnect_and_ban(&self.clock, &peer_id, ban_reason);
                 NetworkResponses::NoResponse

diff --git a/chain/network/src/types.rs b/chain/network/src/types.rs
@@ -6,7 +6,7 @@ pub use crate::network_protocol::{
 /// Exported types, which are part of network protocol.
 pub use crate::network_protocol::{
     Edge, PartialEdgeInfo, PartialEncodedChunkForwardMsg, PartialEncodedChunkRequestMsg,
-    PartialEncodedChunkResponseMsg, PeerChainInfoV2, PeerInfo, StateResponseInfo,
+    PartialEncodedChunkResponseMsg, PeerChainInfoV2, PeerInfo, SnapshotHostInfo, StateResponseInfo,
     StateResponseInfoV1, StateResponseInfoV2,
 };
 use crate::routing::routing_table_view::RoutingTableInfo;
@@ -232,6 +232,8 @@ pub enum NetworkRequests {
     BanPeer { peer_id: PeerId, ban_reason: ReasonForBan },
     /// Announce account
     AnnounceAccount(AnnounceAccount),
+    /// Broadcast information about a hosted snapshot.
+    SnapshotHostInfo(Arc<SnapshotHostInfo>),
 
     /// Request chunk parts and/or receipts
     PartialEncodedChunkRequest {