forget old libp2p peers

1 files changed, 79 insertions, 0 deletions

diff --git a/src/transport/p2p.rs b/src/transport/p2p.rs
index 955c0f7..ebda922 100644
--- a/src/transport/p2p.rs
+++ b/src/transport/p2p.rs
@@ -99,6 +99,12 @@ const RELAY_CHECK: Duration = Duration::from_secs(30);
 /// answers at its address: a relay restarted without its libp2p key) isn't picked again
 /// for this long.
 const RELAY_PAUSE: Duration = Duration::from_secs(600);
+/// A peer nobody mentioned for this long (the member list, identify, mDNS, a
+/// connection) is forgotten, unless connected: mostly old peer IDs of nodes restarted
+/// without their libp2p key, whose new ones the member list has since. Members are
+/// mentioned every few seconds (`swarm::topology`). Checked every `PRUNE_EVERY`.
+const STALE_AFTER: Duration = Duration::from_secs(600);
+const PRUNE_EVERY: Duration = Duration::from_secs(60);
 
 /// How long a message of `bytes` may take on a stream.
 fn stream_timeout(bytes: usize) -> Duration {
@@ -187,6 +193,8 @@ struct Book {
     public_relays: HashMap<PeerId, Multiaddr>,
     /// Relays picked that never granted a slot, and when (`RELAY_PAUSE`).
     failed_relays: HashMap<PeerId, tokio::time::Instant>,
+    /// When each peer was last mentioned (`STALE_AFTER`).
+    seen: HashMap<PeerId, tokio::time::Instant>,
     /// Peers found on the LAN (mDNS) not known yet to speak continuum.
     lan: std::collections::HashSet<PeerId>,
     /// Peers that speak continuum (identify says so).
@@ -197,6 +205,44 @@ struct Book {
 }
 
 impl Book {
+    /// `peer` was mentioned just now.
+    fn touch(&mut self, peer: PeerId) {
+        self.seen.insert(peer, tokio::time::Instant::now());
+    }
+
+    /// Forgets the peers not mentioned within `STALE_AFTER` before `now`, but those
+    /// `connected`. Returns how many.
+    fn prune(&mut self, now: tokio::time::Instant, connected: impl Fn(&PeerId) -> bool) -> usize {
+        let stale: Vec<PeerId> = self
+            .addresses
+            .keys()
+            .chain(self.relays.keys())
+            .chain(self.public_relays.keys())
+            .chain(self.continuum.iter())
+            .chain(self.lan.iter())
+            .filter(|peer| {
+                !connected(peer)
+                    && self
+                        .seen
+                        .get(peer)
+                        .is_none_or(|at| now.saturating_duration_since(*at) >= STALE_AFTER)
+            })
+            .copied()
+            .collect::<std::collections::HashSet<_>>()
+            .into_iter()
+            .collect();
+        for peer in &stale {
+            self.addresses.remove(peer);
+            self.relays.remove(peer);
+            self.public_relays.remove(peer);
+            self.continuum.remove(peer);
+            self.lan.remove(peer);
+            self.failed_relays.remove(peer);
+            self.seen.remove(peer);
+        }
+        stale.len()
+    }
+
     /// `peer`, found on the LAN, speaks continuum: a seed, at the address we know first.
     fn discover(&mut self, peer: PeerId) {
         let Some(address) = self.addresses.get(&peer).and_then(|known| known.first()) else {
@@ -262,6 +308,7 @@ impl Book {
         if !self.addresses.contains_key(&peer) && self.addresses.len() >= MAX_PEERS {
             return;
         }
+        self.touch(peer);
         let addresses = self.addresses.entry(peer).or_default();
         addresses.retain(|known| *known != address);
         addresses.insert(0, address);
@@ -713,8 +760,16 @@ async fn run(
         ConnectionId,
         (Multiaddr, oneshot::Sender<Result<PeerId, String>>),
     > = HashMap::new();
+    let mut prune = tokio::time::interval(PRUNE_EVERY);
+    prune.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
     loop {
         tokio::select! {
+            _ = prune.tick() => {
+                let forgotten = lock(&book).prune(tokio::time::Instant::now(), |peer| swarm.is_connected(peer));
+                if forgotten > 0 {
+                    debug!(forgotten, "Forgot libp2p peers nobody mentions any more");
+                }
+            }
             _ = dht_bootstrap.tick(), if relaying.public => {
                 if let Some(dht) = swarm.behaviour_mut().kad.as_mut()
                     && let Err(error) = dht.bootstrap()
@@ -808,6 +863,7 @@ async fn run(
             },
             event = swarm.select_next_some() => match event {
                 SwarmEvent::ConnectionEstablished { peer_id, connection_id, .. } => {
+                    lock(&book).touch(peer_id);
                     for done in connecting.remove(&peer_id).into_iter().flatten() {
                         let _ = done.send(Ok(()));
                     }
@@ -1496,6 +1552,29 @@ mod tests {
         assert_eq!(book.discovered, [address.with_p2p(lan).unwrap()]);
     }
 
+    /// Peers nobody mentions any more are forgotten, but connected ones.
+    #[test]
+    fn stale_peers_are_forgotten() {
+        let start = tokio::time::Instant::now();
+        let (old, current, connected) = (peer(), peer(), peer());
+        let address: Multiaddr = "/ip6/fd00::1/tcp/9661".parse().unwrap();
+        let mut book = Book::default();
+        for peer in [old, current, connected] {
+            book.remember(peer, address.clone());
+        }
+        book.relays.insert(old, address.clone());
+        book.continuum.insert(old);
+        let later = start + STALE_AFTER;
+        // `current` was mentioned again since.
+        book.seen.insert(current, later);
+        assert_eq!(book.prune(later, |peer| *peer == connected), 1);
+        assert!(!book.addresses.contains_key(&old));
+        assert!(!book.relays.contains_key(&old) && !book.continuum.contains(&old));
+        assert!(book.addresses.contains_key(&current));
+        assert!(book.addresses.contains_key(&connected));
+        assert_eq!(book.prune(later, |_| false), 1); // `connected`, once it isn't.
+    }
+
     /// A relay doesn't dial the nodes it relays for through itself.
     #[test]
     fn no_circuits_through_ourselves() {