refactor: use better expression

nervosnetwork · Dec 18, 2018 · df414b0 · df414b0
1 parent 9616a18
commit df414b0
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 57 deletions.
diff --git a/network/src/memory_peer_store.rs b/network/src/memory_peer_store.rs
@@ -86,12 +86,13 @@ impl PeerStore for MemoryPeerStore {
         if self.is_banned(peer_id) {
             return ReportResult::Banned;
         }
-        let behaviour_score = self.schema.get_score(behaviour);
-        if behaviour_score.is_none() {
-            debug!(target: "network", "behaviour {:?} is undefined", behaviour);
-            return ReportResult::Normal;
-        }
-        let behaviour_score = behaviour_score.unwrap();
+        let behaviour_score = match self.schema.get_score(behaviour) {
+            Some(score) => score,
+            None => {
+                debug!(target: "network", "behaviour {:?} is undefined", behaviour);
+                return ReportResult::Normal;
+            }
+        };
         // apply reported score
         let score = match self.peers.get_mut(peer_id) {
             Some(peer) => {

diff --git a/network/src/peers_registry.rs b/network/src/peers_registry.rs
@@ -155,21 +155,21 @@ pub(crate) struct PeersRegistry {
 fn find_most_peers_in_same_network_group<'a>(
     peers: impl Iterator<Item = (&'a PeerId, &'a PeerConnection)>,
 ) -> Vec<&'a PeerId> {
-    let mut groups: FnvHashMap<Group, Vec<&'a PeerId>> =
-        FnvHashMap::with_capacity_and_hasher(16, Default::default());
-    let mut largest_group_len = 0;
-    let mut largest_group: Group = Default::default();
-
-    for (peer_id, peer) in peers {
-        let group_name = peer.network_group();
-        let group = groups.entry(group_name.clone()).or_insert_with(Vec::new);
-        group.push(peer_id);
-        if group.len() > largest_group_len {
-            largest_group_len = group.len();
-            largest_group = group_name;
-        }
-    }
-    groups[&largest_group].clone()
+    peers
+        .fold(
+            FnvHashMap::with_capacity_and_hasher(16, Default::default()),
+            |mut groups, (peer_id, peer)| {
+                groups
+                    .entry(peer.network_group())
+                    .or_insert_with(Vec::new)
+                    .push(peer_id);
+                groups
+            },
+        )
+        .values()
+        .max_by_key(|group| group.len())
+        .cloned()
+        .unwrap_or_else(Vec::new)
 }
 
 impl PeersRegistry {
@@ -243,15 +243,16 @@ impl PeersRegistry {
     }
 
     fn try_evict_inbound_peer(&mut self, candidate_score: Score) -> bool {
-        let (peer_id, score) = {
+        let peer_id = {
             let inbound_peers = self
                 .peers
                 .iter()
                 .filter(|(peer_id, peer)| peer.is_inbound() && !self.is_reserved(peer_id));
             let candidate_peers = find_most_peers_in_same_network_group(inbound_peers);
             let peer_store = self.peer_store.read();
 
-            let mut lowest_score = std::i32::MAX;
+            // must have less score than candidate_score
+            let mut lowest_score = candidate_score - 1;
             let mut low_score_peers = Vec::new();
             for peer_id in candidate_peers {
                 if let Some(score) = peer_store.peer_score(peer_id) {
@@ -270,22 +271,15 @@ impl PeersRegistry {
                 return false;
             }
             let mut rng = thread_rng();
-            (
-                low_score_peers[..]
-                    .choose(&mut rng)
-                    .unwrap()
-                    .to_owned()
-                    .to_owned(),
-                lowest_score,
-            )
+            low_score_peers[..]
+                .choose(&mut rng)
+                .unwrap()
+                .to_owned()
+                .to_owned()
         };
-        if score < candidate_score {
-            debug!("evict inbound peer {:?}", peer_id);
-            self.drop_peer(&peer_id);
-            true
-        } else {
-            false
-        }
+        debug!("evict inbound peer {:?}", peer_id);
+        self.drop_peer(&peer_id);
+        true
     }
 
     pub fn try_outbound_peer(&mut self, peer_id: PeerId, addr: Multiaddr) -> Result<(), Error> {

diff --git a/network/src/tests/peers_registry.rs b/network/src/tests/peers_registry.rs
@@ -8,22 +8,6 @@ use ckb_util::RwLock;
 use std::default::Default;
 use std::sync::Arc;
 
-fn new_peers_registry(
-    peer_store: Arc<RwLock<Box<PeerStore>>>,
-    max_inbound: u32,
-    max_outbound: u32,
-    reserved_only: bool,
-    reserved_peers: Vec<PeerId>,
-) -> PeersRegistry {
-    PeersRegistry::new(
-        peer_store,
-        max_inbound,
-        max_outbound,
-        reserved_only,
-        reserved_peers,
-    )
-}
-
 #[test]
 fn test_accept_inbound_peer_in_reserve_only_mode() {
     let peer_store: Arc<RwLock<Box<PeerStore>>> = Arc::new(RwLock::new(Box::new(
@@ -33,7 +17,7 @@ fn test_accept_inbound_peer_in_reserve_only_mode() {
     let addr = "/ip4/127.0.0.1".to_multiaddr().unwrap();
 
     // reserved_only mode: only accept reserved_peer
-    let mut peers_registry = new_peers_registry(
+    let mut peers_registry = PeersRegistry::new(
         Arc::clone(&peer_store),
         3,
         3,
@@ -56,7 +40,7 @@ fn test_accept_inbound_peer_until_full() {
     let reserved_peer = random_peer_id().unwrap();
     let addr = "/ip4/127.0.0.1".to_multiaddr().unwrap();
     // accept node until inbound connections is full
-    let mut peers_registry = new_peers_registry(
+    let mut peers_registry = PeersRegistry::new(
         Arc::clone(&peer_store),
         3,
         3,
@@ -99,7 +83,7 @@ fn test_accept_inbound_peer_eviction() {
     let lowest_score_peer = random_peer_id().unwrap();
     let addr1 = "/ip4/127.0.0.1".to_multiaddr().unwrap();
     let addr2 = "/ip4/192.168.0.1".to_multiaddr().unwrap();
-    let mut peers_registry = new_peers_registry(
+    let mut peers_registry = PeersRegistry::new(
         Arc::clone(&peer_store),
         5,
         3,