close duplicate libp2p connections

1 files changed, 59 insertions, 0 deletions

diff --git a/src/transport/libp2p.rs b/src/transport/libp2p.rs
index 58fbf64..077b9ac 100644
--- a/src/transport/libp2p.rs
+++ b/src/transport/libp2p.rs
@@ -21,6 +21,10 @@
 //! that works, or for good if it doesn't, within its limits (`RELAY_CIRCUIT_DURATION`,
 //! `RELAY_CIRCUIT_BYTES`; a circuit that ends is made again on the next send).
 //!
+//! Of several connections to one peer (both dialled at once, or AutoNAT dialled back),
+//! all but one are closed a few seconds later (`DUPLICATE_GRACE`, `Book::duplicates`),
+//! by the side with the smaller peer ID.
+//!
 //! A relay says so where it publishes its addresses: next to each, the same address
 //! ending in `/p2p-circuit` (where to reserve a slot), which other nodes find in the
 //! member list. A node that needs a relay and has none configured picks up to
@@ -113,6 +117,11 @@ const RELAY_PAUSE: Duration = Duration::from_secs(600);
 /// 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);
+/// Several connections to one peer (both dialled at once, or AutoNAT dialled back to
+/// see whether it can) are cut to one (`Book::duplicates`) this long after the last one
+/// opened: long enough for what it was dialled for. Looked for every `DEDUPE_EVERY`.
+const DUPLICATE_GRACE: Duration = Duration::from_secs(10);
+const DEDUPE_EVERY: Duration = Duration::from_secs(5);
 
 /// How long a message of `bytes` may take on a stream.
 fn stream_timeout(bytes: usize) -> Duration {
@@ -258,6 +267,20 @@ impl Book {
         found
     }
 
+    /// The connections to `peer` to close, of several: all but one, the oldest direct
+    /// one (the oldest relayed one if none is direct).
+    fn duplicates(&self, peer: &PeerId) -> Vec<ConnectionId> {
+        let mut connections: Vec<(bool, ConnectionId)> = self
+            .connections
+            .iter()
+            .filter(|(_, (connected, _, _))| connected == peer)
+            .map(|(id, (_, relayed, _))| (*relayed, *id))
+            .collect();
+        // Direct (false) first, then the oldest.
+        connections.sort();
+        connections.into_iter().skip(1).map(|(_, id)| id).collect()
+    }
+
     /// The remote address of a connection to `peer`: a direct one if there is one.
     fn location(&self, peer: &PeerId) -> Option<&Multiaddr> {
         let mut found = None;
@@ -873,8 +896,29 @@ async fn run(
     > = HashMap::new();
     let mut prune = tokio::time::interval(PRUNE_EVERY);
     prune.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
+    // Peers that may have several connections, and when to cut them to one. Only the
+    // side with the smaller peer ID does that, so the two never close different ones.
+    let mut dedupe_due: HashMap<PeerId, tokio::time::Instant> = HashMap::new();
+    let mut dedupe = tokio::time::interval(DEDUPE_EVERY);
+    dedupe.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
     loop {
         tokio::select! {
+            _ = dedupe.tick(), if !dedupe_due.is_empty() => {
+                let now = tokio::time::Instant::now();
+                let due: Vec<PeerId> = dedupe_due
+                    .iter()
+                    .filter(|(_, at)| **at <= now)
+                    .map(|(peer, _)| *peer)
+                    .collect();
+                for peer in due {
+                    dedupe_due.remove(&peer);
+                    let duplicates = lock(&book).duplicates(&peer);
+                    for connection in duplicates {
+                        debug!(peer = %peer, connection = %connection, "Closing a duplicate connection");
+                        swarm.close_connection(connection);
+                    }
+                }
+            }
             _ = prune.tick() => {
                 let forgotten = lock(&book).prune(tokio::time::Instant::now(), |peer| swarm.is_connected(peer));
                 if forgotten > 0 {
@@ -986,6 +1030,11 @@ async fn run(
                             connection_id,
                             (peer_id, relayed, endpoint.get_remote_address().clone()),
                         );
+                        if relaying.local < peer_id
+                            && book.connections.values().filter(|(connected, _, _)| *connected == peer_id).count() > 1
+                        {
+                            dedupe_due.insert(peer_id, tokio::time::Instant::now() + DUPLICATE_GRACE);
+                        }
                         debug!(
                             peer = %peer_id,
                             connection = %connection_id,
@@ -1888,6 +1937,16 @@ mod tests {
             book.location(&both).map(|address| address.to_string()),
             Some("/ip4/10.0.0.3/tcp/1".to_owned())
         );
+        // Of several, the oldest direct one stays; a relayed one only without a direct.
+        connect(&mut book, both, false);
+        assert_eq!(
+            book.duplicates(&both),
+            [
+                ConnectionId::new_unchecked(4),
+                ConnectionId::new_unchecked(2)
+            ]
+        );
+        assert!(book.duplicates(&relayed_only).is_empty());
         let connections = book.stats().connections;
         assert_eq!((connections.direct, connections.relayed), (1, 2));
     }