RAII wrapper for CNode*

PIVX-Project · Jun 23, 2020 · 9c9e55b · 9c9e55b
1 parent e92780d
commit 9c9e55b
Showing 1 changed file with 18 additions and 7 deletions.
diff --git a/src/net.cpp b/src/net.cpp
@@ -889,12 +889,23 @@ void CheckOffsetDisconnectedPeers(const CNetAddr& ip)
 
 static std::list<CNode*> vNodesDisconnected;
 
-static bool ReverseCompareNodeMinPingTime(CNode *a, CNode *b)
+class CNodeRef {
+public:
+    CNodeRef(CNode *pnode) : _pnode(pnode)  {_pnode->AddRef();}
+    ~CNodeRef() {_pnode->Release();}
+
+    CNode& operator *() const {return *_pnode;};
+    CNode* operator ->() const {return _pnode;};
+private:
+    CNode *_pnode;
+};
+
+static bool ReverseCompareNodeMinPingTime(const CNodeRef &a, const CNodeRef &b)
 {
     return a->nMinPingUsecTime > b->nMinPingUsecTime;
 }
 
-static bool ReverseCompareNodeTimeConnected(CNode *a, CNode *b)
+static bool ReverseCompareNodeTimeConnected(const CNodeRef &a, const CNodeRef &b)
 {
     return a->nTimeConnected > b->nTimeConnected;
 }
@@ -909,7 +920,7 @@ class CompareNetGroupKeyed
         GetRandBytes(vchSecretKey.data(), vchSecretKey.size());
     }
 
-    bool operator()(CNode *a, CNode *b)
+    bool operator()(const CNodeRef &a, const CNodeRef &b)
     {
         std::vector<unsigned char> vchGroupA, vchGroupB;
         CSHA256 hashA, hashB;
@@ -932,7 +943,7 @@ class CompareNetGroupKeyed
 };
 
 static bool AttemptToEvictConnection(bool fPreferNewConnection) {
-    std::vector<CNode*> vEvictionCandidates;
+    std::vector<CNodeRef> vEvictionCandidates;
     {
         LOCK(cs_vNodes);
 
@@ -945,7 +956,7 @@ static bool AttemptToEvictConnection(bool fPreferNewConnection) {
                 continue;
             if (node->addr.IsLocal())
                 continue;
-            vEvictionCandidates.push_back(node);
+            vEvictionCandidates.push_back(CNodeRef(node));
         }
     }
 
@@ -973,8 +984,8 @@ static bool AttemptToEvictConnection(bool fPreferNewConnection) {
     // Identify CNetAddr with the most connections
     CNetAddr naMostConnections;
     unsigned int nMostConnections = 0;
-    std::map<CNetAddr, std::vector<CNode*> > mapAddrCounts;
-    for (CNode *node : vEvictionCandidates) {
+    std::map<CNetAddr, std::vector<CNodeRef> > mapAddrCounts;
+    for (const CNodeRef& node : vEvictionCandidates) {
         mapAddrCounts[node->addr].push_back(node);
 
         if (mapAddrCounts[node->addr].size() > nMostConnections) {