Ensure disconnecting peers are added to the peerdb (sigp#2451)

paulhauner · Jul 15, 2021 · 381befb · 381befb
1 parent 059d9ec
commit 381befb
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/beacon_node/eth2_libp2p/src/peer_manager/mod.rs b/beacon_node/eth2_libp2p/src/peer_manager/mod.rs
@@ -319,7 +319,7 @@ impl<TSpec: EthSpec> PeerManager<TSpec> {
             self.network_globals
                 .peers
                 .write()
-                .notify_disconnecting(&peer_id, true);
+                .notify_disconnecting(peer_id, true);
             return;
         }
 
@@ -339,7 +339,7 @@ impl<TSpec: EthSpec> PeerManager<TSpec> {
             self.network_globals
                 .peers
                 .write()
-                .notify_disconnecting(&peer_id, false);
+                .notify_disconnecting(peer_id, false);
             return;
         }
 
@@ -1023,7 +1023,7 @@ impl<TSpec: EthSpec> PeerManager<TSpec> {
 
         let mut peer_db = self.network_globals.peers.write();
         for peer_id in disconnecting_peers {
-            peer_db.notify_disconnecting(&peer_id, false);
+            peer_db.notify_disconnecting(peer_id, false);
             self.events.push(PeerManagerEvent::DisconnectPeer(
                 peer_id,
                 GoodbyeReason::TooManyPeers,

diff --git a/beacon_node/eth2_libp2p/src/peer_manager/peerdb.rs b/beacon_node/eth2_libp2p/src/peer_manager/peerdb.rs
@@ -475,10 +475,11 @@ impl<TSpec: EthSpec> PeerDB<TSpec> {
 
     /// Notifies the peer manager that the peer is undergoing a normal disconnect. Optionally tag
     /// the peer to be banned after the disconnect.
-    pub fn notify_disconnecting(&mut self, peer_id: &PeerId, to_ban_afterwards: bool) {
-        if let Some(info) = self.peers.get_mut(peer_id) {
-            info.disconnecting(to_ban_afterwards);
-        }
+    pub fn notify_disconnecting(&mut self, peer_id: PeerId, to_ban_afterwards: bool) {
+        self.peers
+            .entry(peer_id)
+            .or_default()
+            .disconnecting(to_ban_afterwards);
     }
 
     /// Marks a peer to be disconnected and then banned.