mirror of
https://github.com/pezkuwichain/pezkuwi-subxt.git
synced 2026-06-09 20:11:09 +00:00
Peer prioritization (#2717)
* Added priority groups * Added a test * Whitespace * Added add/remove single peer API * Added a warning * Fixed removing reserved peer * Fixed build * Made some methods private and made get_priority_group return an Option
This commit is contained in:
committed by
Gavin Wood
parent
fff90e86c0
commit
7885068fac
@@ -19,7 +19,7 @@
|
||||
|
||||
mod peersstate;
|
||||
|
||||
use std::{collections::HashMap, collections::VecDeque, time::Instant};
|
||||
use std::{collections::{HashSet, HashMap}, collections::VecDeque, time::Instant};
|
||||
use futures::{prelude::*, sync::mpsc, try_ready};
|
||||
use libp2p::PeerId;
|
||||
use log::{debug, error, trace};
|
||||
@@ -29,6 +29,8 @@ use serde_json::json;
|
||||
const BANNED_THRESHOLD: i32 = 82 * (i32::min_value() / 100);
|
||||
/// Reputation change for a node when we get disconnected from it.
|
||||
const DISCONNECT_REPUTATION_CHANGE: i32 = -10;
|
||||
/// Reserved peers group ID
|
||||
const RESERVED_NODES: &'static str = "reserved";
|
||||
|
||||
#[derive(Debug)]
|
||||
enum Action {
|
||||
@@ -36,6 +38,9 @@ enum Action {
|
||||
RemoveReservedPeer(PeerId),
|
||||
SetReservedOnly(bool),
|
||||
ReportPeer(PeerId, i32),
|
||||
SetPriorityGroup(String, HashSet<PeerId>),
|
||||
AddToPriorityGroup(String, PeerId),
|
||||
RemoveFromPriorityGroup(String, PeerId),
|
||||
}
|
||||
|
||||
/// Shared handle to the peer set manager (PSM). Distributed around the code.
|
||||
@@ -72,6 +77,21 @@ impl PeersetHandle {
|
||||
pub fn report_peer(&self, peer_id: PeerId, score_diff: i32) {
|
||||
let _ = self.tx.unbounded_send(Action::ReportPeer(peer_id, score_diff));
|
||||
}
|
||||
|
||||
/// Modify a priority group.
|
||||
pub fn set_priority_group(&self, group_id: String, peers: HashSet<PeerId>) {
|
||||
let _ = self.tx.unbounded_send(Action::SetPriorityGroup(group_id, peers));
|
||||
}
|
||||
|
||||
/// Add a peer to a priority group.
|
||||
pub fn add_to_priority_group(&self, group_id: String, peer_id: PeerId) {
|
||||
let _ = self.tx.unbounded_send(Action::AddToPriorityGroup(group_id, peer_id));
|
||||
}
|
||||
|
||||
/// Remove a peer from a priority group.
|
||||
pub fn remove_from_priority_group(&self, group_id: String, peer_id: PeerId) {
|
||||
let _ = self.tx.unbounded_send(Action::RemoveFromPriorityGroup(group_id, peer_id));
|
||||
}
|
||||
}
|
||||
|
||||
/// Message that can be sent by the peer set manager (PSM).
|
||||
@@ -161,14 +181,7 @@ impl Peerset {
|
||||
latest_time_update: Instant::now(),
|
||||
};
|
||||
|
||||
for peer_id in config.reserved_nodes {
|
||||
if let peersstate::Peer::Unknown(entry) = peerset.data.peer(&peer_id) {
|
||||
entry.discover().set_reserved(true);
|
||||
} else {
|
||||
debug!(target: "peerset", "Duplicate reserved node in config: {:?}", peer_id);
|
||||
}
|
||||
}
|
||||
|
||||
peerset.data.set_priority_group(RESERVED_NODES, config.reserved_nodes.into_iter().collect());
|
||||
for peer_id in config.bootnodes {
|
||||
if let peersstate::Peer::Unknown(entry) = peerset.data.peer(&peer_id) {
|
||||
entry.discover();
|
||||
@@ -182,32 +195,25 @@ impl Peerset {
|
||||
}
|
||||
|
||||
fn on_add_reserved_peer(&mut self, peer_id: PeerId) {
|
||||
let mut entry = match self.data.peer(&peer_id) {
|
||||
peersstate::Peer::Connected(mut connected) => {
|
||||
connected.set_reserved(true);
|
||||
return
|
||||
}
|
||||
peersstate::Peer::NotConnected(entry) => entry,
|
||||
peersstate::Peer::Unknown(entry) => entry.discover(),
|
||||
};
|
||||
|
||||
// We reach this point if and only if we were not connected to the node.
|
||||
entry.set_reserved(true);
|
||||
entry.force_outgoing();
|
||||
self.message_queue.push_back(Message::Connect(peer_id));
|
||||
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();
|
||||
}
|
||||
|
||||
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(mut peer) => {
|
||||
peer.set_reserved(false);
|
||||
peersstate::Peer::Connected(peer) => {
|
||||
if self.reserved_only {
|
||||
peer.disconnect();
|
||||
self.message_queue.push_back(Message::Drop(peer_id));
|
||||
}
|
||||
}
|
||||
peersstate::Peer::NotConnected(mut peer) => peer.set_reserved(false),
|
||||
peersstate::Peer::Unknown(_) => {}
|
||||
peersstate::Peer::NotConnected(_) => {},
|
||||
peersstate::Peer::Unknown(_) => {},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -215,20 +221,35 @@ impl Peerset {
|
||||
// Disconnect non-reserved nodes.
|
||||
self.reserved_only = reserved_only;
|
||||
if self.reserved_only {
|
||||
let reserved = self.data.get_priority_group(RESERVED_NODES).unwrap_or_default();
|
||||
for peer_id in self.data.connected_peers().cloned().collect::<Vec<_>>().into_iter() {
|
||||
let peer = self.data.peer(&peer_id).into_connected()
|
||||
.expect("We are enumerating connected peers, therefore the peer is connected; qed");
|
||||
if !peer.is_reserved() {
|
||||
if !reserved.contains(&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();
|
||||
}
|
||||
|
||||
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.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();
|
||||
}
|
||||
|
||||
fn on_report_peer(&mut self, peer_id: PeerId, score_diff: i32) {
|
||||
// We want reputations to be up-to-date before adjusting them.
|
||||
self.update_time();
|
||||
@@ -292,7 +313,13 @@ impl Peerset {
|
||||
self.update_time();
|
||||
|
||||
// Try to grab the next node to attempt to connect to.
|
||||
while let Some(next) = self.data.reserved_not_connected_peer() {
|
||||
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()
|
||||
}
|
||||
} {
|
||||
match next.try_outgoing() {
|
||||
Ok(conn) => self.message_queue.push_back(Message::Connect(conn.into_peer_id())),
|
||||
Err(_) => break, // No more slots available.
|
||||
@@ -421,6 +448,11 @@ impl Peerset {
|
||||
"message_queue": self.message_queue.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)
|
||||
}
|
||||
}
|
||||
|
||||
impl Stream for Peerset {
|
||||
@@ -439,6 +471,9 @@ impl Stream for Peerset {
|
||||
Action::RemoveReservedPeer(peer_id) => self.on_remove_reserved_peer(peer_id),
|
||||
Action::SetReservedOnly(reserved) => self.on_set_reserved_only(reserved),
|
||||
Action::ReportPeer(peer_id, score_diff) => self.on_report_peer(peer_id, score_diff),
|
||||
Action::SetPriorityGroup(group_id, peers) => self.on_set_priority_group(&group_id, peers),
|
||||
Action::AddToPriorityGroup(group_id, peer_id) => self.on_add_to_priority_group(&group_id, peer_id),
|
||||
Action::RemoveFromPriorityGroup(group_id, peer_id) => self.on_remove_from_priority_group(&group_id, peer_id),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -590,3 +625,4 @@ mod tests {
|
||||
tokio::runtime::current_thread::Runtime::new().unwrap().block_on(fut).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -17,7 +17,8 @@
|
||||
//! Contains the state storage behind the peerset.
|
||||
|
||||
use libp2p::PeerId;
|
||||
use std::{borrow::Cow, collections::HashMap};
|
||||
use std::{borrow::Cow, collections::{HashSet, HashMap}};
|
||||
use log::warn;
|
||||
|
||||
/// State storage behind the peerset.
|
||||
///
|
||||
@@ -35,17 +36,20 @@ pub struct PeersState {
|
||||
/// sort, to make the logic easier.
|
||||
nodes: HashMap<PeerId, Node>,
|
||||
|
||||
/// Number of non-reserved nodes for which the `ConnectionState` is `In`.
|
||||
/// Number of non-priority nodes for which the `ConnectionState` is `In`.
|
||||
num_in: u32,
|
||||
|
||||
/// Number of non-reserved nodes for which the `ConnectionState` is `In`.
|
||||
/// Number of non-priority nodes for which the `ConnectionState` is `In`.
|
||||
num_out: u32,
|
||||
|
||||
/// Maximum allowed number of non-reserved nodes for which the `ConnectionState` is `In`.
|
||||
/// Maximum allowed number of non-priority nodes for which the `ConnectionState` is `In`.
|
||||
max_in: u32,
|
||||
|
||||
/// Maximum allowed number of non-reserved nodes for which the `ConnectionState` is `Out`.
|
||||
/// Maximum allowed number of non-priority 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>>,
|
||||
}
|
||||
|
||||
/// State of a single node that we know about.
|
||||
@@ -54,14 +58,20 @@ struct Node {
|
||||
/// Whether we are connected to this node.
|
||||
connection_state: ConnectionState,
|
||||
|
||||
/// If true, this node is reserved and should always be connected to.
|
||||
reserved: bool,
|
||||
|
||||
/// Reputation value of the node, between `i32::min_value` (we hate that node) and
|
||||
/// `i32::max_value` (we love that node).
|
||||
reputation: i32,
|
||||
}
|
||||
|
||||
impl Default for Node {
|
||||
fn default() -> Node {
|
||||
Node {
|
||||
connection_state: ConnectionState::NotConnected,
|
||||
reputation: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether we are connected to a node.
|
||||
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
|
||||
enum ConnectionState {
|
||||
@@ -93,42 +103,30 @@ impl PeersState {
|
||||
num_out: 0,
|
||||
max_in: in_peers,
|
||||
max_out: out_peers,
|
||||
priority_nodes: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns an object that grants access to the state of a peer.
|
||||
pub fn peer<'a>(&'a mut self, peer_id: &'a PeerId) -> Peer<'a> {
|
||||
// Note: the Rust borrow checker still has some issues. In particular, we can't put this
|
||||
// block as an `else` below (as the obvious solution would be here), or it will complain
|
||||
// that we borrow `self` while it is already borrowed.
|
||||
if !self.nodes.contains_key(peer_id) {
|
||||
return Peer::Unknown(UnknownPeer {
|
||||
match self.nodes.get_mut(peer_id) {
|
||||
None => return Peer::Unknown(UnknownPeer {
|
||||
parent: self,
|
||||
peer_id: Cow::Borrowed(peer_id),
|
||||
});
|
||||
}
|
||||
|
||||
let state = self.nodes.get_mut(peer_id)
|
||||
.expect("We check that the value is present right above; QED");
|
||||
|
||||
if state.connection_state.is_connected() {
|
||||
Peer::Connected(ConnectedPeer {
|
||||
state,
|
||||
peer_id: Cow::Borrowed(peer_id),
|
||||
num_in: &mut self.num_in,
|
||||
num_out: &mut self.num_out,
|
||||
max_in: self.max_in,
|
||||
max_out: self.max_out,
|
||||
})
|
||||
} else {
|
||||
Peer::NotConnected(NotConnectedPeer {
|
||||
state,
|
||||
peer_id: Cow::Borrowed(peer_id),
|
||||
num_in: &mut self.num_in,
|
||||
num_out: &mut self.num_out,
|
||||
max_in: self.max_in,
|
||||
max_out: self.max_out,
|
||||
})
|
||||
}),
|
||||
Some(peer) => {
|
||||
if peer.connection_state.is_connected() {
|
||||
Peer::Connected(ConnectedPeer {
|
||||
state: self,
|
||||
peer_id: Cow::Borrowed(peer_id),
|
||||
})
|
||||
} else {
|
||||
Peer::NotConnected(NotConnectedPeer {
|
||||
state: self,
|
||||
peer_id: Cow::Borrowed(peer_id),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -148,28 +146,32 @@ impl PeersState {
|
||||
.map(|(p, _)| p)
|
||||
}
|
||||
|
||||
/// Returns the first reserved peer that we are not connected to.
|
||||
/// Returns the first priority peer that we are not connected to.
|
||||
///
|
||||
/// If multiple nodes are reserved, which one is returned is unspecified.
|
||||
pub fn reserved_not_connected_peer(&mut self) -> Option<NotConnectedPeer> {
|
||||
let outcome = self.nodes.iter_mut()
|
||||
.find(|(_, &mut Node { connection_state, reserved, .. })| {
|
||||
reserved && !connection_state.is_connected()
|
||||
})
|
||||
.map(|(peer_id, node)| (peer_id.clone(), node));
|
||||
/// 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),
|
||||
})
|
||||
}
|
||||
|
||||
if let Some((peer_id, state)) = outcome {
|
||||
Some(NotConnectedPeer {
|
||||
state,
|
||||
peer_id: Cow::Owned(peer_id),
|
||||
num_in: &mut self.num_in,
|
||||
num_out: &mut self.num_out,
|
||||
max_in: self.max_in,
|
||||
max_out: self.max_out,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
/// 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.
|
||||
@@ -187,21 +189,160 @@ impl PeersState {
|
||||
}
|
||||
Some(to_try)
|
||||
})
|
||||
.map(|(peer_id, state)| (peer_id.clone(), state));
|
||||
.map(|(peer_id, _)| peer_id.clone());
|
||||
|
||||
if let Some((peer_id, state)) = outcome {
|
||||
if let Some(peer_id) = outcome {
|
||||
Some(NotConnectedPeer {
|
||||
state,
|
||||
state: self,
|
||||
peer_id: Cow::Owned(peer_id),
|
||||
num_in: &mut self.num_in,
|
||||
num_out: &mut self.num_out,
|
||||
max_in: self.max_in,
|
||||
max_out: self.max_out,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
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;
|
||||
} 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 {
|
||||
// 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.
|
||||
let is_priority = self.is_priority(peer_id);
|
||||
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.
|
||||
///
|
||||
/// 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 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);
|
||||
// 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 => {},
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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()
|
||||
}
|
||||
|
||||
/// 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);
|
||||
}
|
||||
}
|
||||
|
||||
/// Grants access to the state of a peer in the `PeersState`.
|
||||
@@ -250,12 +391,8 @@ impl<'a> Peer<'a> {
|
||||
|
||||
/// A peer that is connected to us.
|
||||
pub struct ConnectedPeer<'a> {
|
||||
state: &'a mut Node,
|
||||
state: &'a mut PeersState,
|
||||
peer_id: Cow<'a, PeerId>,
|
||||
num_in: &'a mut u32,
|
||||
num_out: &'a mut u32,
|
||||
max_in: u32,
|
||||
max_out: u32,
|
||||
}
|
||||
|
||||
impl<'a> ConnectedPeer<'a> {
|
||||
@@ -266,87 +403,36 @@ impl<'a> ConnectedPeer<'a> {
|
||||
|
||||
/// Switches the peer to "not connected".
|
||||
pub fn disconnect(self) -> NotConnectedPeer<'a> {
|
||||
let connec_state = &mut self.state.connection_state;
|
||||
|
||||
if !self.state.reserved {
|
||||
match *connec_state {
|
||||
ConnectionState::In => *self.num_in -= 1,
|
||||
ConnectionState::Out => *self.num_out -= 1,
|
||||
ConnectionState::NotConnected =>
|
||||
debug_assert!(false, "State inconsistency: disconnecting a disconnected node")
|
||||
}
|
||||
}
|
||||
|
||||
*connec_state = ConnectionState::NotConnected;
|
||||
|
||||
self.state.disconnect(&self.peer_id);
|
||||
NotConnectedPeer {
|
||||
state: self.state,
|
||||
peer_id: self.peer_id,
|
||||
num_in: self.num_in,
|
||||
num_out: self.num_out,
|
||||
max_in: self.max_in,
|
||||
max_out: self.max_out,
|
||||
}
|
||||
}
|
||||
|
||||
/// Sets whether or not the node is reserved.
|
||||
pub fn set_reserved(&mut self, reserved: bool) {
|
||||
if reserved == self.state.reserved {
|
||||
return;
|
||||
}
|
||||
|
||||
if reserved {
|
||||
self.state.reserved = true;
|
||||
match self.state.connection_state {
|
||||
ConnectionState::In => *self.num_in -= 1,
|
||||
ConnectionState::Out => *self.num_out -= 1,
|
||||
ConnectionState::NotConnected => debug_assert!(false, "State inconsistency: \
|
||||
connected node is in fact not connected"),
|
||||
}
|
||||
|
||||
} else {
|
||||
self.state.reserved = false;
|
||||
match self.state.connection_state {
|
||||
ConnectionState::In => *self.num_in += 1,
|
||||
ConnectionState::Out => *self.num_out += 1,
|
||||
ConnectionState::NotConnected => debug_assert!(false, "State inconsistency: \
|
||||
connected node is in fact not connected"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns whether or not the node is reserved.
|
||||
pub fn is_reserved(&self) -> bool {
|
||||
self.state.reserved
|
||||
}
|
||||
|
||||
/// Returns the reputation value of the node.
|
||||
pub fn reputation(&self) -> i32 {
|
||||
self.state.reputation
|
||||
self.state.reputation(&self.peer_id)
|
||||
}
|
||||
|
||||
/// Sets the reputation of the peer.
|
||||
pub fn set_reputation(&mut self, value: i32) {
|
||||
self.state.reputation = value;
|
||||
self.state.set_reputation(&self.peer_id, value)
|
||||
}
|
||||
|
||||
/// 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) {
|
||||
let reputation = &mut self.state.reputation;
|
||||
*reputation = reputation.saturating_add(modifier);
|
||||
self.state.add_reputation(&self.peer_id, modifier)
|
||||
}
|
||||
}
|
||||
|
||||
/// A peer that is not connected to us.
|
||||
#[derive(Debug)]
|
||||
pub struct NotConnectedPeer<'a> {
|
||||
state: &'a mut Node,
|
||||
state: &'a mut PeersState,
|
||||
peer_id: Cow<'a, PeerId>,
|
||||
num_in: &'a mut u32,
|
||||
num_out: &'a mut u32,
|
||||
max_in: u32,
|
||||
max_out: u32,
|
||||
}
|
||||
|
||||
impl<'a> NotConnectedPeer<'a> {
|
||||
@@ -360,41 +446,16 @@ 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`.
|
||||
/// If the node is reserved, this method always succeeds.
|
||||
///
|
||||
/// Note that reserved nodes don't count towards the number of slots.
|
||||
/// Note that priority nodes don't count towards the number of slots.
|
||||
pub fn try_outgoing(self) -> Result<ConnectedPeer<'a>, NotConnectedPeer<'a>> {
|
||||
if self.is_reserved() {
|
||||
return Ok(self.force_outgoing())
|
||||
}
|
||||
|
||||
// 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, or if the user used
|
||||
// `force_outgoing`.
|
||||
if *self.num_out >= self.max_out {
|
||||
return Err(self);
|
||||
}
|
||||
|
||||
Ok(self.force_outgoing())
|
||||
}
|
||||
|
||||
/// Sets the peer as connected as an outgoing connection.
|
||||
pub fn force_outgoing(self) -> ConnectedPeer<'a> {
|
||||
let connec_state = &mut self.state.connection_state;
|
||||
debug_assert!(!connec_state.is_connected());
|
||||
*connec_state = ConnectionState::Out;
|
||||
|
||||
if !self.state.reserved {
|
||||
*self.num_out += 1;
|
||||
}
|
||||
|
||||
ConnectedPeer {
|
||||
state: self.state,
|
||||
peer_id: self.peer_id,
|
||||
num_in: self.num_in,
|
||||
num_out: self.num_out,
|
||||
max_in: self.max_in,
|
||||
max_out: self.max_out,
|
||||
if self.state.try_outgoing(&self.peer_id) {
|
||||
Ok(ConnectedPeer {
|
||||
state: self.state,
|
||||
peer_id: self.peer_id,
|
||||
})
|
||||
} else {
|
||||
Err(self)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -403,59 +464,26 @@ 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 reserved nodes don't count towards the number of slots.
|
||||
/// Note that priority nodes don't count towards the number of slots.
|
||||
pub fn try_accept_incoming(self) -> Result<ConnectedPeer<'a>, NotConnectedPeer<'a>> {
|
||||
if self.is_reserved() {
|
||||
return Ok(self.force_ingoing())
|
||||
if self.state.try_accept_incoming(&self.peer_id) {
|
||||
Ok(ConnectedPeer {
|
||||
state: self.state,
|
||||
peer_id: self.peer_id,
|
||||
})
|
||||
} else {
|
||||
Err(self)
|
||||
}
|
||||
|
||||
// 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 {
|
||||
return Err(self);
|
||||
}
|
||||
|
||||
Ok(self.force_ingoing())
|
||||
}
|
||||
|
||||
/// Sets the peer as connected as an ingoing connection.
|
||||
pub fn force_ingoing(self) -> ConnectedPeer<'a> {
|
||||
let connec_state = &mut self.state.connection_state;
|
||||
debug_assert!(!connec_state.is_connected());
|
||||
*connec_state = ConnectionState::In;
|
||||
|
||||
if !self.state.reserved {
|
||||
*self.num_in += 1;
|
||||
}
|
||||
|
||||
ConnectedPeer {
|
||||
state: self.state,
|
||||
peer_id: self.peer_id,
|
||||
num_in: self.num_in,
|
||||
num_out: self.num_out,
|
||||
max_in: self.max_in,
|
||||
max_out: self.max_out,
|
||||
}
|
||||
}
|
||||
|
||||
/// Sets whether or not the node is reserved.
|
||||
pub fn set_reserved(&mut self, reserved: bool) {
|
||||
self.state.reserved = reserved;
|
||||
}
|
||||
|
||||
/// Returns true if the the node is reserved.
|
||||
pub fn is_reserved(&self) -> bool {
|
||||
self.state.reserved
|
||||
}
|
||||
|
||||
/// Returns the reputation value of the node.
|
||||
pub fn reputation(&self) -> i32 {
|
||||
self.state.reputation
|
||||
self.state.reputation(&self.peer_id)
|
||||
}
|
||||
|
||||
/// Sets the reputation of the peer.
|
||||
pub fn set_reputation(&mut self, value: i32) {
|
||||
self.state.reputation = value;
|
||||
self.state.set_reputation(&self.peer_id, value)
|
||||
}
|
||||
|
||||
/// Performs an arithmetic addition on the reputation score of that peer.
|
||||
@@ -463,8 +491,7 @@ impl<'a> NotConnectedPeer<'a> {
|
||||
/// 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) {
|
||||
let reputation = &mut self.state.reputation;
|
||||
*reputation = reputation.saturating_add(modifier);
|
||||
self.state.add_reputation(&self.peer_id, modifier)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -477,25 +504,18 @@ pub struct UnknownPeer<'a> {
|
||||
impl<'a> UnknownPeer<'a> {
|
||||
/// Inserts the peer identity in our list.
|
||||
///
|
||||
/// The node is not reserved and starts with a reputation of 0. You can adjust these default
|
||||
/// The node starts with a reputation of 0. You can adjust these default
|
||||
/// values using the `NotConnectedPeer` that this method returns.
|
||||
pub fn discover(self) -> NotConnectedPeer<'a> {
|
||||
self.parent.nodes.insert(self.peer_id.clone().into_owned(), Node {
|
||||
connection_state: ConnectionState::NotConnected,
|
||||
reputation: 0,
|
||||
reserved: false,
|
||||
});
|
||||
|
||||
let state = self.parent.nodes.get_mut(&self.peer_id)
|
||||
.expect("We insert that key into the HashMap right above; QED");
|
||||
|
||||
let state = self.parent;
|
||||
NotConnectedPeer {
|
||||
state,
|
||||
peer_id: self.peer_id,
|
||||
num_in: &mut self.parent.num_in,
|
||||
num_out: &mut self.parent.num_out,
|
||||
max_in: self.parent.max_in,
|
||||
max_out: self.parent.max_out,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -521,14 +541,13 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reserved_node_doesnt_use_slot() {
|
||||
fn priority_node_doesnt_use_slot() {
|
||||
let mut peers_state = PeersState::new(1, 1);
|
||||
let id1 = PeerId::random();
|
||||
let id2 = PeerId::random();
|
||||
|
||||
if let Peer::Unknown(e) = peers_state.peer(&id1) {
|
||||
let mut p = e.discover();
|
||||
p.set_reserved(true);
|
||||
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());
|
||||
} else { panic!() }
|
||||
|
||||
@@ -550,23 +569,22 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reserved_not_connected_peer() {
|
||||
fn priority_not_connected_peer() {
|
||||
let mut peers_state = PeersState::new(25, 25);
|
||||
let id1 = PeerId::random();
|
||||
let id2 = PeerId::random();
|
||||
|
||||
assert!(peers_state.reserved_not_connected_peer().is_none());
|
||||
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.reserved_not_connected_peer().is_none());
|
||||
peers_state.peer(&id1).into_not_connected().unwrap().set_reserved(true);
|
||||
assert!(peers_state.reserved_not_connected_peer().is_some());
|
||||
peers_state.peer(&id2).into_not_connected().unwrap().set_reserved(true);
|
||||
peers_state.peer(&id1).into_not_connected().unwrap().set_reserved(false);
|
||||
assert!(peers_state.reserved_not_connected_peer().is_some());
|
||||
peers_state.peer(&id2).into_not_connected().unwrap().set_reserved(false);
|
||||
assert!(peers_state.reserved_not_connected_peer().is_none());
|
||||
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]
|
||||
@@ -581,7 +599,7 @@ mod tests {
|
||||
assert_eq!(peers_state.highest_not_connected_peer().map(|p| p.into_peer_id()), Some(id1.clone()));
|
||||
peers_state.peer(&id2).into_not_connected().unwrap().set_reputation(75);
|
||||
assert_eq!(peers_state.highest_not_connected_peer().map(|p| p.into_peer_id()), Some(id2.clone()));
|
||||
peers_state.peer(&id2).into_not_connected().unwrap().force_ingoing();
|
||||
peers_state.peer(&id2).into_not_connected().unwrap().try_accept_incoming().unwrap();
|
||||
assert_eq!(peers_state.highest_not_connected_peer().map(|p| p.into_peer_id()), Some(id1.clone()));
|
||||
peers_state.peer(&id1).into_not_connected().unwrap().set_reputation(100);
|
||||
peers_state.peer(&id2).into_connected().unwrap().disconnect();
|
||||
@@ -591,24 +609,31 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disconnect_reserved_doesnt_panic() {
|
||||
fn disconnect_priority_doesnt_panic() {
|
||||
let mut peers_state = PeersState::new(1, 1);
|
||||
let id = PeerId::random();
|
||||
let mut peer = peers_state.peer(&id).into_unknown().unwrap().discover()
|
||||
.force_outgoing();
|
||||
peer.set_reserved(true);
|
||||
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();
|
||||
peer.disconnect();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multiple_set_reserved_calls_doesnt_panic() {
|
||||
fn multiple_priority_groups_slot_count() {
|
||||
let mut peers_state = PeersState::new(1, 1);
|
||||
let id = PeerId::random();
|
||||
let mut peer = peers_state.peer(&id)
|
||||
.into_unknown().unwrap().discover()
|
||||
.force_outgoing();
|
||||
peer.set_reserved(true);
|
||||
peer.set_reserved(true);
|
||||
peer.disconnect();
|
||||
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user