From 381befbf827725ded7184591d8f017dca348db60 Mon Sep 17 00:00:00 2001
From: Age Manning <Age@AgeManning.com>
Date: Wed, 14 Jul 2021 12:59:24 +1000
Subject: [PATCH] Ensure disconnecting peers are added to the peerdb (#2451)

---
 beacon_node/eth2_libp2p/src/peer_manager/mod.rs    | 6 +++---
 beacon_node/eth2_libp2p/src/peer_manager/peerdb.rs | 9 +++++----
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/beacon_node/eth2_libp2p/src/peer_manager/mod.rs b/beacon_node/eth2_libp2p/src/peer_manager/mod.rs
index bfe8150f5..deb6de949 100644
--- 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
index 760d0245d..8819fdb19 100644
--- 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.