Performance: Store weighted transactions in a tree

qa/rpc-tests/mempool_limit.py

@@ -8,8 +8,10 @@ import sys; assert sys.version_info < (3,), ur"This script does not run under Py
 from test_framework.test_framework import BitcoinTestFramework
 from test_framework.util import (
     assert_equal,
+    get_coinbase_address,
     initialize_chain_clean,
     start_nodes,
+    wait_and_assert_operationid_status,
 )
 
 from decimal import Decimal
@@ -51,13 +53,15 @@ class MempoolLimit(BitcoinTestFramework):
         assert_equal(Decimal("10.00"), Decimal(self.nodes[2].z_gettotalbalance()['transparent']))
         assert_equal(Decimal("10.00"), Decimal(self.nodes[3].z_gettotalbalance()['transparent']))
 
-        taddr1 = self.nodes[0].getnewaddress()
+        zaddr1 = self.nodes[0].z_getnewaddress('sapling')
-        taddr2 = self.nodes[0].getnewaddress()
+        zaddr2 = self.nodes[0].z_getnewaddress('sapling')
-        taddr3 = self.nodes[0].getnewaddress()
+        zaddr3 = self.nodes[0].z_getnewaddress('sapling')
 
         print("Filling mempool...")
-        self.nodes[1].sendtoaddress(taddr1, 9.999)
+        opid1 = self.nodes[1].z_sendmany(get_coinbase_address(self.nodes[1]), [{"address": zaddr1, "amount": Decimal('9.999')}])
-        self.nodes[2].sendtoaddress(taddr2, 9.999)
+        wait_and_assert_operationid_status(self.nodes[1], opid1)
+        opid2 = self.nodes[2].z_sendmany(get_coinbase_address(self.nodes[2]), [{"address": zaddr2, "amount": Decimal('9.999')}])
+        wait_and_assert_operationid_status(self.nodes[2], opid2)
         self.sync_all()
 
         for i in range(0, 4):
@@ -66,12 +70,13 @@ class MempoolLimit(BitcoinTestFramework):
             assert_equal(2, len(mempool), "node {}".format(i))
 
         print("Adding one more transaction...")
-        self.nodes[3].sendtoaddress(taddr3, 9.999)
+        opid3 = self.nodes[3].z_sendmany(get_coinbase_address(self.nodes[3]), [{"address": zaddr3, "amount": Decimal('9.999')}])
+        wait_and_assert_operationid_status(self.nodes[3], opid3)
         # The mempools are no longer guaranteed to be in a consistent state, so we cannot sync
         sleep(5)
-        mempool_node3 = self.nodes[i].getrawmempool()
+        mempool_node3 = self.nodes[3].getrawmempool()
         print("Mempool for node 3: {}".format(mempool_node3))
-        assert_equal(3, len(mempool_node3), "node {}".format(i))
+        assert_equal(3, len(mempool_node3), "node {}".format(3))
 
         print("Checking mempool size...")
         # Due to the size limit, there should only be 2 transactions in the mempool
@@ -85,9 +90,11 @@ class MempoolLimit(BitcoinTestFramework):
 
         # The mempool sizes should be reset
         print("Checking mempool size reset after block mined...")
-        taddr4 = self.nodes[0].getnewaddress()
+        zaddr4 = self.nodes[0].z_getnewaddress('sapling')
-        self.nodes[0].sendtoaddress(taddr4, 9.999)
+        opid4 = self.nodes[0].z_sendmany(zaddr1, [{"address": zaddr4, "amount": Decimal('9.998')}])
-        self.nodes[0].sendtoaddress(taddr4, 9.999)
+        wait_and_assert_operationid_status(self.nodes[0], opid4)
+        opid5 = self.nodes[0].z_sendmany(zaddr2, [{"address": zaddr4, "amount": Decimal('9.998')}])
+        wait_and_assert_operationid_status(self.nodes[0], opid5)
         self.sync_all()
 
         for i in range(0, 4):

src/gtest/test_mempoollimit.cpp

@@ -52,33 +52,33 @@ TEST(MempoolLimitTests, RecentlyEvictedList_DoesNotContainAfterExpiry)
     EXPECT_FALSE(recentlyEvicted.contains(TX_ID3));
 }
 
-TEST(MempoolLimitTests, WeightedTransactionList_CheckSizeAfterDropping)
+TEST(MempoolLimitTests, WeightedTxTree_CheckSizeAfterDropping)
 {
     std::set<uint256> testedDropping;
     // Run the test until we have tested dropping each of the elements
     int trialNum = 0;
     while (testedDropping.size() < 3) {
-        WeightedTransactionList list(MIN_TX_COST * 2);
+        WeightedTxTree tree(MIN_TX_WEIGHT * 2);
-        EXPECT_EQ(0, list.getTotalCost());
+        EXPECT_EQ(0, tree.getTotalWeight().weight);
-        EXPECT_EQ(0, list.getTotalLowFeePenaltyCost());
+        EXPECT_EQ(0, tree.getTotalWeight().lowFeePenaltyWeight);
-        list.add(WeightedTxInfo(TX_ID1, MIN_TX_COST, MIN_TX_COST));
+        tree.add(WeightedTxInfo(TX_ID1, TxWeight(MIN_TX_WEIGHT, MIN_TX_WEIGHT)));
-        EXPECT_EQ(4000, list.getTotalCost());
+        EXPECT_EQ(4000, tree.getTotalWeight().weight);
-        EXPECT_EQ(4000, list.getTotalLowFeePenaltyCost());
+        EXPECT_EQ(4000, tree.getTotalWeight().lowFeePenaltyWeight);
-        list.add(WeightedTxInfo(TX_ID2, MIN_TX_COST, MIN_TX_COST));
+        tree.add(WeightedTxInfo(TX_ID2, TxWeight(MIN_TX_WEIGHT, MIN_TX_WEIGHT)));
-        EXPECT_EQ(8000, list.getTotalCost());
+        EXPECT_EQ(8000, tree.getTotalWeight().weight);
-        EXPECT_EQ(8000, list.getTotalLowFeePenaltyCost());
+        EXPECT_EQ(8000, tree.getTotalWeight().lowFeePenaltyWeight);
-        EXPECT_FALSE(list.maybeDropRandom(true).is_initialized());
+        EXPECT_FALSE(tree.maybeDropRandom().is_initialized());
-        list.add(WeightedTxInfo(TX_ID3, MIN_TX_COST, MIN_TX_COST + LOW_FEE_PENALTY));
+        tree.add(WeightedTxInfo(TX_ID3, TxWeight(MIN_TX_WEIGHT, MIN_TX_WEIGHT + LOW_FEE_PENALTY)));
-        EXPECT_EQ(12000, list.getTotalCost());
+        EXPECT_EQ(12000, tree.getTotalWeight().weight);
-        EXPECT_EQ(12000 + LOW_FEE_PENALTY, list.getTotalLowFeePenaltyCost());
+        EXPECT_EQ(12000 + LOW_FEE_PENALTY, tree.getTotalWeight().lowFeePenaltyWeight);
-        boost::optional<WeightedTxInfo> drop = list.maybeDropRandom(true);
+        boost::optional<uint256> drop = tree.maybeDropRandom();
         ASSERT_TRUE(drop.is_initialized());
-        uint256 txid = drop.get().txId;
+        uint256 txid = drop.get();
         std::cerr << "Trial " << trialNum++ << ": dropped " << txid.ToString() << std::endl;
         testedDropping.insert(txid);
         // Do not continue to test if a particular trial fails
-        ASSERT_EQ(8000, list.getTotalCost());
+        ASSERT_EQ(8000, tree.getTotalWeight().weight);
-        ASSERT_EQ(txid == TX_ID3 ? 8000 : 8000 + LOW_FEE_PENALTY, list.getTotalLowFeePenaltyCost());
+        ASSERT_EQ(txid == TX_ID3 ? 8000 : 8000 + LOW_FEE_PENALTY, tree.getTotalWeight().lowFeePenaltyWeight);
     }
     std::cerr << "All 3 scenarios tested in " << trialNum << " trials" << std::endl;
 }
@@ -99,8 +99,8 @@ TEST(MempoolLimitTests, WeightedTXInfo_FromTx)
         builder.AddSaplingOutput(sk.full_viewing_key().ovk, sk.default_address(), 25000, {});
         
         WeightedTxInfo info = WeightedTxInfo::from(builder.Build().GetTxOrThrow(), 10000);
-        EXPECT_EQ(MIN_TX_COST, info.cost);
+        EXPECT_EQ(MIN_TX_WEIGHT, info.txWeight.weight);
-        EXPECT_EQ(MIN_TX_COST, info.lowFeePenaltyCost);
+        EXPECT_EQ(MIN_TX_WEIGHT, info.txWeight.lowFeePenaltyWeight);
     }
     
     // Lower than standard fee
@@ -111,8 +111,8 @@ TEST(MempoolLimitTests, WeightedTXInfo_FromTx)
         builder.SetFee(9999);
 
         WeightedTxInfo info = WeightedTxInfo::from(builder.Build().GetTxOrThrow(), 9999);
-        EXPECT_EQ(MIN_TX_COST, info.cost);
+        EXPECT_EQ(MIN_TX_WEIGHT, info.txWeight.weight);
-        EXPECT_EQ(MIN_TX_COST + LOW_FEE_PENALTY, info.lowFeePenaltyCost);
+        EXPECT_EQ(MIN_TX_WEIGHT + LOW_FEE_PENALTY, info.txWeight.lowFeePenaltyWeight);
     }
 
     // Larger Tx
@@ -129,8 +129,8 @@ TEST(MempoolLimitTests, WeightedTXInfo_FromTx)
             std::cerr << result.GetError() << std::endl;
         }
         WeightedTxInfo info = WeightedTxInfo::from(result.GetTxOrThrow(), 10000);
-        EXPECT_EQ(5124, info.cost);
+        EXPECT_EQ(5124, info.txWeight.weight);
-        EXPECT_EQ(5124, info.lowFeePenaltyCost);
+        EXPECT_EQ(5124, info.txWeight.lowFeePenaltyWeight);
     }
     
     RegtestDeactivateSapling();

src/init.cpp

@@ -385,7 +385,7 @@ std::string HelpMessage(HelpMessageMode mode)
     strUsage += HelpMessageOpt("-maxreceivebuffer=<n>", strprintf(_("Maximum per-connection receive buffer, <n>*1000 bytes (default: %u)"), 5000));
     strUsage += HelpMessageOpt("-maxsendbuffer=<n>", strprintf(_("Maximum per-connection send buffer, <n>*1000 bytes (default: %u)"), 1000));
     strUsage += HelpMessageOpt("-mempoolevictionmemoryminutes=<n>", strprintf(_("The number of minutes before allowing rejected transactions to re-enter the mempool. (default: %u)"), DEFAULT_MEMPOOL_EVICTION_MEMORY_MINUTES));
-    strUsage += HelpMessageOpt("-mempooltotalcostlimit=<n>",strprintf(_("An upper bound on the maximum size in bytes of all transactions in the mempool. (default: %s)"), DEFAULT_MEMPOOL_TOTAL_COST_LIMIT));
+    strUsage += HelpMessageOpt("-mempooltotalcostlimit=<n>",strprintf(_("An upper bound on the maximum size in bytes of all transactions in the mempool. (default: %s)"), DEFAULT_MEMPOOL_TOTAL_WEIGHT_LIMIT));
     strUsage += HelpMessageOpt("-onion=<ip:port>", strprintf(_("Use separate SOCKS5 proxy to reach peers via Tor hidden services (default: %s)"), "-proxy"));
     strUsage += HelpMessageOpt("-onlynet=<net>", _("Only connect to nodes in network <net> (ipv4, ipv6 or onion)"));
     strUsage += HelpMessageOpt("-permitbaremultisig", strprintf(_("Relay non-P2SH multisig (default: %u)"), 1));
@@ -978,7 +978,7 @@ bool AppInit2(boost::thread_group& threadGroup, CScheduler& scheduler)
         mempool.setSanityCheck(1.0 / ratio);
     }
 
-    int64_t mempoolTotalCostLimit = GetArg("-mempooltotalcostlimit", DEFAULT_MEMPOOL_TOTAL_COST_LIMIT);
+    int64_t mempoolTotalCostLimit = GetArg("-mempooltotalcostlimit", DEFAULT_MEMPOOL_TOTAL_WEIGHT_LIMIT);
     int64_t mempoolEvictionMemorySeconds = GetArg("-mempoolevictionmemoryminutes", DEFAULT_MEMPOOL_EVICTION_MEMORY_MINUTES) * 60;
     mempool.setMempoolCostLimit(mempoolTotalCostLimit, mempoolEvictionMemorySeconds);
 

src/mempoollimit.cpp

@@ -10,6 +10,7 @@
 #include "version.h"
 
 const CAmount DEFAULT_FEE = 10000;
+const TxWeight ZERO_WEIGHT = TxWeight(0, 0);
 
 void RecentlyEvictedList::pruneList()
 {
@@ -17,15 +18,16 @@ void RecentlyEvictedList::pruneList()
         return;
     }
     int64_t now = GetAdjustedTime();
-    size_t startIndex = (txIdsAndTimesIndex + maxSize - txIdSet.size()) % maxSize;
+    size_t startIndex = (txIdsAndTimesIndex + capacity - txIdSet.size()) % capacity;
     boost::optional<std::pair<uint256, int64_t>> txIdAndTime;
     while ((txIdAndTime = txIdsAndTimes[startIndex]).is_initialized() && (now - txIdAndTime.get().second) > timeToKeep) {
         txIdsAndTimes[startIndex] = boost::none;
         txIdSet.erase(txIdAndTime.get().first);
-        startIndex = (startIndex + 1) % maxSize;
+        startIndex = (startIndex + 1) % capacity;
     }
 }
-void RecentlyEvictedList::add(uint256 txId)
+
+void RecentlyEvictedList::add(const uint256& txId)
 {
     pruneList();
     if (txIdsAndTimes[txIdsAndTimesIndex].is_initialized()) {
@@ -34,7 +36,7 @@ void RecentlyEvictedList::add(uint256 txId)
     }
     txIdsAndTimes[txIdsAndTimesIndex] = std::make_pair(txId, GetAdjustedTime());
     txIdSet.insert(txId);
-    txIdsAndTimesIndex = (txIdsAndTimesIndex + 1) % maxSize;
+    txIdsAndTimesIndex = (txIdsAndTimesIndex + 1) % capacity;
 }
 
 bool RecentlyEvictedList::contains(const uint256& txId)
@@ -44,61 +46,99 @@ bool RecentlyEvictedList::contains(const uint256& txId)
 }
 
 
-void WeightedTransactionList::clear() {
+TxWeight WeightedTxTree::getWeightAt(size_t index) const
-    weightedTxInfos.clear();
+{
+    return index < size ? txIdAndWeights[index].txWeight.add(childWeights[index]) : ZERO_WEIGHT;
 }
 
-int64_t WeightedTransactionList::getTotalCost() const
+void WeightedTxTree::backPropagate(size_t fromIndex, const TxWeight& weightDelta)
 {
-    return weightedTxInfos.empty() ? 0 : weightedTxInfos.back().cost;
+    while (fromIndex > 0) {
+        fromIndex = (fromIndex - 1) / 2;
+        childWeights[fromIndex] = childWeights[fromIndex].add(weightDelta);
+    }
 }
 
-int64_t WeightedTransactionList::getTotalLowFeePenaltyCost() const
+size_t WeightedTxTree::findByWeight(size_t fromIndex, uint64_t weightToFind) const
 {
-    return weightedTxInfos.empty() ? 0 : weightedTxInfos.back().lowFeePenaltyCost;
+    int leftWeight = getWeightAt(fromIndex * 2 + 1).lowFeePenaltyWeight;
+    int rightWeight = getWeightAt(fromIndex).lowFeePenaltyWeight - getWeightAt(fromIndex * 2 + 2).lowFeePenaltyWeight;
+    // On Left
+    if (weightToFind < leftWeight) {
+        return findByWeight(fromIndex * 2 + 1, weightToFind);
+    }
+    // Found
+    if (weightToFind < rightWeight) {
+        return fromIndex;
+    }
+    // On Right
+    return findByWeight(fromIndex * 2 + 2, weightToFind - rightWeight);
 }
 
-void WeightedTransactionList::add(WeightedTxInfo weightedTxInfo)
+TxWeight WeightedTxTree::getTotalWeight() const
 {
-    if (weightedTxInfos.empty()) {
+    return getWeightAt(0);
-        weightedTxInfos.push_back(weightedTxInfo);
+}
+
+
+void WeightedTxTree::add(const WeightedTxInfo& weightedTxInfo)
+{
+    txIdAndWeights.push_back(weightedTxInfo);
+    childWeights.push_back(ZERO_WEIGHT);
+    txIdToIndexMap[weightedTxInfo.txId] = size;
+    backPropagate(size++, weightedTxInfo.txWeight);
+}
+
+void WeightedTxTree::remove(const uint256& txId)
+{
+    if (txIdToIndexMap.find(txId) == txIdToIndexMap.end()) {
         return;
     }
-    weightedTxInfo.plusEquals(weightedTxInfos.back());
+
-    weightedTxInfos.push_back(weightedTxInfo);
+    size_t removeIndex = txIdToIndexMap[txId];
-    for (int i =0; i < weightedTxInfos.size(); ++i) {
+
-        WeightedTxInfo info = weightedTxInfos[i];
+    TxWeight lastChildWeight = txIdAndWeights[--size].txWeight;
+    backPropagate(size, lastChildWeight.negate());
+
+    if (removeIndex < size) {
+        TxWeight weightDelta = lastChildWeight.add(txIdAndWeights[removeIndex].txWeight.negate());
+        txIdAndWeights[removeIndex] = txIdAndWeights[size];
+        txIdToIndexMap[txIdAndWeights[removeIndex].txId] = removeIndex;
+        backPropagate(removeIndex, weightDelta);
     }
+
+    txIdToIndexMap.erase(txId);
+    txIdAndWeights.pop_back();
+    childWeights.pop_back();
 }
- 
+
-boost::optional<WeightedTxInfo> WeightedTransactionList::maybeDropRandom(bool rebuildList)
+boost::optional<uint256> WeightedTxTree::maybeDropRandom()
 {
-    int64_t totalCost = getTotalCost();
+    uint64_t totalPenaltyWeight = getTotalWeight().lowFeePenaltyWeight;
-    if (totalCost <= maxTotalCost) {
+    if (totalPenaltyWeight <= capacity) {
         return boost::none;
     }
-    LogPrint("mempool", "Mempool cost limit exceeded (cost=%d, limit=%d)\n", totalCost, maxTotalCost);
+    LogPrint("mempool", "Mempool cost limit exceeded (cost=%d, limit=%d)\n", totalPenaltyWeight, capacity);
-    int randomWeight = GetRand(getTotalLowFeePenaltyCost());
+    int randomWeight = GetRand(totalPenaltyWeight);
-    int i = 0;
+    WeightedTxInfo drop = txIdAndWeights[findByWeight(0, randomWeight)];
-    while (randomWeight > weightedTxInfos[i].lowFeePenaltyCost) {
+    LogPrint("mempool", "Evicting transaction (txid=%s, cost=%d, penaltyCost=%d)\n",
-        ++i;
+        drop.txId.ToString(), drop.txWeight.weight, drop.txWeight.lowFeePenaltyWeight);
-    }
+    remove(drop.txId);
-    WeightedTxInfo drop = weightedTxInfos[i];
+    return drop.txId;
-    if (i > 0) {
-        drop.minusEquals(weightedTxInfos[i - 1]);
-    }
-    if (rebuildList) {
-        while (++i < weightedTxInfos.size()) {
-            WeightedTxInfo nextTx = weightedTxInfos[i];
-            nextTx.minusEquals(drop);
-            weightedTxInfos[i - 1] = nextTx;
-        }
-        weightedTxInfos.pop_back();
-    }
-    LogPrint("mempool", "Evicting transaction (txid=%s, cost=%d, penaltyCost=%d)\n", drop.txId.ToString(), drop.cost, drop.lowFeePenaltyCost);
-    return drop;
 }
 
+
+TxWeight TxWeight::add(const TxWeight& other) const
+{
+    return TxWeight(weight + other.weight, lowFeePenaltyWeight + other.lowFeePenaltyWeight);
+}
+
+TxWeight TxWeight::negate() const
+{
+    return TxWeight(-weight, -lowFeePenaltyWeight);
+}
+
+
 // These are also defined in rpcwallet.cpp
 #define JOINSPLIT_SIZE GetSerializeSize(JSDescription(), SER_NETWORK, PROTOCOL_VERSION)
 #define OUTPUTDESCRIPTION_SIZE GetSerializeSize(OutputDescription(), SER_NETWORK, PROTOCOL_VERSION)
@@ -110,22 +150,10 @@ WeightedTxInfo WeightedTxInfo::from(const CTransaction& tx, const CAmount& fee)
     memUsage += tx.vJoinSplit.size() * JOINSPLIT_SIZE;
     memUsage += tx.vShieldedOutput.size() * OUTPUTDESCRIPTION_SIZE;
     memUsage += tx.vShieldedSpend.size() * SPENDDESCRIPTION_SIZE;
-    int64_t cost = std::max(memUsage, MIN_TX_COST);
+    uint64_t cost = std::max(memUsage, MIN_TX_WEIGHT);
-    int64_t lowFeePenaltyCost = cost;
+    uint64_t lowFeePenaltyCost = cost;
     if (fee < DEFAULT_FEE) {
         lowFeePenaltyCost += LOW_FEE_PENALTY;
     }
-    return WeightedTxInfo(tx.GetHash(), cost, lowFeePenaltyCost);
+    return WeightedTxInfo(tx.GetHash(), TxWeight(cost, lowFeePenaltyCost));
-}
-
-void WeightedTxInfo::plusEquals(const WeightedTxInfo& other)
-{
-    cost += other.cost;
-    lowFeePenaltyCost += other.lowFeePenaltyCost;
-}
-
-void WeightedTxInfo::minusEquals(const WeightedTxInfo& other)
-{
-    cost -= other.cost;
-    lowFeePenaltyCost -= other.lowFeePenaltyCost;
 }

src/mempoollimit.h

@@ -5,77 +5,91 @@
 #ifndef MEMPOOLLIMIT_H
 #define MEMPOOLLIMIT_H
 
-#include <vector>
+#include <map>
 #include <set>
-#include "uint256.h"
+#include <vector>
-#include "primitives/transaction.h"
 
 #include "boost/optional.hpp"
+#include "primitives/transaction.h"
+#include "uint256.h"
 
-const size_t DEFAULT_MEMPOOL_TOTAL_COST_LIMIT = 80000000;
+const size_t DEFAULT_MEMPOOL_TOTAL_WEIGHT_LIMIT = 80000000;
 const int64_t DEFAULT_MEMPOOL_EVICTION_MEMORY_MINUTES = 60;
 
 const size_t RECENTLY_EVICTED_SIZE = 10000;
-const uint64_t MIN_TX_COST = 4000;
+const uint64_t MIN_TX_WEIGHT = 4000;
 const uint64_t LOW_FEE_PENALTY = 16000; 
 
-struct WeightedTxInfo;
-
 class RecentlyEvictedList
 {
-    const size_t maxSize;
+    const size_t capacity;
+    size_t txIdsAndTimesIndex = 0;
+
     const int64_t timeToKeep;
     // Pairs of txid and time (seconds since epoch)
     boost::optional<std::pair<uint256, int64_t>> txIdsAndTimes[RECENTLY_EVICTED_SIZE];
-    size_t txIdsAndTimesIndex;
     std::set<uint256> txIdSet;
 
     void pruneList();
+
 public:
-    RecentlyEvictedList(size_t maxSize_, int64_t timeToKeep_) :
+    RecentlyEvictedList(size_t capacity_, int64_t timeToKeep_) : capacity(capacity_), timeToKeep(timeToKeep_) 
-        maxSize(maxSize_),
-        timeToKeep(timeToKeep_),
-        txIdsAndTimesIndex(0)    
     {
-        assert(maxSize <= RECENTLY_EVICTED_SIZE);
+        assert(capacity <= RECENTLY_EVICTED_SIZE);
-        std::fill_n(txIdsAndTimes, maxSize, boost::none);
+        std::fill_n(txIdsAndTimes, capacity, boost::none);
     }
     RecentlyEvictedList(int64_t timeToKeep_) : RecentlyEvictedList(RECENTLY_EVICTED_SIZE, timeToKeep_) {}
 
-    void add(uint256 txId);
+    void add(const uint256& txId);
     bool contains(const uint256& txId);
 };
 
 
-class WeightedTransactionList
+struct TxWeight {
-{
+    uint64_t weight;
-    const uint64_t maxTotalCost;
+    uint64_t lowFeePenaltyWeight;
-    std::vector<WeightedTxInfo> weightedTxInfos;
-public:
-    WeightedTransactionList(int64_t maxTotalCost_) : maxTotalCost(maxTotalCost_) {}
 
-    void clear();
+    TxWeight(uint64_t weight_, uint64_t lowFeePenaltyWeight_)
+        : weight(weight_), lowFeePenaltyWeight(lowFeePenaltyWeight_) {}
 
-    int64_t getTotalCost() const;
+    TxWeight add(const TxWeight& other) const;
-    int64_t getTotalLowFeePenaltyCost() const;
+    TxWeight negate() const;
-
-    void add(WeightedTxInfo weightedTxInfo);
-    boost::optional<WeightedTxInfo> maybeDropRandom(bool rebuildList);
 };
 
 
 struct WeightedTxInfo {
     uint256 txId;
-    uint64_t cost;
+    TxWeight txWeight;
-    uint64_t lowFeePenaltyCost;
 
-    WeightedTxInfo(uint256 txId_, uint64_t cost_, uint64_t lowFeePenaltyCost_)
+    WeightedTxInfo(uint256 txId_, TxWeight txWeight_) : txId(txId_), txWeight(txWeight_) {}
-        : txId(txId_), cost(cost_), lowFeePenaltyCost(lowFeePenaltyCost_) {}
 
     static WeightedTxInfo from(const CTransaction& tx, const CAmount& fee);
-    
-    void plusEquals(const WeightedTxInfo& other);
-    void minusEquals(const WeightedTxInfo& other);
 };
 
+
+class WeightedTxTree
+{
+    const uint64_t capacity;
+    size_t size = 0;
+    
+    std::vector<WeightedTxInfo> txIdAndWeights;
+    std::vector<TxWeight> childWeights;
+    std::map<uint256, size_t> txIdToIndexMap;
+
+    TxWeight getWeightAt(size_t index) const;
+    void backPropagate(size_t fromIndex, const TxWeight& weightDelta);
+    size_t findByWeight(size_t fromIndex, uint64_t weightToFind) const;
+
+public:
+    WeightedTxTree(uint64_t capacity_) : capacity(capacity_) {}
+
+    TxWeight getTotalWeight() const;
+
+    void add(const WeightedTxInfo& weightedTxInfo);
+    void remove(const uint256& txId);
+
+    boost::optional<uint256> maybeDropRandom();
+};
+
+
 #endif // MEMPOOLLIMIT_H

src/txmempool.cpp

@@ -69,7 +69,7 @@ CTxMemPool::~CTxMemPool()
 {
     delete minerPolicyEstimator;
     delete recentlyEvicted;
-    delete weightedTxList;
+    delete weightedTxTree;
 }
 
 void CTxMemPool::pruneSpent(const uint256 &hashTx, CCoins &coins)
@@ -104,9 +104,8 @@ bool CTxMemPool::addUnchecked(const uint256& hash, const CTxMemPoolEntry &entry,
     // Used by main.cpp AcceptToMemoryPool(), which DOES do
     // all the appropriate checks.
     LOCK(cs);
-    if (weightedTxList) {
+    assert(weightedTxTree);
-        weightedTxList->add(WeightedTxInfo::from(entry.GetTx(), entry.GetFee()));
+    weightedTxTree->add(WeightedTxInfo::from(entry.GetTx(), entry.GetFee()));
-    }
     mapTx.insert(entry);
     const CTransaction& tx = mapTx.find(hash)->GetTx();
     mapRecentlyAddedTx[tx.GetHash()] = &tx;
@@ -293,12 +292,10 @@ void CTxMemPool::remove(const CTransaction &origTx, std::list<CTransaction>& rem
                 removeAddressIndex(hash);
             if (fSpentIndex)
                 removeSpentIndex(hash);
-        } 
+        }
-        if (weightedTxList) {
+        assert(weightedTxTree);
-            weightedTxList->clear();
+        for (CTransaction tx : removed) {
-            for (const CTxMemPoolEntry& e : mapTx) {
+            weightedTxTree->remove(tx.GetHash());
-                weightedTxList->add(WeightedTxInfo::from(e.GetTx(), e.GetFee()));
-            }
         }
     }
 }
@@ -816,9 +813,9 @@ void CTxMemPool::setMempoolCostLimit(int64_t totalCostLimit, int64_t evictionMem
     LogPrint("mempool", "Setting mempool cost limit: (limit=%d, time=%d)\n", totalCostLimit, evictionMemorySeconds);
     // This method should not be called more than once
     assert(!recentlyEvicted);
-    assert(!weightedTxList);
+    assert(!weightedTxTree);
     recentlyEvicted = new RecentlyEvictedList(evictionMemorySeconds);
-    weightedTxList = new WeightedTransactionList(totalCostLimit);
+    weightedTxTree = new WeightedTxTree(totalCostLimit);
 }
 
 bool CTxMemPool::isRecentlyEvicted(const uint256& txId) {
@@ -830,12 +827,12 @@ bool CTxMemPool::isRecentlyEvicted(const uint256& txId) {
 void CTxMemPool::ensureSizeLimit() {
     AssertLockHeld(cs);
     assert(recentlyEvicted);
-    assert(weightedTxList);
+    assert(weightedTxTree);
-    boost::optional<WeightedTxInfo> maybeDrop;
+    boost::optional<uint256> maybeDropTxId;
-    std::list<CTransaction> removed;
+    while ((maybeDropTxId = weightedTxTree->maybeDropRandom()).is_initialized()) {
-    while ((maybeDrop = weightedTxList->maybeDropRandom(false)).is_initialized()) {
+        uint256 txId = maybeDropTxId.get();
-        uint256 txId = maybeDrop.get().txId;
         recentlyEvicted->add(txId);
+        std::list<CTransaction> removed;
         remove(mapTx.find(txId)->GetTx(), removed, true);
     }
 }

src/txmempool.h

@@ -141,7 +141,7 @@ private:
     std::map<uint256, const CTransaction*> mapSproutNullifiers;
     std::map<uint256, const CTransaction*> mapSaplingNullifiers;
     RecentlyEvictedList* recentlyEvicted = nullptr;
-    WeightedTransactionList* weightedTxList = nullptr;
+    WeightedTxTree* weightedTxTree = nullptr;
 
     void checkNullifiers(ShieldedType type) const;