[net processing] Add RemovePeer()

jnewbery · jnewbery · commit 3025ca9e7743 · 2020-12-07T11:59:24.000Z
This allows us to avoid repeated locking in FinalizeNode()
diff --git a/src/net_processing.cpp b/src/net_processing.cpp
@@ -790,11 +790,9 @@ void PeerManager::FinalizeNode(const CNode& node, bool& fUpdateConnectionTime) {
     LOCK(cs_main);
     int misbehavior{0};
     {
-        PeerRef peer = GetPeerRef(nodeid);
+        PeerRef peer = RemovePeer(nodeid);
         assert(peer != nullptr);
         misbehavior = WITH_LOCK(peer->m_misbehavior_mutex, return peer->m_misbehavior_score);
-        LOCK(m_peer_mutex);
-        m_peer_map.erase(nodeid);
     }
     CNodeState *state = State(nodeid);
     assert(state != nullptr);
@@ -842,6 +840,18 @@ PeerRef PeerManager::GetPeerRef(NodeId id) const
     return it != m_peer_map.end() ? it->second : nullptr;
 }
 
+PeerRef PeerManager::RemovePeer(NodeId id)
+{
+    PeerRef ret;
+    LOCK(m_peer_mutex);
+    auto it = m_peer_map.find(id);
+    if (it != m_peer_map.end()) {
+        ret = std::move(it->second);
+        m_peer_map.erase(it);
+    }
+    return ret;
+}
+
 bool PeerManager::GetNodeStateStats(NodeId nodeid, CNodeStateStats &stats) {
     {
         LOCK(cs_main);
diff --git a/src/net_processing.h b/src/net_processing.h
@@ -143,6 +143,10 @@ class PeerManager final : public CValidationInterface, public NetEventsInterface
      *  May return an empty shared_ptr if the Peer object can't be found. */
     PeerRef GetPeerRef(NodeId id) const;
 
+    /** Get a shared pointer to the Peer object and remove it from m_peer_map.
+     *  May return an empty shared_ptr if the Peer object can't be found. */
+    PeerRef RemovePeer(NodeId id);
+
     /**
      * Potentially mark a node discouraged based on the contents of a BlockValidationState object
      *
