added tests for prune_paths, added EndpointActorState::prune_paths method, and called it whenever new paths are added

Author: “ramfox”

iroh/src/magicsock/endpoint_map.rs

@@ -252,11 +252,6 @@ pub enum Source {
         _0: Private,
     },
     /// We established a connection on this address.
-    ///
-    /// Currently this means the path was in uses as [`PathId::ZERO`] when the a connection
-    /// was added to the `EndpointStateActor`.
-    ///
-    /// [`PathId::ZERO`]: quinn_proto::PathId::ZERO
     #[strum(serialize = "Connection")]
     Connection {
         /// private marker

iroh/src/magicsock/endpoint_map/endpoint_state.rs

@@ -1,5 +1,5 @@
 use std::{
-    collections::{BTreeSet, VecDeque},
+    collections::{BTreeSet, HashSet, VecDeque},
     net::SocketAddr,
     pin::Pin,
     sync::Arc,
@@ -410,6 +410,8 @@ impl EndpointStateActor {
                     .sources
                     .insert(Source::Connection { _0: Private }, Instant::now());
                 self.select_path();
+                // TODO(ramfox): do we need to prune paths here?
+                self.prune_paths();
 
                 if path_remote_is_ip {
                     // We may have raced this with a relay address.  Try and add any
@@ -453,6 +455,9 @@ impl EndpointStateActor {
                 .sources
                 .insert(source.clone(), Instant::now());
         }
+        // potentially prune inactive, unused paths now that we have added new
+        // paths
+        self.prune_paths();
         trace!("added addressing information");
     }
 
@@ -483,6 +488,9 @@ impl EndpointStateActor {
             self.send_disco_message(dst, disco::Message::Ping(ping))
                 .await;
         }
+        // prune any unused/inactive paths, now that we have added potential
+        // new ones
+        self.prune_paths()
     }
 
     /// Handles [`EndpointStateMessage::PingReceived`].
@@ -518,6 +526,7 @@ impl EndpointStateActor {
 
         trace!("ping received, triggering holepunching");
         self.trigger_holepunching().await;
+        // TODO(ramfox): potentially prune addrs here?
     }
 
     /// Handles [`EndpointStateMessage::PongReceived`].
@@ -707,6 +716,7 @@ impl EndpointStateActor {
             self.paths.entry(addr.clone()).or_default().ping_sent = Some(msg.tx_id);
             self.send_disco_message(addr, disco::Message::Ping(msg))
                 .await;
+            // TODO(ramfox): potentially prune addrs here?
         }
 
         // Send the DISCO CallMeMaybe message over the relay.
@@ -861,6 +871,7 @@ impl EndpointStateActor {
                 }
 
                 self.select_path();
+                self.prune_paths();
             }
             PathEvent::Abandoned { id, path_stats } => {
                 trace!(?path_stats, "path abandoned");
@@ -1012,6 +1023,25 @@ impl EndpointStateActor {
             }
         }
     }
+
+    fn prune_paths(&mut self) {
+        // if the total number of paths, relay or ip, is less
+        // than the max inactive ip addrs we allow, bail early
+        if self.paths.len() < MAX_INACTIVE_IP_ADDRESSES {
+            return;
+        }
+        let open_paths = self
+            .connections
+            .values()
+            .map(|state| state.open_paths.values())
+            .flatten();
+        prune_paths(
+            &mut self.paths,
+            &self.pending_open_paths,
+            &self.selected_path.get(),
+            open_paths,
+        );
+    }
 }
 
 /// Messages to send to the [`EndpointStateActor`].
@@ -1373,82 +1403,251 @@ impl Future for OnClosed {
     }
 }
 
-fn prune_paths(
+fn prune_paths<'a>(
     paths: &mut FxHashMap<transports::Addr, PathState>,
     pending: &VecDeque<transports::Addr>,
     selected_path: &Option<transports::Addr>,
-    open_paths: &Vec<transports::Addr>,
+    open_paths: impl Iterator<Item = &'a transports::Addr>,
 ) {
-    let ip_paths: BTreeSet<_> = paths
-        .keys()
-        .filter(|p| {
-            if p.is_ip() {
-                return true;
-            }
-            return false;
-        })
-        .cloned()
-        .collect();
+    let ip_count = paths.keys().filter(|p| p.is_ip()).count();
     // if the total number of ip paths is less than the allowed number of inactive
     // paths, just return early;
-    if ip_paths.len() < MAX_INACTIVE_IP_ADDRESSES {
+    if ip_count < MAX_INACTIVE_IP_ADDRESSES {
         return;
     }
 
-    let mut protected_paths = std::collections::BTreeSet::new();
+    let ip_paths: HashSet<_> = paths.keys().filter(|p| p.is_ip()).collect();
+
+    let mut protected_paths = HashSet::new();
     for addr in pending {
-        protected_paths.insert(addr.clone());
+        protected_paths.insert(addr);
     }
     if let Some(path) = selected_path {
-        protected_paths.insert(path.clone());
+        protected_paths.insert(path);
     }
     for path in open_paths {
-        protected_paths.insert(path.clone());
+        protected_paths.insert(path);
     }
 
-    let inactive_paths: Vec<_> = ip_paths.difference(&protected_paths).collect();
+    let inactive_paths: HashSet<_> = ip_paths
+        .difference(&protected_paths)
+        // cloned here so we can use `paths.retain` later
+        .map(|&addr| addr.clone())
+        .collect();
 
     if inactive_paths.len() < MAX_INACTIVE_IP_ADDRESSES {
         return;
     }
 
-    let mut keep_paths = Vec::new();
     let now = Instant::now();
-    // if the last instance in the source was CONST time ago, it can be pruned
-    for (addr, state) in paths {
-        if inactive_paths.contains(&&addr) {
-            let mut is_expired = true;
-            for (_source, instant) in &state.sources {
-                // it's been less than LAST_SOURCE_PRUNE_DURATION since we
-                // last learned about this source
-                if *instant + LAST_SOURCE_PRUNE_DURATION < now {
-                    is_expired = false;
-                    break;
-                }
-            }
-            if !is_expired {
-                keep_paths.push(addr);
-            }
-            continue;
+
+    paths.retain(|addr, state| {
+        if inactive_paths.contains(addr) {
+            keep_path(state, &now)
         } else {
-            keep_paths.push(addr);
+            // keep all active paths
+            true
         }
-    }
+    });
+}
 
-    *paths = paths
-        .iter()
-        .to_owned()
-        .filter(|(addr, _)| keep_paths.contains(addr))
-        .map(|(addr, state)| (addr.clone(), state.clone()))
-        .collect();
+/// Based on the [`PathState`], returns true if we should keep this path.
+///
+/// Currently we have two criteria:
+/// 1) This path has sent a Ping
+/// 2) The last time we learned about this address was greater than LAST_SOURCE_PRUNE_DURATION
+fn keep_path(state: &PathState, now: &Instant) -> bool {
+    // if we have never sent a ping, don't remove it
+    state.ping_sent.is_none()
+        || state
+            .sources
+            .values()
+            // only keep it if this path contains recent sources
+            .any(|instant| *instant + LAST_SOURCE_PRUNE_DURATION > *now)
 }
 
 #[cfg(test)]
 mod tests {
+    use super::super::Source;
+    use super::{PathState, prune_paths};
+    use crate::disco::TransactionId;
+    use crate::magicsock::{endpoint_map::Private, transports};
     use n0_error::Result;
+    use n0_future::time::{Duration, Instant};
+    use rustc_hash::FxHashMap;
+    use std::collections::VecDeque;
+    use std::net::{Ipv4Addr, SocketAddr};
+
+    /// Create a test IP address with specific port
+    fn test_ip_addr(port: u16) -> transports::Addr {
+        transports::Addr::Ip(SocketAddr::new(
+            std::net::IpAddr::V4(Ipv4Addr::LOCALHOST),
+            port,
+        ))
+    }
+
+    /// Create a PathState with sources at a specific time offset
+    fn test_path_state(time_offset: Duration, sent_ping: bool) -> PathState {
+        let mut state = PathState::default();
+        if sent_ping {
+            state.ping_sent = Some(TransactionId::default());
+        }
+        state.sources.insert(
+            Source::Connection { _0: Private },
+            Instant::now() - time_offset,
+        );
+        state
+    }
+
+    #[test]
+    fn test_prune_paths_too_few_total_paths() -> Result {
+        // create fewer than MAX_INACTIVE_IP_ADDRESSES paths
+        let mut paths = FxHashMap::default();
+        for i in 0..15 {
+            paths.insert(
+                test_ip_addr(i),
+                test_path_state(Duration::from_secs(0), false),
+            );
+        }
+
+        let pending = VecDeque::new();
+        let selected_path = None;
+        let open_paths = Vec::new();
+
+        let initial_len = paths.len();
+        // should not prune because we have fewer than MAX_INACTIVE_IP_ADDRESSES paths
+        prune_paths(&mut paths, &pending, &selected_path, open_paths.iter());
+        assert_eq!(
+            paths.len(),
+            initial_len,
+            "Expected no paths to be pruned when total IP paths < MAX_INACTIVE_IP_ADDRESSES"
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_prune_paths_too_few_inactive_paths() -> Result {
+        // create MAX_INACTIVE_IP_ADDRESSES + 5 paths
+        let mut paths = FxHashMap::default();
+        for i in 0..25 {
+            paths.insert(
+                test_ip_addr(i),
+                test_path_state(Duration::from_secs(0), false),
+            );
+        }
+
+        // mark 10 of them as "active" by adding them to open_paths
+        let open_paths: Vec<transports::Addr> = (0..10).map(|i| test_ip_addr(i)).collect();
+
+        let pending = VecDeque::new();
+        let selected_path = None;
+
+        let initial_len = paths.len();
+        // now we have 25 total paths, but only 15 inactive paths (25 - 10 = 15)
+        // which is less than MAX_INACTIVE_IP_ADDRESSES (20)
+        prune_paths(&mut paths, &pending, &selected_path, open_paths.iter());
+        assert_eq!(
+            paths.len(),
+            initial_len,
+            "Expected no paths to be pruned when inactive paths < MAX_INACTIVE_IP_ADDRESSES"
+        );
+
+        Ok(())
+    }
 
     #[test]
-    fn test_prune_paths() -> Result {
-        todo!();
+    fn test_prune_paths_prunes_old_inactive_paths() -> Result {
+        // create MAX_INACTIVE_IP_ADDRESSES + 10 paths
+        let mut paths = FxHashMap::default();
+
+        // add 20 paths with recent sources (within 2 minutes)
+        for i in 0..20 {
+            paths.insert(
+                test_ip_addr(i),
+                test_path_state(Duration::from_secs(60), true), // 1 minute ago
+            );
+        }
+
+        // add 10 paths with old sources (more than 2 minutes ago)
+        for i in 20..30 {
+            paths.insert(
+                test_ip_addr(i),
+                test_path_state(Duration::from_secs(180), true), // 3 minutes ago
+            );
+        }
+
+        let pending = VecDeque::new();
+        let selected_path = None;
+        let open_paths = Vec::new();
+
+        // we have 30 total paths, all inactive
+        // paths with sources older than LAST_SOURCE_PRUNE_DURATION should be pruned
+        prune_paths(&mut paths, &pending, &selected_path, open_paths.iter());
+
+        // we should have kept the 20 recent paths
+        assert_eq!(
+            paths.len(),
+            20,
+            "Expected to keep 20 paths with recent sources"
+        );
+
+        // verify that the kept paths are the ones with recent sources
+        for i in 0..20 {
+            let addr = test_ip_addr(i);
+            assert!(
+                paths.contains_key(&addr),
+                "Expected to keep path with recent source: {:?}",
+                addr
+            );
+        }
+
+        // verify that the old paths were removed
+        for i in 20..30 {
+            let addr = test_ip_addr(i);
+            assert!(
+                !paths.contains_key(&addr),
+                "Expected to prune path with old source: {:?}",
+                addr
+            );
+        }
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_prune_paths_protects_selected_and_open_paths() -> Result {
+        // create MAX_INACTIVE_IP_ADDRESSES + 10 paths, all with old sources
+        let mut paths = FxHashMap::default();
+        for i in 0..30 {
+            paths.insert(
+                test_ip_addr(i),
+                test_path_state(Duration::from_secs(180), true), // 3 minutes ago
+            );
+        }
+
+        let pending = VecDeque::new();
+        // mark one path as selected
+        let selected_path = Some(test_ip_addr(0));
+        // mark a few paths as open
+        let open_paths = vec![test_ip_addr(1), test_ip_addr(2)];
+
+        prune_paths(&mut paths, &pending, &selected_path, open_paths.iter());
+
+        // protected paths should still be in the result even though they have old sources
+        assert!(
+            paths.contains_key(&test_ip_addr(0)),
+            "Expected to keep selected path even with old source"
+        );
+        assert!(
+            paths.contains_key(&test_ip_addr(1)),
+            "Expected to keep open path even with old source"
+        );
+        assert!(
+            paths.contains_key(&test_ip_addr(2)),
+            "Expected to keep open path even with old source"
+        );
+
+        Ok(())
     }
 }