remove peer from connected_peers if connection fails

sigp · Jan 8, 2024 · b64bb06 · b64bb06
1 parent d355e42
commit b64bb06
Showing 1 changed file with 26 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,30 @@ 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 {
+                    if let Some(mut peer) = self.connected_peers.remove(&peer_id) {
+                        if let Some(index) =
+                            peer.connections.iter().position(|v| v == &connection_id)
+                        {
+                            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),
             _ => {}
         }