Peerset cleanup (#6078)

* Move methods from Peerset to peers structs

* Remove priority_only from peersstate

* Refactor PSM

* Don't test private fields

* Update sc_network

* Remove wrong comment

* Also fix small stupidity when setting reserved_only

* Put back priority_group

* Restore priority groups as before

* Apply suggestions from code review

Co-authored-by: André Silva <123550+andresilva@users.noreply.github.com>

* Do the reserved only change

* Update client/peerset/src/lib.rs

Co-authored-by: Arkadiy Paronyan <arkady.paronyan@gmail.com>

* Use HashSet::difference

Co-authored-by: André Silva <123550+andresilva@users.noreply.github.com>
Co-authored-by: Arkadiy Paronyan <arkady.paronyan@gmail.com>

substrate/client/peerset/src/lib.rs

@@ -182,9 +182,13 @@ pub struct PeersetConfig {
 /// errors.
 #[derive(Debug)]
 pub struct Peerset {
+	/// Underlying data structure for the nodes's states.
 	data: peersstate::PeersState,
 	/// If true, we only accept reserved nodes.
 	reserved_only: bool,
+	/// Lists of nodes that don't occupy slots and that we should try to always be connected to.
+	/// Is kept in sync with the list of reserved nodes in [`Peerset::data`].
+	priority_groups: HashMap<String, HashSet<PeerId>>,
 	/// Receiver for messages from the `PeersetHandle` and from `tx`.
 	rx: TracingUnboundedReceiver<Action>,
 	/// Sending side of `rx`.
@@ -209,17 +213,18 @@ impl Peerset {
 		let now = Instant::now();
 
 		let mut peerset = Peerset {
-			data: peersstate::PeersState::new(config.in_peers, config.out_peers, config.reserved_only),
+			data: peersstate::PeersState::new(config.in_peers, config.out_peers),
 			tx,
 			rx,
 			reserved_only: config.reserved_only,
+			priority_groups: config.priority_groups.clone().into_iter().collect(),
 			message_queue: VecDeque::new(),
 			created: now,
 			latest_time_update: now,
 		};
 
-		for (group, nodes) in config.priority_groups {
-			peerset.data.set_priority_group(&group, nodes);
+		for node in config.priority_groups.into_iter().flat_map(|(_, l)| l) {
+			peerset.data.add_no_slot_node(node);
 		}
 
 		for peer_id in config.bootnodes {
@@ -235,61 +240,92 @@ impl Peerset {
 	}
 
 	fn on_add_reserved_peer(&mut self, peer_id: PeerId) {
-		let mut reserved = self.data.get_priority_group(RESERVED_NODES).unwrap_or_default();
-		reserved.insert(peer_id);
-		self.data.set_priority_group(RESERVED_NODES, reserved);
-		self.alloc_slots();
+		self.on_add_to_priority_group(RESERVED_NODES, peer_id);
 	}
 
 	fn on_remove_reserved_peer(&mut self, peer_id: PeerId) {
-		let mut reserved = self.data.get_priority_group(RESERVED_NODES).unwrap_or_default();
-		reserved.remove(&peer_id);
-		self.data.set_priority_group(RESERVED_NODES, reserved);
-		match self.data.peer(&peer_id) {
-			peersstate::Peer::Connected(peer) => {
-				if self.reserved_only {
-					peer.disconnect();
-					self.message_queue.push_back(Message::Drop(peer_id));
-				}
-			}
-			peersstate::Peer::NotConnected(_) => {},
-			peersstate::Peer::Unknown(_) => {},
-		}
+		self.on_remove_from_priority_group(RESERVED_NODES, peer_id);
 	}
 
 	fn on_set_reserved_only(&mut self, reserved_only: bool) {
 		self.reserved_only = reserved_only;
-		self.data.set_priority_only(reserved_only);
 
 		if self.reserved_only {
-			// Disconnect non-reserved nodes.
-			let reserved = self.data.get_priority_group(RESERVED_NODES).unwrap_or_default();
+			// Disconnect all the nodes that aren't reserved.
 			for peer_id in self.data.connected_peers().cloned().collect::<Vec<_>>().into_iter() {
+				if self.priority_groups.get(RESERVED_NODES).map_or(false, |g| g.contains(&peer_id)) {
+					continue;
+				}
+
 				let peer = self.data.peer(&peer_id).into_connected()
 					.expect("We are enumerating connected peers, therefore the peer is connected; qed");
-				if !reserved.contains(&peer_id) {
-					peer.disconnect();
-					self.message_queue.push_back(Message::Drop(peer_id));
-				}
+				peer.disconnect();
+				self.message_queue.push_back(Message::Drop(peer_id));
 			}
+
 		} else {
 			self.alloc_slots();
 		}
 	}
 
 	fn on_set_priority_group(&mut self, group_id: &str, peers: HashSet<PeerId>) {
-		self.data.set_priority_group(group_id, peers);
-		self.alloc_slots();
+		// Determine the difference between the current group and the new list.
+		let (to_insert, to_remove) = {
+			let current_group = self.priority_groups.entry(group_id.to_owned()).or_default();
+			let to_insert = peers.difference(current_group)
+				.cloned().collect::<Vec<_>>();
+			let to_remove = current_group.difference(&peers)
+				.cloned().collect::<Vec<_>>();
+			(to_insert, to_remove)
+		};
+
+		// Enumerate elements in `peers` not in `current_group`.
+		for peer_id in &to_insert {
+			// We don't call `on_add_to_priority_group` here in order to avoid calling
+			// `alloc_slots` all the time.
+			self.priority_groups.entry(group_id.to_owned()).or_default().insert(peer_id.clone());
+			self.data.add_no_slot_node(peer_id.clone());
+		}
+
+		// Enumerate elements in `current_group` not in `peers`.
+		for peer in to_remove {
+			self.on_remove_from_priority_group(group_id, peer);
+		}
+
+		if !to_insert.is_empty() {
+			self.alloc_slots();
+		}
 	}
 
 	fn on_add_to_priority_group(&mut self, group_id: &str, peer_id: PeerId) {
-		self.data.add_to_priority_group(group_id, peer_id);
+		self.priority_groups.entry(group_id.to_owned()).or_default().insert(peer_id.clone());
+		self.data.add_no_slot_node(peer_id);
 		self.alloc_slots();
 	}
 
 	fn on_remove_from_priority_group(&mut self, group_id: &str, peer_id: PeerId) {
-		self.data.remove_from_priority_group(group_id, &peer_id);
-		self.alloc_slots();
+		if let Some(priority_group) = self.priority_groups.get_mut(group_id) {
+			if !priority_group.remove(&peer_id) {
+				// `PeerId` wasn't in the group in the first place.
+				return;
+			}
+		} else {
+			// Group doesn't exist, so the `PeerId` can't be in it.
+			return;
+		}
+
+		// If that `PeerId` isn't in any other group, then it is no longer no-slot-occupying.
+		if !self.priority_groups.values().any(|l| l.contains(&peer_id)) {
+			self.data.remove_no_slot_node(&peer_id);
+		}
+
+		// Disconnect the peer if necessary.
+		if group_id != RESERVED_NODES && self.reserved_only {
+			if let peersstate::Peer::Connected(peer) = self.data.peer(&peer_id) {
+				peer.disconnect();
+				self.message_queue.push_back(Message::Drop(peer_id));
+			}
+		}
 	}
 
 	fn on_report_peer(&mut self, peer_id: PeerId, change: ReputationChange) {
@@ -376,25 +412,82 @@ impl Peerset {
 	fn alloc_slots(&mut self) {
 		self.update_time();
 
-		// Try to grab the next node to attempt to connect to.
-		while let Some(next) = {
-			if self.reserved_only {
-				self.data.priority_not_connected_peer_from_group(RESERVED_NODES)
-			} else {
-				self.data.priority_not_connected_peer()
-			}
-		} {
+		// Try to connect to all the reserved nodes that we are not connected to.
+		loop {
+			let next = {
+				let data = &mut self.data;
+				self.priority_groups
+					.get(RESERVED_NODES)
+					.into_iter()
+					.flatten()
+					.filter(move |n| {
+						data.peer(n).into_connected().is_none()
+					})
+					.next()
+					.cloned()
+			};
+
+			let next = match next {
+				Some(n) => n,
+				None => break,
+			};
+
+			let next = match self.data.peer(&next) {
+				peersstate::Peer::Unknown(n) => n.discover(),
+				peersstate::Peer::NotConnected(n) => n,
+				peersstate::Peer::Connected(_) => {
+					debug_assert!(false, "State inconsistency: not connected state");
+					break;
+				}
+			};
+
 			match next.try_outgoing() {
 				Ok(conn) => self.message_queue.push_back(Message::Connect(conn.into_peer_id())),
 				Err(_) => break,	// No more slots available.
 			}
 		}
 
-		loop {
-			if self.reserved_only {
-				break
-			}
+		// Nothing more to do if we're in reserved mode.
+		if self.reserved_only {
+			return;
+		}
 
+		// Try to connect to all the nodes in priority groups and that we are not connected to.
+		loop {
+			let next = {
+				let data = &mut self.data;
+				self.priority_groups
+					.values()
+					.flatten()
+					.filter(move |n| {
+						data.peer(n).into_connected().is_none()
+					})
+					.next()
+					.cloned()
+			};
+
+			let next = match next {
+				Some(n) => n,
+				None => break,
+			};
+
+			let next = match self.data.peer(&next) {
+				peersstate::Peer::Unknown(n) => n.discover(),
+				peersstate::Peer::NotConnected(n) => n,
+				peersstate::Peer::Connected(_) => {
+					debug_assert!(false, "State inconsistency: not connected state");
+					break;
+				}
+			};
+
+			match next.try_outgoing() {
+				Ok(conn) => self.message_queue.push_back(Message::Connect(conn.into_peer_id())),
+				Err(_) => break,	// No more slots available.
+			}
+		}
+
+		// Now, we try to connect to non-priority nodes.
+		loop {
 			// Try to grab the next node to attempt to connect to.
 			let next = match self.data.highest_not_connected_peer() {
 				Some(p) => p,
@@ -529,9 +622,9 @@ impl Peerset {
 		self.data.peers().len()
 	}
 
-	/// Returns priority group by id.
-	pub fn get_priority_group(&self, group_id: &str) -> Option<HashSet<PeerId>> {
-		self.data.get_priority_group(group_id)
+	/// Returns the content of a priority group.
+	pub fn priority_group(&self, group_id: &str) -> Option<impl ExactSizeIterator<Item = &PeerId>> {
+		self.priority_groups.get(group_id).map(|l| l.iter())
 	}
 }
 
@@ -583,7 +676,6 @@ mod tests {
 			assert_eq!(message, expected_message);
 			peerset = p;
 		}
-		assert!(peerset.message_queue.is_empty(), peerset.message_queue);
 		peerset
 	}
 
@@ -713,4 +805,3 @@ mod tests {
 		futures::executor::block_on(fut);
 	}
 }
-

substrate/client/peerset/src/peersstate.rs

@@ -14,10 +14,19 @@
 // You should have received a copy of the GNU General Public License
 // along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
 
-//! Contains the state storage behind the peerset.
+//! Reputation and slots allocation system behind the peerset.
+//!
+//! The [`PeersState`] state machine is responsible for managing the reputation and allocating
+//! slots. It holds a list of nodes, each associated with a reputation value and whether we are
+//! connected or not to this node. Thanks to this list, it knows how many slots are occupied. It
+//! also holds a list of nodes which don't occupy slots.
+//!
+//! > Note: This module is purely dedicated to managing slots and reputations. Features such as
+//! >       for example connecting to some nodes in priority should be added outside of this
+//! >       module, rather than inside.
 
 use libp2p::PeerId;
-use log::{error, warn};
+use log::error;
 use std::{borrow::Cow, collections::{HashSet, HashMap}};
 use wasm_timer::Instant;
 
@@ -37,23 +46,24 @@ pub struct PeersState {
 	/// 			sort, to make the logic easier.
 	nodes: HashMap<PeerId, Node>,
 
-	/// Number of non-priority nodes for which the `ConnectionState` is `In`.
+	/// Number of slot-occupying nodes for which the `ConnectionState` is `In`.
 	num_in: u32,
 
-	/// Number of non-priority nodes for which the `ConnectionState` is `In`.
+	/// Number of slot-occupying nodes for which the `ConnectionState` is `In`.
 	num_out: u32,
 
-	/// Maximum allowed number of non-priority nodes for which the `ConnectionState` is `In`.
+	/// Maximum allowed number of slot-occupying nodes for which the `ConnectionState` is `In`.
 	max_in: u32,
 
-	/// Maximum allowed number of non-priority nodes for which the `ConnectionState` is `Out`.
+	/// Maximum allowed number of slot-occupying nodes for which the `ConnectionState` is `Out`.
 	max_out: u32,
 
-	/// Priority groups. Each group is identified by a string ID and contains a set of peer IDs.
-	priority_nodes: HashMap<String, HashSet<PeerId>>,
-
-	/// Only allow connections to/from peers in a priority group.
-	priority_only: bool,
+	/// List of node identities (discovered or not) that don't occupy slots.
+	///
+	/// Note for future readers: this module is purely dedicated to managing slots. If you are
+	/// considering adding more features, please consider doing so outside of this module rather
+	/// than inside.
+	no_slot_nodes: HashSet<PeerId>,
 }
 
 /// State of a single node that we know about.
@@ -106,15 +116,14 @@ impl ConnectionState {
 
 impl PeersState {
 	/// Builds a new empty `PeersState`.
-	pub fn new(in_peers: u32, out_peers: u32, priority_only: bool) -> Self {
+	pub fn new(in_peers: u32, out_peers: u32) -> Self {
 		PeersState {
 			nodes: HashMap::new(),
 			num_in: 0,
 			num_out: 0,
 			max_in: in_peers,
 			max_out: out_peers,
-			priority_nodes: HashMap::new(),
-			priority_only,
+			no_slot_nodes: HashSet::new(),
 		}
 	}
 
@@ -157,34 +166,6 @@ impl PeersState {
 			.map(|(p, _)| p)
 	}
 
-	/// Returns the first priority peer that we are not connected to.
-	///
-	/// If multiple nodes are prioritized, which one is returned is unspecified.
-	pub fn priority_not_connected_peer(&mut self) -> Option<NotConnectedPeer> {
-		let id = self.priority_nodes.values()
-			.flatten()
-			.find(|&id| self.nodes.get(id).map_or(false, |node| !node.connection_state.is_connected()))
-			.cloned();
-		id.map(move |id| NotConnectedPeer {
-			state: self,
-			peer_id: Cow::Owned(id),
-		})
-	}
-
-	/// Returns the first priority peer that we are not connected to.
-	///
-	/// If multiple nodes are prioritized, which one is returned is unspecified.
-	pub fn priority_not_connected_peer_from_group(&mut self, group_id: &str) -> Option<NotConnectedPeer> {
-		let id = self.priority_nodes.get(group_id)
-			.and_then(|group| group.iter()
-				.find(|&id| self.nodes.get(id).map_or(false, |node| !node.connection_state.is_connected()))
-				.cloned());
-		id.map(move |id| NotConnectedPeer {
-			state: self,
-			peer_id: Cow::Owned(id),
-		})
-	}
-
 	/// Returns the peer with the highest reputation and that we are not connected to.
 	///
 	/// If multiple nodes have the same reputation, which one is returned is unspecified.
@@ -212,170 +193,40 @@ impl PeersState {
 		}
 	}
 
-	fn disconnect(&mut self, peer_id: &PeerId) {
-		let is_priority = self.is_priority(peer_id);
-		if let Some(mut node) = self.nodes.get_mut(peer_id) {
-			if !is_priority {
-				match node.connection_state {
-					ConnectionState::In => self.num_in -= 1,
-					ConnectionState::Out => self.num_out -= 1,
-					ConnectionState::NotConnected { .. } =>
-						debug_assert!(false, "State inconsistency: disconnecting a disconnected node")
-				}
-			}
-			node.connection_state = ConnectionState::NotConnected {
-				last_connected: Instant::now(),
-			};
-		} else {
-			warn!(target: "peerset", "Attempting to disconnect unknown peer {}", peer_id);
-		}
-	}
-
-	/// Sets the peer as connected with an outgoing connection.
-	fn try_outgoing(&mut self, peer_id: &PeerId) -> bool {
-		let is_priority = self.is_priority(peer_id);
-
-		// We are only accepting connections from priority nodes.
-		if !is_priority && self.priority_only {
-			return false;
-		}
-
-		// Note that it is possible for num_out to be strictly superior to the max, in case we were
-		// connected to reserved node then marked them as not reserved.
-		if self.num_out >= self.max_out && !is_priority {
-			return false;
-		}
-
-		if let Some(mut peer) = self.nodes.get_mut(peer_id) {
-			peer.connection_state = ConnectionState::Out;
-			if !is_priority {
-				self.num_out += 1;
-			}
-			return true;
-		}
-		false
-	}
-
-	/// Tries to accept the peer as an incoming connection.
+	/// Add a node to the list of nodes that don't occupy slots.
 	///
-	/// If there are enough slots available, switches the node to "connected" and returns `true`. If
-	/// the slots are full, the node stays "not connected" and we return `false`.
+	/// Has no effect if the peer was already in the group.
+	pub fn add_no_slot_node(&mut self, peer_id: PeerId) {
+		// Reminder: `HashSet::insert` returns false if the node was already in the set
+		if !self.no_slot_nodes.insert(peer_id.clone()) {
+			return;
+		}
+
+		if let Some(peer) = self.nodes.get_mut(&peer_id) {
+			match peer.connection_state {
+				ConnectionState::In => self.num_in -= 1,
+				ConnectionState::Out => self.num_out -= 1,
+				ConnectionState::NotConnected { .. } => {},
+			}
+		}
+	}
+
+	/// Removes a node from the list of nodes that don't occupy slots.
 	///
-	/// Note that reserved nodes don't count towards the number of slots.
-	fn try_accept_incoming(&mut self, peer_id: &PeerId) -> bool {
-		let is_priority = self.is_priority(peer_id);
-
-		// We are only accepting connections from priority nodes.
-		if !is_priority && self.priority_only {
-			return false;
+	/// Has no effect if the peer was not in the group.
+	pub fn remove_no_slot_node(&mut self, peer_id: &PeerId) {
+		// Reminder: `HashSet::remove` returns false if the node was already not in the set
+		if !self.no_slot_nodes.remove(peer_id) {
+			return;
 		}
 
-		// Note that it is possible for num_in to be strictly superior to the max, in case we were
-		// connected to reserved node then marked them as not reserved.
-		if self.num_in >= self.max_in && !is_priority {
-			return false;
-		}
-		if let Some(mut peer) = self.nodes.get_mut(peer_id) {
-			peer.connection_state = ConnectionState::In;
-			if !is_priority {
-				self.num_in += 1;
-			}
-			return true;
-		}
-		false
-	}
-
-	/// Sets priority group
-	pub fn set_priority_group(&mut self, group_id: &str, peers: HashSet<PeerId>) {
-		// update slot counters
-		let all_other_groups: HashSet<_> = self.priority_nodes
-			.iter()
-			.filter(|(g, _)| *g != group_id)
-			.flat_map(|(_, id)| id.clone())
-			.collect();
-		let existing_group = self.priority_nodes.remove(group_id).unwrap_or_default();
-		for id in existing_group {
-			// update slots for nodes that are no longer priority
-			if !all_other_groups.contains(&id) {
-				if let Some(peer) = self.nodes.get_mut(&id) {
-					match peer.connection_state {
-						ConnectionState::In => self.num_in += 1,
-						ConnectionState::Out => self.num_out += 1,
-						ConnectionState::NotConnected { .. } => {},
-					}
-				}
+		if let Some(peer) = self.nodes.get_mut(peer_id) {
+			match peer.connection_state {
+				ConnectionState::In => self.num_in += 1,
+				ConnectionState::Out => self.num_out += 1,
+				ConnectionState::NotConnected { .. } => {},
 			}
 		}
-
-		for id in &peers {
-			// update slots for nodes that become priority
-			if !all_other_groups.contains(id) {
-				let peer = self.nodes.entry(id.clone()).or_default();
-				match peer.connection_state {
-					ConnectionState::In => self.num_in -= 1,
-					ConnectionState::Out => self.num_out -= 1,
-					ConnectionState::NotConnected { .. } => {},
-				}
-			}
-		}
-		self.priority_nodes.insert(group_id.into(), peers);
-	}
-
-	/// Add a peer to a priority group.
-	pub fn add_to_priority_group(&mut self, group_id: &str, peer_id: PeerId) {
-		let mut peers = self.priority_nodes.get(group_id).cloned().unwrap_or_default();
-		peers.insert(peer_id);
-		self.set_priority_group(group_id, peers);
-	}
-
-	/// Remove a peer from a priority group.
-	pub fn remove_from_priority_group(&mut self, group_id: &str, peer_id: &PeerId) {
-		let mut peers = self.priority_nodes.get(group_id).cloned().unwrap_or_default();
-		peers.remove(peer_id);
-		self.set_priority_group(group_id, peers);
-	}
-
-	/// Get priority group content.
-	pub fn get_priority_group(&self, group_id: &str) -> Option<HashSet<PeerId>> {
-		self.priority_nodes.get(group_id).cloned()
-	}
-
-	/// Set whether to only allow connections to/from peers in a priority group.
-	/// Calling this method does not affect any existing connection, e.g.
-	/// enabling priority only will not disconnect from any non-priority peers
-	/// we are already connected to, only future incoming/outgoing connection
-	/// attempts will be affected.
-	pub fn set_priority_only(&mut self, priority: bool) {
-		self.priority_only = priority;
-	}
-
-	/// Check that node is any priority group.
-	fn is_priority(&self, peer_id: &PeerId) -> bool {
-		self.priority_nodes.iter().any(|(_, group)| group.contains(peer_id))
-	}
-
-	/// Returns the reputation value of the node.
-	fn reputation(&self, peer_id: &PeerId) -> i32 {
-		self.nodes.get(peer_id).map_or(0, |p| p.reputation)
-	}
-
-	/// Sets the reputation of the peer.
-	fn set_reputation(&mut self, peer_id: &PeerId, value: i32) {
-		let node = self.nodes
-			.entry(peer_id.clone())
-			.or_default();
-		node.reputation = value;
-	}
-
-	/// Performs an arithmetic addition on the reputation score of that peer.
-	///
-	/// In case of overflow, the value will be capped.
-	/// If the peer is unknown to us, we insert it and consider that it has a reputation of 0.
-	fn add_reputation(&mut self, peer_id: &PeerId, modifier: i32) {
-		let node = self.nodes
-			.entry(peer_id.clone())
-			.or_default();
-		node.reputation = node.reputation.saturating_add(modifier);
 	}
 }
 
@@ -437,7 +288,23 @@ impl<'a> ConnectedPeer<'a> {
 
 	/// Switches the peer to "not connected".
 	pub fn disconnect(self) -> NotConnectedPeer<'a> {
-		self.state.disconnect(&self.peer_id);
+		let is_no_slot_occupy = self.state.no_slot_nodes.contains(&*self.peer_id);
+		if let Some(mut node) = self.state.nodes.get_mut(&*self.peer_id) {
+			if !is_no_slot_occupy {
+				match node.connection_state {
+					ConnectionState::In => self.state.num_in -= 1,
+					ConnectionState::Out => self.state.num_out -= 1,
+					ConnectionState::NotConnected { .. } =>
+						debug_assert!(false, "State inconsistency: disconnecting a disconnected node")
+				}
+			}
+			node.connection_state = ConnectionState::NotConnected {
+				last_connected: Instant::now(),
+			};
+		} else {
+			debug_assert!(false, "State inconsistency: disconnecting a disconnected node");
+		}
+
 		NotConnectedPeer {
 			state: self.state,
 			peer_id: self.peer_id,
@@ -446,19 +313,27 @@ impl<'a> ConnectedPeer<'a> {
 
 	/// Returns the reputation value of the node.
 	pub fn reputation(&self) -> i32 {
-		self.state.reputation(&self.peer_id)
+		self.state.nodes.get(&*self.peer_id).map_or(0, |p| p.reputation)
 	}
 
 	/// Sets the reputation of the peer.
 	pub fn set_reputation(&mut self, value: i32) {
-		self.state.set_reputation(&self.peer_id, value)
+		if let Some(node) = self.state.nodes.get_mut(&*self.peer_id) {
+			node.reputation = value;
+		} else {
+			debug_assert!(false, "State inconsistency: set_reputation on an unknown node");
+		}
 	}
 
 	/// Performs an arithmetic addition on the reputation score of that peer.
 	///
 	/// In case of overflow, the value will be capped.
 	pub fn add_reputation(&mut self, modifier: i32) {
-		self.state.add_reputation(&self.peer_id, modifier)
+		if let Some(node) = self.state.nodes.get_mut(&*self.peer_id) {
+			node.reputation = node.reputation.saturating_add(modifier);
+		} else {
+			debug_assert!(false, "State inconsistency: add_reputation on an unknown node");
+		}
 	}
 }
 
@@ -520,16 +395,29 @@ impl<'a> NotConnectedPeer<'a> {
 	/// If there are enough slots available, switches the node to "connected" and returns `Ok`. If
 	/// the slots are full, the node stays "not connected" and we return `Err`.
 	///
-	/// Note that priority nodes don't count towards the number of slots.
+	/// Non-slot-occupying nodes don't count towards the number of slots.
 	pub fn try_outgoing(self) -> Result<ConnectedPeer<'a>, NotConnectedPeer<'a>> {
-		if self.state.try_outgoing(&self.peer_id) {
-			Ok(ConnectedPeer {
-				state: self.state,
-				peer_id: self.peer_id,
-			})
-		} else {
-			Err(self)
+		let is_no_slot_occupy = self.state.no_slot_nodes.contains(&*self.peer_id);
+
+		// Note that it is possible for num_out to be strictly superior to the max, in case we were
+		// connected to reserved node then marked them as not reserved.
+		if self.state.num_out >= self.state.max_out && !is_no_slot_occupy {
+			return Err(self);
 		}
+
+		if let Some(mut peer) = self.state.nodes.get_mut(&*self.peer_id) {
+			peer.connection_state = ConnectionState::Out;
+			if !is_no_slot_occupy {
+				self.state.num_out += 1;
+			}
+		} else {
+			debug_assert!(false, "State inconsistency: try_outgoing on an unknown node");
+		}
+
+		Ok(ConnectedPeer {
+			state: self.state,
+			peer_id: self.peer_id,
+		})
 	}
 
 	/// Tries to accept the peer as an incoming connection.
@@ -537,34 +425,54 @@ impl<'a> NotConnectedPeer<'a> {
 	/// If there are enough slots available, switches the node to "connected" and returns `Ok`. If
 	/// the slots are full, the node stays "not connected" and we return `Err`.
 	///
-	/// Note that priority nodes don't count towards the number of slots.
+	/// Non-slot-occupying nodes don't count towards the number of slots.
 	pub fn try_accept_incoming(self) -> Result<ConnectedPeer<'a>, NotConnectedPeer<'a>> {
-		if self.state.try_accept_incoming(&self.peer_id) {
-			Ok(ConnectedPeer {
-				state: self.state,
-				peer_id: self.peer_id,
-			})
-		} else {
-			Err(self)
+		let is_no_slot_occupy = self.state.no_slot_nodes.contains(&*self.peer_id);
+
+		// Note that it is possible for num_in to be strictly superior to the max, in case we were
+		// connected to reserved node then marked them as not reserved.
+		if self.state.num_in >= self.state.max_in && !is_no_slot_occupy {
+			return Err(self);
 		}
+
+		if let Some(mut peer) = self.state.nodes.get_mut(&*self.peer_id) {
+			peer.connection_state = ConnectionState::In;
+			if !is_no_slot_occupy {
+				self.state.num_in += 1;
+			}
+		} else {
+			debug_assert!(false, "State inconsistency: try_accept_incoming on an unknown node");
+		}
+
+		Ok(ConnectedPeer {
+			state: self.state,
+			peer_id: self.peer_id,
+		})
 	}
 
 	/// Returns the reputation value of the node.
 	pub fn reputation(&self) -> i32 {
-		self.state.reputation(&self.peer_id)
+		self.state.nodes.get(&*self.peer_id).map_or(0, |p| p.reputation)
 	}
 
 	/// Sets the reputation of the peer.
 	pub fn set_reputation(&mut self, value: i32) {
-		self.state.set_reputation(&self.peer_id, value)
+		if let Some(node) = self.state.nodes.get_mut(&*self.peer_id) {
+			node.reputation = value;
+		} else {
+			debug_assert!(false, "State inconsistency: set_reputation on an unknown node");
+		}
 	}
 
 	/// Performs an arithmetic addition on the reputation score of that peer.
 	///
 	/// In case of overflow, the value will be capped.
-	/// If the peer is unknown to us, we insert it and consider that it has a reputation of 0.
 	pub fn add_reputation(&mut self, modifier: i32) {
-		self.state.add_reputation(&self.peer_id, modifier)
+		if let Some(node) = self.state.nodes.get_mut(&*self.peer_id) {
+			node.reputation = node.reputation.saturating_add(modifier);
+		} else {
+			debug_assert!(false, "State inconsistency: add_reputation on an unknown node");
+		}
 	}
 
 	/// Un-discovers the peer. Removes it from the list.
@@ -618,7 +526,7 @@ mod tests {
 
 	#[test]
 	fn full_slots_in() {
-		let mut peers_state = PeersState::new(1, 1, false);
+		let mut peers_state = PeersState::new(1, 1);
 		let id1 = PeerId::random();
 		let id2 = PeerId::random();
 
@@ -632,14 +540,14 @@ mod tests {
 	}
 
 	#[test]
-	fn priority_node_doesnt_use_slot() {
-		let mut peers_state = PeersState::new(1, 1, false);
+	fn no_slot_node_doesnt_use_slot() {
+		let mut peers_state = PeersState::new(1, 1);
 		let id1 = PeerId::random();
 		let id2 = PeerId::random();
 
-		peers_state.set_priority_group("test", vec![id1.clone()].into_iter().collect());
-		if let Peer::NotConnected(p) = peers_state.peer(&id1) {
-			assert!(p.try_accept_incoming().is_ok());
+		peers_state.add_no_slot_node(id1.clone());
+		if let Peer::Unknown(p) = peers_state.peer(&id1) {
+			assert!(p.discover().try_accept_incoming().is_ok());
 		} else { panic!() }
 
 		if let Peer::Unknown(e) = peers_state.peer(&id2) {
@@ -649,7 +557,7 @@ mod tests {
 
 	#[test]
 	fn disconnecting_frees_slot() {
-		let mut peers_state = PeersState::new(1, 1, false);
+		let mut peers_state = PeersState::new(1, 1);
 		let id1 = PeerId::random();
 		let id2 = PeerId::random();
 
@@ -659,28 +567,9 @@ mod tests {
 		assert!(peers_state.peer(&id2).into_not_connected().unwrap().try_accept_incoming().is_ok());
 	}
 
-	#[test]
-	fn priority_not_connected_peer() {
-		let mut peers_state = PeersState::new(25, 25, false);
-		let id1 = PeerId::random();
-		let id2 = PeerId::random();
-
-		assert!(peers_state.priority_not_connected_peer().is_none());
-		peers_state.peer(&id1).into_unknown().unwrap().discover();
-		peers_state.peer(&id2).into_unknown().unwrap().discover();
-
-		assert!(peers_state.priority_not_connected_peer().is_none());
-		peers_state.set_priority_group("test", vec![id1.clone()].into_iter().collect());
-		assert!(peers_state.priority_not_connected_peer().is_some());
-		peers_state.set_priority_group("test", vec![id2.clone(), id2.clone()].into_iter().collect());
-		assert!(peers_state.priority_not_connected_peer().is_some());
-		peers_state.set_priority_group("test", vec![].into_iter().collect());
-		assert!(peers_state.priority_not_connected_peer().is_none());
-	}
-
 	#[test]
 	fn highest_not_connected_peer() {
-		let mut peers_state = PeersState::new(25, 25, false);
+		let mut peers_state = PeersState::new(25, 25);
 		let id1 = PeerId::random();
 		let id2 = PeerId::random();
 
@@ -700,87 +589,11 @@ mod tests {
 	}
 
 	#[test]
-	fn disconnect_priority_doesnt_panic() {
-		let mut peers_state = PeersState::new(1, 1, false);
+	fn disconnect_no_slot_doesnt_panic() {
+		let mut peers_state = PeersState::new(1, 1);
 		let id = PeerId::random();
-		peers_state.set_priority_group("test", vec![id.clone()].into_iter().collect());
-		let peer = peers_state.peer(&id).into_not_connected().unwrap().try_outgoing().unwrap();
+		peers_state.add_no_slot_node(id.clone());
+		let peer = peers_state.peer(&id).into_unknown().unwrap().discover().try_outgoing().unwrap();
 		peer.disconnect();
 	}
-
-	#[test]
-	fn multiple_priority_groups_slot_count() {
-		let mut peers_state = PeersState::new(1, 1, false);
-		let id = PeerId::random();
-
-		if let Peer::Unknown(p) = peers_state.peer(&id) {
-			assert!(p.discover().try_accept_incoming().is_ok());
-		} else { panic!() }
-
-		assert_eq!(peers_state.num_in, 1);
-		peers_state.set_priority_group("test1", vec![id.clone()].into_iter().collect());
-		assert_eq!(peers_state.num_in, 0);
-		peers_state.set_priority_group("test2", vec![id.clone()].into_iter().collect());
-		assert_eq!(peers_state.num_in, 0);
-		peers_state.set_priority_group("test1", vec![].into_iter().collect());
-		assert_eq!(peers_state.num_in, 0);
-		peers_state.set_priority_group("test2", vec![].into_iter().collect());
-		assert_eq!(peers_state.num_in, 1);
-	}
-
-	#[test]
-	fn priority_only_mode_ignores_drops_unknown_nodes() {
-		// test whether connection to/from given peer is allowed
-		let test_connection = |peers_state: &mut PeersState, id| {
-			if let Peer::Unknown(p) = peers_state.peer(id) {
-				p.discover();
-			}
-
-			let incoming = if let Peer::NotConnected(p) = peers_state.peer(id) {
-				p.try_accept_incoming().is_ok()
-			} else {
-				panic!()
-			};
-
-			if incoming {
-				peers_state.peer(id).into_connected().map(|p| p.disconnect());
-			}
-
-			let outgoing = if let Peer::NotConnected(p) = peers_state.peer(id) {
-				p.try_outgoing().is_ok()
-			} else {
-				panic!()
-			};
-
-			if outgoing {
-				peers_state.peer(id).into_connected().map(|p| p.disconnect());
-			}
-
-			incoming || outgoing
-		};
-
-		let mut peers_state = PeersState::new(1, 1, true);
-		let id = PeerId::random();
-
-		// this is an unknown peer and our peer state is set to only allow
-		// priority peers so any connection attempt should be denied.
-		assert!(!test_connection(&mut peers_state, &id));
-
-		// disabling priority only mode should allow the connection to go
-		// through.
-		peers_state.set_priority_only(false);
-		assert!(test_connection(&mut peers_state, &id));
-
-		// re-enabling it we should again deny connections from the peer.
-		peers_state.set_priority_only(true);
-		assert!(!test_connection(&mut peers_state, &id));
-
-		// but if we add the peer to a priority group it should be accepted.
-		peers_state.set_priority_group("TEST_GROUP", vec![id.clone()].into_iter().collect());
-		assert!(test_connection(&mut peers_state, &id));
-
-		// and removing it will cause the connection to once again be denied.
-		peers_state.remove_from_priority_group("TEST_GROUP", &id);
-		assert!(!test_connection(&mut peers_state, &id));
-	}
 }

substrate/client/peerset/tests/fuzz.rs

@@ -46,13 +46,13 @@ fn test_once() {
 			id
 		}).collect(),
 		priority_groups: {
-			let list = (0 .. Uniform::new_inclusive(0, 2).sample(&mut rng)).map(|_| {
+			let nodes = (0 .. Uniform::new_inclusive(0, 2).sample(&mut rng)).map(|_| {
 				let id = PeerId::random();
 				known_nodes.insert(id.clone());
 				reserved_nodes.insert(id.clone());
 				id
 			}).collect();
-			vec![("reserved".to_owned(), list)]
+			vec![("foo".to_string(), nodes)]
 		},
 		reserved_only: Uniform::new_inclusive(0, 10).sample(&mut rng) == 0,
 		in_peers: Uniform::new_inclusive(0, 25).sample(&mut rng),