Merge pull request #2423 from TheBlueMatt/limitedmapalreadyaskedfor

Limited mapAlreadyAskedFor

4 files changed, 113 insertions, 4 deletions

diff --git a/src/limitedmap.h b/src/limitedmap.h
new file mode 100644
index 0000000000..7049d68e5a
--- /dev/null
+++ b/src/limitedmap.h
@@ -0,0 +1,100 @@
+// Copyright (c) 2012 The Bitcoin developers
+// Distributed under the MIT/X11 software license, see the accompanying
+// file COPYING or http://www.opensource.org/licenses/mit-license.php.
+#ifndef BITCOIN_LIMITEDMAP_H
+#define BITCOIN_LIMITEDMAP_H
+
+#include <map>
+#include <deque>
+
+/** STL-like map container that only keeps the N elements with the highest value. */
+template <typename K, typename V> class limitedmap
+{
+public:
+    typedef K key_type;
+    typedef V mapped_type;
+    typedef std::pair<const key_type, mapped_type> value_type;
+    typedef typename std::map<K, V>::const_iterator const_iterator;
+    typedef typename std::map<K, V>::size_type size_type;
+
+protected:
+    std::map<K, V> map;
+    typedef typename std::map<K, V>::iterator iterator;
+    std::multimap<V, iterator> rmap;
+    typedef typename std::multimap<V, iterator>::iterator rmap_iterator;
+    size_type nMaxSize;
+
+public:
+    limitedmap(size_type nMaxSizeIn = 0) { nMaxSize = nMaxSizeIn; }
+    const_iterator begin() const { return map.begin(); }
+    const_iterator end() const { return map.end(); }
+    size_type size() const { return map.size(); }
+    bool empty() const { return map.empty(); }
+    const_iterator find(const key_type& k) const { return map.find(k); }
+    size_type count(const key_type& k) const { return map.count(k); }
+    void insert(const value_type& x)
+    {
+        std::pair<iterator, bool> ret = map.insert(x);
+        if (ret.second)
+        {
+            if (nMaxSize && map.size() == nMaxSize)
+            {
+                map.erase(rmap.begin()->second);
+                rmap.erase(rmap.begin());
+            }
+            rmap.insert(make_pair(x.second, ret.first));
+        }
+        return;
+    }
+    void erase(const key_type& k)
+    {
+        iterator itTarget = map.find(k);
+        if (itTarget == map.end())
+            return;
+        std::pair<rmap_iterator, rmap_iterator> itPair = rmap.equal_range(itTarget->second);
+        for (rmap_iterator it = itPair.first; it != itPair.second; ++it)
+            if (it->second == itTarget)
+            {
+                rmap.erase(it);
+                map.erase(itTarget);
+                return;
+            }
+        // Shouldn't ever get here
+        assert(0); //TODO remove me
+        map.erase(itTarget);
+    }
+    void update(const_iterator itIn, const mapped_type& v)
+    {
+        //TODO: When we switch to C++11, use map.erase(itIn, itIn) to get the non-const iterator
+        iterator itTarget = map.find(itIn->first);
+        if (itTarget == map.end())
+            return;
+        std::pair<rmap_iterator, rmap_iterator> itPair = rmap.equal_range(itTarget->second);
+        for (rmap_iterator it = itPair.first; it != itPair.second; ++it)
+            if (it->second == itTarget)
+            {
+                rmap.erase(it);
+                itTarget->second = v;
+                rmap.insert(make_pair(v, itTarget));
+                return;
+            }
+        // Shouldn't ever get here
+        assert(0); //TODO remove me
+        itTarget->second = v;
+        rmap.insert(make_pair(v, itTarget));
+    }
+    size_type max_size() const { return nMaxSize; }
+    size_type max_size(size_type s)
+    {
+        if (s)
+            while (map.size() > s)
+            {
+                map.erase(rmap.begin()->second);
+                rmap.erase(rmap.begin());
+            }
+        nMaxSize = s;
+        return nMaxSize;
+    }
+};
+
+#endif
diff --git a/src/main.cpp b/src/main.cpp
index 4f5febc459..fb8f1b3b1f 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -3990,7 +3990,6 @@ bool SendMessages(CNode* pto, bool fSendTrickle)
                     pto->PushMessage("getdata", vGetData);
                     vGetData.clear();
                 }
-                mapAlreadyAskedFor[inv] = nNow;
             }
             pto->mapAskFor.erase(pto->mapAskFor.begin());
         }
diff --git a/src/net.cpp b/src/net.cpp
index f1ece2c2e9..5f8b5ba32b 100644
--- a/src/net.cpp
+++ b/src/net.cpp
@@ -53,7 +53,7 @@ CCriticalSection cs_vNodes;
 map<CInv, CDataStream> mapRelay;
 deque<pair<int64, CInv> > vRelayExpiration;
 CCriticalSection cs_mapRelay;
-map<CInv, int64> mapAlreadyAskedFor;
+limitedmap<CInv, int64> mapAlreadyAskedFor(MAX_INV_SZ);
 
 static deque<string> vOneShots;
 CCriticalSection cs_vOneShots;
diff --git a/src/net.h b/src/net.h
index 65e83f909c..30b9ac8663 100644
--- a/src/net.h
+++ b/src/net.h
@@ -15,6 +15,7 @@
 #endif
 
 #include "mruset.h"
+#include "limitedmap.h"
 #include "netbase.h"
 #include "protocol.h"
 #include "addrman.h"
@@ -79,7 +80,7 @@ extern CCriticalSection cs_vNodes;
 extern std::map<CInv, CDataStream> mapRelay;
 extern std::deque<std::pair<int64, CInv> > vRelayExpiration;
 extern CCriticalSection cs_mapRelay;
-extern std::map<CInv, int64> mapAlreadyAskedFor;
+extern limitedmap<CInv, int64> mapAlreadyAskedFor;
 
 extern std::vector<std::string> vAddedNodes;
 extern CCriticalSection cs_vAddedNodes;
@@ -346,7 +347,12 @@ public:
     {
         // We're using mapAskFor as a priority queue,
         // the key is the earliest time the request can be sent
-        int64& nRequestTime = mapAlreadyAskedFor[inv];
+        int64 nRequestTime;
+        limitedmap<CInv, int64>::const_iterator it = mapAlreadyAskedFor.find(inv);
+        if (it != mapAlreadyAskedFor.end())
+            nRequestTime = it->second;
+        else
+            nRequestTime = 0;
         if (fDebugNet)
             printf("askfor %s   %"PRI64d" (%s)\n", inv.ToString().c_str(), nRequestTime, DateTimeStrFormat("%H:%M:%S", nRequestTime/1000000).c_str());
 
@@ -359,6 +365,10 @@ public:
 
         // Each retry is 2 minutes after the last
         nRequestTime = std::max(nRequestTime + 2 * 60 * 1000000, nNow);
+        if (it != mapAlreadyAskedFor.end())
+            mapAlreadyAskedFor.update(it, nRequestTime);
+        else
+            mapAlreadyAskedFor.insert(std::make_pair(inv, nRequestTime));
         mapAskFor.insert(std::make_pair(nRequestTime, inv));
     }