From b31bcdfcc38f01025e96e6c3cf865504e2656ba3 Mon Sep 17 00:00:00 2001
From: Arkadiy Paronyan <arkady.paronyan@gmail.com>
Date: Thu, 21 Mar 2019 12:55:30 +0100
Subject: [PATCH] Keep peer info up to date (#2067)

* Keep peer info up to date

* Docs

Co-Authored-By: arkpar <arkady.paronyan@gmail.com>
---
 substrate/core/network/src/protocol.rs | 20 +++++++++++++++++++-
 substrate/core/network/src/sync.rs     | 19 +++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/substrate/core/network/src/protocol.rs b/substrate/core/network/src/protocol.rs
index bcfd1d07e5..d33e0d0af1 100644
--- a/substrate/core/network/src/protocol.rs
+++ b/substrate/core/network/src/protocol.rs
@@ -445,6 +445,20 @@ impl<B: BlockT, S: NetworkSpecialization<B>, H: ExHashT> Protocol<B, S, H> {
 		None
 	}
 
+	fn update_peer_info(&mut self, who: NodeIndex) {
+		if let Some(info) = self.sync.peer_info(who) {
+			if let Some(ref mut peer) = self.context_data.peers.get_mut(&who) {
+				peer.info.best_hash = info.best_hash;
+				peer.info.best_number = info.best_number;
+			}
+			let mut peers = self.connected_peers.write();
+			if let Some(ref mut peer) = peers.get_mut(&who) {
+				peer.peer_info.best_hash = info.best_hash;
+				peer.peer_info.best_number = info.best_number;
+			}
+		}
+	}
+
 	/// Propagates protocol statuses.
 	fn on_status(&mut self) {
 		let status = ProtocolStatus {
@@ -467,9 +481,13 @@ impl<B: BlockT, S: NetworkSpecialization<B>, H: ExHashT> Protocol<B, S, H> {
 			GenericMessage::BlockResponse(r) => {
 				if let Some(request) = self.handle_response(who, &r) {
 					self.on_block_response(who, request, r);
+					self.update_peer_info(who);
 				}
 			},
-			GenericMessage::BlockAnnounce(announce) => self.on_block_announce(who, announce),
+			GenericMessage::BlockAnnounce(announce) => {
+				self.on_block_announce(who, announce);
+				self.update_peer_info(who);
+			},
 			GenericMessage::Transactions(m) => self.on_extrinsics(who, m),
 			GenericMessage::RemoteCallRequest(request) => self.on_remote_call_request(who, request),
 			GenericMessage::RemoteCallResponse(response) => self.on_remote_call_response(who, response),
diff --git a/substrate/core/network/src/sync.rs b/substrate/core/network/src/sync.rs
index a4dc1855ef..76dbeabe79 100644
--- a/substrate/core/network/src/sync.rs
+++ b/substrate/core/network/src/sync.rs
@@ -58,6 +58,15 @@ struct PeerSync<B: BlockT> {
 	pub recently_announced: VecDeque<B::Hash>,
 }
 
+#[derive(Debug)]
+/// Peer sync status.
+pub(crate) struct PeerInfo<B: BlockT> {
+	/// Their best block hash.
+	pub best_hash: B::Hash,
+	/// Their best block number.
+	pub best_number: NumberFor<B>,
+}
+
 #[derive(Copy, Clone, Eq, PartialEq, Debug)]
 enum AncestorSearchState<B: BlockT> {
 	/// Use exponential backoff to find an ancestor, then switch to binary search.
@@ -408,6 +417,16 @@ impl<B: BlockT> ChainSync<B> {
 		}
 	}
 
+	/// Returns peer sync status (if any).
+	pub(crate) fn peer_info(&self, who: NodeIndex) -> Option<PeerInfo<B>> {
+		self.peers.get(&who).map(|peer| {
+			PeerInfo {
+				best_hash: peer.best_hash,
+				best_number: peer.best_number,
+			}
+		})
+	}
+
 	/// Returns sync status.
 	pub(crate) fn status(&self) -> Status<B> {
 		let best_seen = self.best_seen_block();