Make sure we remove a peer on disconnect in gossip (paritytech#5104)

* Make sure we remove peers on disconnect in gossip state machine

* Clear up the code

* Add a comment

Author: bkchr

client/network-gossip/src/state_machine.rs

@@ -258,6 +258,7 @@ impl<B: BlockT> ConsensusGossip<B> {
 			let mut context = NetworkContext { gossip: self, network, engine_id: engine_id.clone() };
 			v.peer_disconnected(&mut context, &who);
 		}
+		self.peers.remove(&who);
 	}
 
 	/// Perform periodic maintenance
@@ -644,4 +645,52 @@ mod tests {
 		let _ = consensus.live_message_sinks.remove(&([0, 0, 0, 0], topic));
 		assert_eq!(stream.next(), None);
 	}
+
+	#[test]
+	fn peer_is_removed_on_disconnect() {
+		struct TestNetwork;
+		impl Network<Block> for TestNetwork {
+			fn event_stream(
+				&self,
+			) -> std::pin::Pin<Box<dyn futures::Stream<Item = crate::Event> + Send>> {
+				unimplemented!("Not required in tests")
+			}
+
+			fn report_peer(&self, _: PeerId, _: crate::ReputationChange) {
+				unimplemented!("Not required in tests")
+			}
+
+			fn disconnect_peer(&self, _: PeerId) {
+				unimplemented!("Not required in tests")
+			}
+
+			fn write_notification(&self, _: PeerId, _: crate::ConsensusEngineId, _: Vec<u8>) {
+				unimplemented!("Not required in tests")
+			}
+
+			fn register_notifications_protocol(
+				&self,
+				_: ConsensusEngineId,
+				_: std::borrow::Cow<'static, [u8]>,
+			) {
+				unimplemented!("Not required in tests")
+			}
+
+			fn announce(&self, _: H256, _: Vec<u8>) {
+				unimplemented!("Not required in tests")
+			}
+		}
+
+		let mut consensus = ConsensusGossip::<Block>::new();
+		consensus.register_validator_internal([0, 0, 0, 0], Arc::new(AllowAll));
+
+		let mut network = TestNetwork;
+
+		let peer_id = PeerId::random();
+		consensus.new_peer(&mut network, peer_id.clone(), Roles::FULL);
+		assert!(consensus.peers.contains_key(&peer_id));
+
+		consensus.peer_disconnected(&mut network, peer_id.clone());
+		assert!(!consensus.peers.contains_key(&peer_id));
+	}
 }

client/network/src/protocol/sync.rs

@@ -1167,8 +1167,7 @@ impl<B: BlockT> ChainSync<B> {
 	}
 
 	/// Restart the sync process.
-	fn restart<'a>(&'a mut self) -> impl Iterator<Item = Result<(PeerId, BlockRequest<B>), BadPeer>> + 'a
-	{
+	fn restart<'a>(&'a mut self) -> impl Iterator<Item = Result<(PeerId, BlockRequest<B>), BadPeer>> + 'a {
 		self.queue_blocks.clear();
 		self.blocks.clear();
 		let info = self.client.info();

client/network/src/protocol/sync/blocks.rs

@@ -104,8 +104,7 @@ impl<B: BlockT> BlockCollection<B> {
 		common: NumberFor<B>,
 		max_parallel: u32,
 		max_ahead: u32,
-	) -> Option<Range<NumberFor<B>>>
-	{
+	) -> Option<Range<NumberFor<B>>> {
 		if peer_best <= common {
 			// Bail out early
 			return None;
@@ -165,20 +164,20 @@ impl<B: BlockT> BlockCollection<B> {
 	pub fn drain(&mut self, from: NumberFor<B>) -> Vec<BlockData<B>> {
 		let mut drained = Vec::new();
 		let mut ranges = Vec::new();
-		{
-			let mut prev = from;
-			for (start, range_data) in &mut self.blocks {
-				match range_data {
-					&mut BlockRangeState::Complete(ref mut blocks) if *start <= prev => {
-							prev = *start + (blocks.len() as u32).into();
-							let mut blocks = mem::replace(blocks, Vec::new());
-							drained.append(&mut blocks);
-							ranges.push(*start);
-					},
-					_ => break,
-				}
+
+		let mut prev = from;
+		for (start, range_data) in &mut self.blocks {
+			match range_data {
+				&mut BlockRangeState::Complete(ref mut blocks) if *start <= prev => {
+					prev = *start + (blocks.len() as u32).into();
+					// Remove all elements from `blocks` and add them to `drained`
+					drained.append(blocks);
+					ranges.push(*start);
+				},
+				_ => break,
 			}
 		}
+
 		for r in ranges {
 			self.blocks.remove(&r);
 		}