1 files changed, 16 insertions, 6 deletions
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;