feat: remove nodes that have been connected within 15s

nervosnetwork · Aug 25, 2021 · 5f77b69 · 5f77b69
1 parent a2a5601
commit 5f77b69
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 7 deletions.
diff --git a/network/src/peer_store/peer_store_impl.rs b/network/src/peer_store/peer_store_impl.rs
@@ -136,7 +136,8 @@ impl PeerStore {
                 extract_peer_id(&peer_addr.addr)
                     .map(|peer_id| !peers.contains_key(&peer_id))
                     .unwrap_or_default()
-                    && peer_addr.had_connected(addr_expired_ms)
+                    && peer_addr
+                        .connected(|t| t > addr_expired_ms && t <= now_ms.saturating_sub(15_000))
             })
     }
 
@@ -157,7 +158,7 @@ impl PeerStore {
                     .map(|peer_id| !peers.contains_key(&peer_id))
                     .unwrap_or_default()
                     && !peer_addr.tried_in_last_minute(now_ms)
-                    && !peer_addr.had_connected(addr_expired_ms)
+                    && !peer_addr.connected(|t| t >= addr_expired_ms)
             })
     }
 
@@ -175,7 +176,7 @@ impl PeerStore {
                 extract_peer_id(&peer_addr.addr)
                     .map(|peer_id| peers.contains_key(&peer_id))
                     .unwrap_or_default()
-                    || peer_addr.had_connected(addr_expired_ms)
+                    || peer_addr.connected(|t| t >= addr_expired_ms)
             })
     }
 

diff --git a/network/src/peer_store/types.rs b/network/src/peer_store/types.rs
@@ -61,9 +61,9 @@ impl AddrInfo {
         }
     }
 
-    /// Whether already connected
-    pub fn had_connected(&self, expires_ms: u64) -> bool {
-        self.last_connected_at_ms > expires_ms
+    /// Connection information
+    pub fn connected<F: FnOnce(u64) -> bool>(&self, f: F) -> bool {
+        f(self.last_connected_at_ms)
     }
 
     /// Whether already try dail within a minute

diff --git a/network/src/tests/peer_store.rs b/network/src/tests/peer_store.rs
@@ -55,8 +55,14 @@ fn test_ban_peer() {
     assert!(peer_store.is_addr_banned(&addr));
 }
 
+#[cfg(not(disable_faketime))]
 #[test]
 fn test_attempt_ban() {
+    let faketime_file = faketime::millis_tempfile(0).expect("create faketime file");
+    faketime::enable(&faketime_file);
+
+    faketime::write_millis(&faketime_file, 1).expect("write millis");
+
     let mut peer_store: PeerStore = Default::default();
     let addr = random_addr();
     peer_store.add_addr(addr.clone()).unwrap();
@@ -65,13 +71,22 @@ fn test_attempt_ban() {
         .get_mut(&addr)
         .unwrap()
         .last_connected_at_ms = faketime::unix_time_as_millis();
+
+    faketime::write_millis(&faketime_file, 100_000).expect("write millis");
+
     assert_eq!(peer_store.fetch_addrs_to_attempt(2).len(), 1);
     peer_store.ban_addr(&addr, 10_000, "no reason".into());
     assert_eq!(peer_store.fetch_addrs_to_attempt(2).len(), 0);
 }
 
+#[cfg(not(disable_faketime))]
 #[test]
 fn test_fetch_addrs_to_attempt() {
+    let faketime_file = faketime::millis_tempfile(0).expect("create faketime file");
+    faketime::enable(&faketime_file);
+
+    faketime::write_millis(&faketime_file, 1).expect("write millis");
+
     let mut peer_store: PeerStore = Default::default();
     assert!(peer_store.fetch_addrs_to_attempt(1).is_empty());
     let addr = random_addr();
@@ -81,6 +96,8 @@ fn test_fetch_addrs_to_attempt() {
         .get_mut(&addr)
         .unwrap()
         .last_connected_at_ms = faketime::unix_time_as_millis();
+    faketime::write_millis(&faketime_file, 100_000).expect("write millis");
+
     assert_eq!(peer_store.fetch_addrs_to_attempt(2).len(), 1);
     peer_store.add_connected_peer(addr, SessionType::Outbound);
     assert!(peer_store.fetch_addrs_to_attempt(1).is_empty());
@@ -97,17 +114,27 @@ fn test_fetch_addrs_to_attempt_or_feeler() {
     let mut peer_store: PeerStore = Default::default();
     let addr = random_addr();
     peer_store.add_outbound_addr(addr);
+
+    faketime::write_millis(&faketime_file, 100_000).expect("write millis");
+
     assert_eq!(peer_store.fetch_addrs_to_attempt(2).len(), 1);
     assert!(peer_store.fetch_addrs_to_feeler(2).is_empty());
 
-    faketime::write_millis(&faketime_file, ADDR_TRY_TIMEOUT_MS + 1).expect("write millis");
+    faketime::write_millis(&faketime_file, 100_000 + ADDR_TRY_TIMEOUT_MS + 1)
+        .expect("write millis");
 
     assert!(peer_store.fetch_addrs_to_attempt(2).is_empty());
     assert_eq!(peer_store.fetch_addrs_to_feeler(2).len(), 1);
 }
 
+#[cfg(not(disable_faketime))]
 #[test]
 fn test_fetch_addrs_to_attempt_in_last_minutes() {
+    let faketime_file = faketime::millis_tempfile(0).expect("create faketime file");
+    faketime::enable(&faketime_file);
+
+    faketime::write_millis(&faketime_file, 100_000).expect("write millis");
+
     let mut peer_store: PeerStore = Default::default();
     let addr = random_addr();
     peer_store.add_addr(addr.clone()).unwrap();
@@ -127,6 +154,8 @@ fn test_fetch_addrs_to_attempt_in_last_minutes() {
         .get_mut(&addr)
         .unwrap()
         .last_connected_at_ms = now;
+    faketime::write_millis(&faketime_file, 200_000).expect("write millis");
+
     assert_eq!(peer_store.fetch_addrs_to_attempt(1).len(), 1);
     if let Some(paddr) = peer_store.mut_addr_manager().get_mut(&addr) {
         paddr.mark_tried(now);