Implement accurate memory accounting for mempool

8 files changed, 110 insertions, 2 deletions

diff --git a/src/memusage.h b/src/memusage.h
index 9f7de9e2e1..7a831e6d33 100644
--- a/src/memusage.h
+++ b/src/memusage.h
@@ -11,6 +11,7 @@
 #include <set>
 #include <vector>
 
+#include <boost/foreach.hpp>
 #include <boost/unordered_set.hpp>
 #include <boost/unordered_map.hpp>
 
@@ -20,12 +21,28 @@ namespace memusage
 /** Compute the total memory used by allocating alloc bytes. */
 static size_t MallocUsage(size_t alloc);
 
+/** Dynamic memory usage for built-in types is zero. */
+static inline size_t DynamicUsage(const int8_t& v) { return 0; }
+static inline size_t DynamicUsage(const uint8_t& v) { return 0; }
+static inline size_t DynamicUsage(const int16_t& v) { return 0; }
+static inline size_t DynamicUsage(const uint16_t& v) { return 0; }
+static inline size_t DynamicUsage(const int32_t& v) { return 0; }
+static inline size_t DynamicUsage(const uint32_t& v) { return 0; }
+static inline size_t DynamicUsage(const int64_t& v) { return 0; }
+static inline size_t DynamicUsage(const uint64_t& v) { return 0; }
+static inline size_t DynamicUsage(const float& v) { return 0; }
+static inline size_t DynamicUsage(const double& v) { return 0; }
+template<typename X> static inline size_t DynamicUsage(X * const &v) { return 0; }
+template<typename X> static inline size_t DynamicUsage(const X * const &v) { return 0; }
+template<typename X, typename Y> static inline size_t DynamicUsage(std::pair<X, Y> &p) { return 0; }
+
 /** Compute the memory used for dynamically allocated but owned data structures.
  *  For generic data types, this is *not* recursive. DynamicUsage(vector<vector<int> >)
  *  will compute the memory used for the vector<int>'s, but not for the ints inside.
  *  This is for efficiency reasons, as these functions are intended to be fast. If
  *  application data structures require more accurate inner accounting, they should
- *  do the recursion themselves, or use more efficient caching + updating on modification.
+ *  use RecursiveDynamicUsage, iterate themselves, or use more efficient caching +
+ *  updating on modification.
  */
 template<typename X> static size_t DynamicUsage(const std::vector<X>& v);
 template<typename X> static size_t DynamicUsage(const std::set<X>& s);
@@ -34,6 +51,12 @@ template<typename X, typename Y> static size_t DynamicUsage(const boost::unorder
 template<typename X, typename Y, typename Z> static size_t DynamicUsage(const boost::unordered_map<X, Y, Z>& s);
 template<typename X> static size_t DynamicUsage(const X& x);
 
+template<typename X> static size_t RecursiveDynamicUsage(const std::vector<X>& v);
+template<typename X> static size_t RecursiveDynamicUsage(const std::set<X>& v);
+template<typename X, typename Y> static size_t RecursiveDynamicUsage(const std::map<X, Y>& v);
+template<typename X, typename Y> static size_t RecursiveDynamicUsage(const std::pair<X, Y>& v);
+template<typename X> static size_t RecursiveDynamicUsage(const X& v);
+
 static inline size_t MallocUsage(size_t alloc)
 {
     // Measured on libc6 2.19 on Linux.
@@ -66,17 +89,53 @@ static inline size_t DynamicUsage(const std::vector<X>& v)
 }
 
 template<typename X>
+static inline size_t RecursiveDynamicUsage(const std::vector<X>& v)
+{
+    size_t usage = DynamicUsage(v);
+    BOOST_FOREACH(const X& x, v) {
+        usage += RecursiveDynamicUsage(x);
+    }
+    return usage;
+}
+
+template<typename X>
 static inline size_t DynamicUsage(const std::set<X>& s)
 {
     return MallocUsage(sizeof(stl_tree_node<X>)) * s.size();
 }
 
+template<typename X>
+static inline size_t RecursiveDynamicUsage(const std::set<X>& v)
+{
+    size_t usage = DynamicUsage(v);
+    BOOST_FOREACH(const X& x, v) {
+        usage += RecursiveDynamicUsage(x);
+    }
+    return usage;
+}
+
 template<typename X, typename Y>
 static inline size_t DynamicUsage(const std::map<X, Y>& m)
 {
     return MallocUsage(sizeof(stl_tree_node<std::pair<const X, Y> >)) * m.size();
 }
 
+template<typename X, typename Y>
+static inline size_t RecursiveDynamicUsage(const std::map<X, Y>& v)
+{
+    size_t usage = DynamicUsage(v);
+    for (typename std::map<X, Y>::const_iterator it = v.begin(); it != v.end(); it++) {
+        usage += RecursiveDynamicUsage(*it);
+    }
+    return usage;
+}
+
+template<typename X, typename Y>
+static inline size_t RecursiveDynamicUsage(const std::pair<X, Y>& v)
+{
+    return RecursiveDynamicUsage(v.first) + RecursiveDynamicUsage(v.second);
+}
+
 // Boost data structures
 
 template<typename X>
@@ -106,6 +165,12 @@ static inline size_t DynamicUsage(const X& x)
     return x.DynamicMemoryUsage();
 }
 
+template<typename X>
+static inline size_t RecursiveDynamicUsage(const X& x)
+{
+    return DynamicUsage(x);
+}
+
 }
 
 #endif
diff --git a/src/primitives/transaction.cpp b/src/primitives/transaction.cpp
index 606dbea798..7ed2d45973 100644
--- a/src/primitives/transaction.cpp
+++ b/src/primitives/transaction.cpp
@@ -72,6 +72,11 @@ void CTransaction::UpdateHash() const
     *const_cast<uint256*>(&hash) = SerializeHash(*this);
 }
 
+size_t CTransaction::DynamicMemoryUsage() const
+{
+    return memusage::RecursiveDynamicUsage(vin) + memusage::RecursiveDynamicUsage(vout);
+}
+
 CTransaction::CTransaction() : nVersion(CTransaction::CURRENT_VERSION), vin(), vout(), nLockTime(0) { }
 
 CTransaction::CTransaction(const CMutableTransaction &tx) : nVersion(tx.nVersion), vin(tx.vin), vout(tx.vout), nLockTime(tx.nLockTime) {
diff --git a/src/primitives/transaction.h b/src/primitives/transaction.h
index 6cfd93a9a1..77326c64b0 100644
--- a/src/primitives/transaction.h
+++ b/src/primitives/transaction.h
@@ -7,6 +7,7 @@
 #define BITCOIN_PRIMITIVES_TRANSACTION_H
 
 #include "amount.h"
+#include "memusage.h"
 #include "script/script.h"
 #include "serialize.h"
 #include "uint256.h"
@@ -48,6 +49,8 @@ public:
     }
 
     std::string ToString() const;
+
+    size_t DynamicMemoryUsage() const { return 0; }
 };
 
 /** An input of a transaction.  It contains the location of the previous
@@ -96,6 +99,8 @@ public:
     }
 
     std::string ToString() const;
+
+    size_t DynamicMemoryUsage() const { return scriptSig.DynamicMemoryUsage(); }
 };
 
 /** An output of a transaction.  It contains the public key that the next input
@@ -166,6 +171,8 @@ public:
     }
 
     std::string ToString() const;
+
+    size_t DynamicMemoryUsage() const { return scriptPubKey.DynamicMemoryUsage(); }
 };
 
 struct CMutableTransaction;
@@ -249,6 +256,8 @@ public:
     }
 
     std::string ToString() const;
+
+    size_t DynamicMemoryUsage() const;
 };
 
 /** A mutable version of CTransaction. */
diff --git a/src/rpcblockchain.cpp b/src/rpcblockchain.cpp
index 9cdd0770e3..c2de6cb244 100644
--- a/src/rpcblockchain.cpp
+++ b/src/rpcblockchain.cpp
@@ -758,6 +758,7 @@ UniValue getmempoolinfo(const UniValue& params, bool fHelp)
             "{\n"
             "  \"size\": xxxxx                (numeric) Current tx count\n"
             "  \"bytes\": xxxxx               (numeric) Sum of all tx sizes\n"
+            "  \"usage\": xxxxx               (numeric) Total memory usage for the mempool\n"
             "}\n"
             "\nExamples:\n"
             + HelpExampleCli("getmempoolinfo", "")
@@ -767,6 +768,7 @@ UniValue getmempoolinfo(const UniValue& params, bool fHelp)
     UniValue ret(UniValue::VOBJ);
     ret.push_back(Pair("size", (int64_t) mempool.size()));
     ret.push_back(Pair("bytes", (int64_t) mempool.GetTotalTxSize()));
+    ret.push_back(Pair("usage", (int64_t) mempool.DynamicMemoryUsage()));
 
     return ret;
 }
diff --git a/src/script/script.cpp b/src/script/script.cpp
index fd33924732..b1d2ceeb9f 100644
--- a/src/script/script.cpp
+++ b/src/script/script.cpp
@@ -260,3 +260,8 @@ std::string CScript::ToString() const
     }
     return str;
 }
+
+size_t CScript::DynamicMemoryUsage() const
+{
+    return memusage::DynamicUsage(*(static_cast<const std::vector<unsigned char>*>(this)));
+}
diff --git a/src/script/script.h b/src/script/script.h
index e39ca57f4f..aea34d05f4 100644
--- a/src/script/script.h
+++ b/src/script/script.h
@@ -6,6 +6,7 @@
 #ifndef BITCOIN_SCRIPT_SCRIPT_H
 #define BITCOIN_SCRIPT_SCRIPT_H
 
+#include "memusage.h"
 #include "crypto/common.h"
 
 #include <assert.h>
@@ -607,6 +608,8 @@ public:
         // The default std::vector::clear() does not release memory.
         std::vector<unsigned char>().swap(*this);
     }
+
+    size_t DynamicMemoryUsage() const;
 };
 
 class CReserveScript
diff --git a/src/txmempool.cpp b/src/txmempool.cpp
index 2292191be4..4caa5fc821 100644
--- a/src/txmempool.cpp
+++ b/src/txmempool.cpp
@@ -18,7 +18,7 @@
 using namespace std;
 
 CTxMemPoolEntry::CTxMemPoolEntry():
-    nFee(0), nTxSize(0), nModSize(0), nTime(0), dPriority(0.0), hadNoDependencies(false)
+    nFee(0), nTxSize(0), nModSize(0), nUsageSize(0), nTime(0), dPriority(0.0), hadNoDependencies(false)
 {
     nHeight = MEMPOOL_HEIGHT;
 }
@@ -31,6 +31,7 @@ CTxMemPoolEntry::CTxMemPoolEntry(const CTransaction& _tx, const CAmount& _nFee,
 {
     nTxSize = ::GetSerializeSize(tx, SER_NETWORK, PROTOCOL_VERSION);
     nModSize = tx.CalculateModifiedSize(nTxSize);
+    nUsageSize = tx.DynamicMemoryUsage();
 }
 
 CTxMemPoolEntry::CTxMemPoolEntry(const CTxMemPoolEntry& other)
@@ -101,6 +102,7 @@ bool CTxMemPool::addUnchecked(const uint256& hash, const CTxMemPoolEntry &entry,
         mapNextTx[tx.vin[i].prevout] = CInPoint(&tx, i);
     nTransactionsUpdated++;
     totalTxSize += entry.GetTxSize();
+    cachedInnerUsage += entry.DynamicMemoryUsage();
     minerPolicyEstimator->processTransaction(entry, fCurrentEstimate);
 
     return true;
@@ -146,6 +148,7 @@ void CTxMemPool::remove(const CTransaction &origTx, std::list<CTransaction>& rem
 
             removed.push_back(tx);
             totalTxSize -= mapTx[hash].GetTxSize();
+            cachedInnerUsage -= mapTx[hash].DynamicMemoryUsage();
             mapTx.erase(hash);
             nTransactionsUpdated++;
             minerPolicyEstimator->removeTx(hash);
@@ -226,6 +229,7 @@ void CTxMemPool::clear()
     mapTx.clear();
     mapNextTx.clear();
     totalTxSize = 0;
+    cachedInnerUsage = 0;
     ++nTransactionsUpdated;
 }
 
@@ -237,6 +241,7 @@ void CTxMemPool::check(const CCoinsViewCache *pcoins) const
     LogPrint("mempool", "Checking mempool with %u transactions and %u inputs\n", (unsigned int)mapTx.size(), (unsigned int)mapNextTx.size());
 
     uint64_t checkTotal = 0;
+    uint64_t innerUsage = 0;
 
     CCoinsViewCache mempoolDuplicate(const_cast<CCoinsViewCache*>(pcoins));
 
@@ -245,6 +250,7 @@ void CTxMemPool::check(const CCoinsViewCache *pcoins) const
     for (std::map<uint256, CTxMemPoolEntry>::const_iterator it = mapTx.begin(); it != mapTx.end(); it++) {
         unsigned int i = 0;
         checkTotal += it->second.GetTxSize();
+        innerUsage += it->second.DynamicMemoryUsage();
         const CTransaction& tx = it->second.GetTx();
         bool fDependsWait = false;
         BOOST_FOREACH(const CTxIn &txin, tx.vin) {
@@ -299,6 +305,7 @@ void CTxMemPool::check(const CCoinsViewCache *pcoins) const
     }
 
     assert(totalTxSize == checkTotal);
+    assert(innerUsage == cachedInnerUsage);
 }
 
 void CTxMemPool::queryHashes(vector<uint256>& vtxid)
@@ -419,3 +426,8 @@ bool CCoinsViewMemPool::GetCoins(const uint256 &txid, CCoins &coins) const {
 bool CCoinsViewMemPool::HaveCoins(const uint256 &txid) const {
     return mempool.exists(txid) || base->HaveCoins(txid);
 }
+
+size_t CTxMemPool::DynamicMemoryUsage() const {
+    LOCK(cs);
+    return memusage::DynamicUsage(mapTx) + memusage::DynamicUsage(mapNextTx) + memusage::DynamicUsage(mapDeltas) + cachedInnerUsage;
+}
diff --git a/src/txmempool.h b/src/txmempool.h
index 7271a5f603..ea36ce1ad5 100644
--- a/src/txmempool.h
+++ b/src/txmempool.h
@@ -40,6 +40,7 @@ private:
     CAmount nFee; //! Cached to avoid expensive parent-transaction lookups
     size_t nTxSize; //! ... and avoid recomputing tx size
     size_t nModSize; //! ... and modified size for priority
+    size_t nUsageSize; //! ... and total memory usage
     int64_t nTime; //! Local time when entering the mempool
     double dPriority; //! Priority when entering the mempool
     unsigned int nHeight; //! Chain height when entering the mempool
@@ -58,6 +59,7 @@ public:
     int64_t GetTime() const { return nTime; }
     unsigned int GetHeight() const { return nHeight; }
     bool WasClearAtEntry() const { return hadNoDependencies; }
+    size_t DynamicMemoryUsage() const { return nUsageSize; }
 };
 
 class CBlockPolicyEstimator;
@@ -73,6 +75,7 @@ public:
     CInPoint(const CTransaction* ptxIn, uint32_t nIn) { ptx = ptxIn; n = nIn; }
     void SetNull() { ptx = NULL; n = (uint32_t) -1; }
     bool IsNull() const { return (ptx == NULL && n == (uint32_t) -1); }
+    size_t DynamicMemoryUsage() const { return 0; }
 };
 
 /**
@@ -93,6 +96,7 @@ private:
     CBlockPolicyEstimator* minerPolicyEstimator;
 
     uint64_t totalTxSize; //! sum of all mempool tx' byte sizes
+    uint64_t cachedInnerUsage; //! sum of dynamic memory usage of all the map elements (NOT the maps themselves)
 
 public:
     mutable CCriticalSection cs;
@@ -139,6 +143,7 @@ public:
         LOCK(cs);
         return mapTx.size();
     }
+
     uint64_t GetTotalTxSize()
     {
         LOCK(cs);
@@ -162,6 +167,8 @@ public:
     /** Write/Read estimates to disk */
     bool WriteFeeEstimates(CAutoFile& fileout) const;
     bool ReadFeeEstimates(CAutoFile& filein);
+
+    size_t DynamicMemoryUsage() const;
 };
 
 /**