2 files changed, 44 insertions, 18 deletions

diff --git a/src/policy/fees.cpp b/src/policy/fees.cpp
index 6d550b1d5a..74c688060d 100644
--- a/src/policy/fees.cpp
+++ b/src/policy/fees.cpp
@@ -515,15 +515,10 @@ void TxConfirmStats::removeTx(unsigned int entryHeight, unsigned int nBestSeenHe
     }
 }
 
-// This function is called from CTxMemPool::removeUnchecked to ensure
-// txs removed from the mempool for any reason are no longer
-// tracked. Txs that were part of a block have already been removed in
-// processBlockTx to ensure they are never double tracked, but it is
-// of no harm to try to remove them again.
-bool CBlockPolicyEstimator::removeTx(uint256 hash, bool inBlock)
+bool CBlockPolicyEstimator::removeTx(uint256 hash)
 {
     LOCK(m_cs_fee_estimator);
-    return _removeTx(hash, inBlock);
+    return _removeTx(hash, /*inBlock=*/false);
 }
 
 bool CBlockPolicyEstimator::_removeTx(const uint256& hash, bool inBlock)
@@ -579,11 +574,26 @@ CBlockPolicyEstimator::CBlockPolicyEstimator(const fs::path& estimation_filepath
 
 CBlockPolicyEstimator::~CBlockPolicyEstimator() = default;
 
-void CBlockPolicyEstimator::processTransaction(const CTxMemPoolEntry& entry, bool validFeeEstimate)
+void CBlockPolicyEstimator::TransactionAddedToMempool(const NewMempoolTransactionInfo& tx, uint64_t /*unused*/)
+{
+    processTransaction(tx);
+}
+
+void CBlockPolicyEstimator::TransactionRemovedFromMempool(const CTransactionRef& tx, MemPoolRemovalReason /*unused*/, uint64_t /*unused*/)
+{
+    removeTx(tx->GetHash());
+}
+
+void CBlockPolicyEstimator::MempoolTransactionsRemovedForBlock(const std::vector<RemovedMempoolTransactionInfo>& txs_removed_for_block, unsigned int nBlockHeight)
+{
+    processBlock(txs_removed_for_block, nBlockHeight);
+}
+
+void CBlockPolicyEstimator::processTransaction(const NewMempoolTransactionInfo& tx)
 {
     LOCK(m_cs_fee_estimator);
-    unsigned int txHeight = entry.GetHeight();
-    uint256 hash = entry.GetTx().GetHash();
+    const unsigned int txHeight = tx.info.txHeight;
+    const auto& hash = tx.info.m_tx->GetHash();
     if (mapMemPoolTxs.count(hash)) {
         LogPrint(BCLog::ESTIMATEFEE, "Blockpolicy error mempool tx %s already being tracked\n",
                  hash.ToString());
@@ -597,17 +607,23 @@ void CBlockPolicyEstimator::processTransaction(const CTxMemPoolEntry& entry, boo
         // It will be synced next time a block is processed.
         return;
     }
+    // This transaction should only count for fee estimation if:
+    // - it's not being re-added during a reorg which bypasses typical mempool fee limits
+    // - the node is not behind
+    // - the transaction is not dependent on any other transactions in the mempool
+    // - it's not part of a package.
+    const bool validForFeeEstimation = !tx.m_from_disconnected_block && !tx.m_submitted_in_package && tx.m_chainstate_is_current && tx.m_has_no_mempool_parents;
 
     // Only want to be updating estimates when our blockchain is synced,
     // otherwise we'll miscalculate how many blocks its taking to get included.
-    if (!validFeeEstimate) {
+    if (!validForFeeEstimation) {
         untrackedTxs++;
         return;
     }
     trackedTxs++;
 
     // Feerates are stored and reported as BTC-per-kb:
-    CFeeRate feeRate(entry.GetFee(), entry.GetTxSize());
+    const CFeeRate feeRate(tx.info.m_fee, tx.info.m_virtual_transaction_size);
 
     mapMemPoolTxs[hash].blockHeight = txHeight;
     unsigned int bucketIndex = feeStats->NewTx(txHeight, static_cast<double>(feeRate.GetFeePerK()));
diff --git a/src/policy/fees.h b/src/policy/fees.h
index cff0c4dbe2..f34f66d3f0 100644
--- a/src/policy/fees.h
+++ b/src/policy/fees.h
@@ -12,6 +12,7 @@
 #include <threadsafety.h>
 #include <uint256.h>
 #include <util/fs.h>
+#include <validationinterface.h>
 
 #include <array>
 #include <chrono>
@@ -35,9 +36,9 @@ static constexpr std::chrono::hours MAX_FILE_AGE{60};
 static constexpr bool DEFAULT_ACCEPT_STALE_FEE_ESTIMATES{false};
 
 class AutoFile;
-class CTxMemPoolEntry;
 class TxConfirmStats;
 struct RemovedMempoolTransactionInfo;
+struct NewMempoolTransactionInfo;
 
 /* Identifier for each of the 3 different TxConfirmStats which will track
  * history over different time horizons. */
@@ -144,7 +145,7 @@ struct FeeCalculation
  * a certain number of blocks.  Every time a block is added to the best chain, this class records
  * stats on the transactions included in that block
  */
-class CBlockPolicyEstimator
+class CBlockPolicyEstimator : public CValidationInterface
 {
 private:
     /** Track confirm delays up to 12 blocks for short horizon */
@@ -199,7 +200,7 @@ private:
 public:
     /** Create new BlockPolicyEstimator and initialize stats tracking classes with default values */
     CBlockPolicyEstimator(const fs::path& estimation_filepath, const bool read_stale_estimates);
-    ~CBlockPolicyEstimator();
+    virtual ~CBlockPolicyEstimator();
 
     /** Process all the transactions that have been included in a block */
     void processBlock(const std::vector<RemovedMempoolTransactionInfo>& txs_removed_for_block,
@@ -207,11 +208,11 @@ public:
         EXCLUSIVE_LOCKS_REQUIRED(!m_cs_fee_estimator);
 
     /** Process a transaction accepted to the mempool*/
-    void processTransaction(const CTxMemPoolEntry& entry, bool validFeeEstimate)
+    void processTransaction(const NewMempoolTransactionInfo& tx)
         EXCLUSIVE_LOCKS_REQUIRED(!m_cs_fee_estimator);
 
-    /** Remove a transaction from the mempool tracking stats*/
-    bool removeTx(uint256 hash, bool inBlock)
+    /** Remove a transaction from the mempool tracking stats for non BLOCK removal reasons*/
+    bool removeTx(uint256 hash)
         EXCLUSIVE_LOCKS_REQUIRED(!m_cs_fee_estimator);
 
     /** DEPRECATED. Return a feerate estimate */
@@ -261,6 +262,15 @@ public:
     /** Calculates the age of the file, since last modified */
     std::chrono::hours GetFeeEstimatorFileAge();
 
+protected:
+    /** Overridden from CValidationInterface. */
+    void TransactionAddedToMempool(const NewMempoolTransactionInfo& tx, uint64_t /*unused*/) override
+        EXCLUSIVE_LOCKS_REQUIRED(!m_cs_fee_estimator);
+    void TransactionRemovedFromMempool(const CTransactionRef& tx, MemPoolRemovalReason /*unused*/, uint64_t /*unused*/) override
+        EXCLUSIVE_LOCKS_REQUIRED(!m_cs_fee_estimator);
+    void MempoolTransactionsRemovedForBlock(const std::vector<RemovedMempoolTransactionInfo>& txs_removed_for_block, unsigned int nBlockHeight) override
+        EXCLUSIVE_LOCKS_REQUIRED(!m_cs_fee_estimator);
+
 private:
     mutable Mutex m_cs_fee_estimator;