remove peer from connected_peers if connection fails

sigp · Jan 8, 2024 · baa38e9 · baa38e9
1 parent d355e42
commit baa38e9
Showing 1 changed file with 28 additions and 2 deletions.
diff --git a/protocols/gossipsub/src/behaviour.rs b/protocols/gossipsub/src/behaviour.rs
@@ -41,8 +41,8 @@ use libp2p_identity::PeerId;
 use libp2p_swarm::{
     behaviour::{AddressChange, ConnectionClosed, ConnectionEstablished, FromSwarm},
     dial_opts::DialOpts,
-    ConnectionDenied, ConnectionId, NetworkBehaviour, NotifyHandler, THandler, THandlerInEvent,
-    THandlerOutEvent, ToSwarm,
+    ConnectionDenied, ConnectionId, DialFailure, ListenFailure, NetworkBehaviour, NotifyHandler,
+    THandler, THandlerInEvent, THandlerOutEvent, ToSwarm,
 };
 
 use crate::peer_score::{PeerScore, PeerScoreParams, PeerScoreThresholds, RejectReason};
@@ -3269,6 +3269,32 @@ where
             FromSwarm::ConnectionClosed(connection_closed) => {
                 self.on_connection_closed(connection_closed)
             }
+            FromSwarm::DialFailure(DialFailure {
+                peer_id,
+                connection_id,
+                ..
+            })
+            | FromSwarm::ListenFailure(ListenFailure {
+                connection_id,
+                peer_id,
+                ..
+            }) => {
+                if let Some(peer_id) = peer_id {
+                    let mut peer = self
+                        .connected_peers
+                        .remove(&peer_id)
+                        .expect("To be connected to peer");
+                    let index = peer
+                        .connections
+                        .iter()
+                        .position(|v| v == &connection_id)
+                        .expect("Previously established connection to peer must be present");
+                    peer.connections.remove(index);
+                    if !peer.connections.is_empty() {
+                        self.connected_peers.insert(peer_id, peer);
+                    }
+                }
+            }
             FromSwarm::AddressChange(address_change) => self.on_address_change(address_change),
             _ => {}
         }