From 1333ad7f783d496975d47eaf7b42c616f8476793 Mon Sep 17 00:00:00 2001
From: Daniel Krawisz <daniel@stashcrypto.com>
Date: Mon, 13 Nov 2017 22:37:35 -0600
Subject: [PATCH] FeeEstimator added to server. Mempool alerts the fee
 estimator of new txs that it observes. The block manager alerts the fee
 estimator of new and orphaned blocks.

Check for invalid state and recreate FeeEstimator if necessary.
---
 mempool/estimatefee.go |  9 +++++++++
 mempool/mempool.go     | 11 ++++++++++-
 netsync/interface.go   |  2 ++
 netsync/manager.go     | 23 +++++++++++++++++++++++
 server.go              |  5 +++++
 5 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/mempool/estimatefee.go b/mempool/estimatefee.go
index 354eabca..7f56403d 100644
--- a/mempool/estimatefee.go
+++ b/mempool/estimatefee.go
@@ -34,6 +34,15 @@ const (
 	// can be made by the txs found in a given block.
 	estimateFeeMaxReplacements = 10
 
+	// DefaultEstimateFeeMaxRollback is the default number of rollbacks
+	// allowed by the fee estimator for orphaned blocks.
+	DefaultEstimateFeeMaxRollback = 2
+
+	// DefaultEstimateFeeMinRegisteredBlocks is the default minimum
+	// number of blocks which must be observed by the fee estimator before
+	// it will provide fee estimations.
+	DefaultEstimateFeeMinRegisteredBlocks = 3
+
 	bytePerKb = 1024
 
 	btcPerSatoshi = 1E-8
diff --git a/mempool/mempool.go b/mempool/mempool.go
index 31d294a6..511ee6b4 100644
--- a/mempool/mempool.go
+++ b/mempool/mempool.go
@@ -89,6 +89,10 @@ type Config struct {
 	// indexing the unconfirmed transactions in the memory pool.
 	// This can be nil if the address index is not enabled.
 	AddrIndex *indexers.AddrIndex
+
+	// FeeEstimatator provides a feeEstimator. If it is not nil, the mempool
+	// records all new transactions it observes into the feeEstimator.
+	FeeEstimator *FeeEstimator
 }
 
 // Policy houses the policy (configuration parameters) which is used to
@@ -527,8 +531,8 @@ func (mp *TxPool) addTransaction(utxoView *blockchain.UtxoViewpoint, tx *btcutil
 		},
 		StartingPriority: mining.CalcPriority(tx.MsgTx(), utxoView, height),
 	}
-	mp.pool[*tx.Hash()] = txD
 
+	mp.pool[*tx.Hash()] = txD
 	for _, txIn := range tx.MsgTx().TxIn {
 		mp.outpoints[txIn.PreviousOutPoint] = tx
 	}
@@ -540,6 +544,11 @@ func (mp *TxPool) addTransaction(utxoView *blockchain.UtxoViewpoint, tx *btcutil
 		mp.cfg.AddrIndex.AddUnconfirmedTx(tx, utxoView)
 	}
 
+	// Record this tx for fee estimation if enabled.
+	if mp.cfg.FeeEstimator != nil {
+		mp.cfg.FeeEstimator.ObserveTransaction(txD)
+	}
+
 	return txD
 }
 
diff --git a/netsync/interface.go b/netsync/interface.go
index a177cb96..2e6edce8 100644
--- a/netsync/interface.go
+++ b/netsync/interface.go
@@ -36,4 +36,6 @@ type Config struct {
 
 	DisableCheckpoints bool
 	MaxPeers           int
+
+	FeeEstimator *mempool.FeeEstimator
 }
diff --git a/netsync/manager.go b/netsync/manager.go
index 7f1d3605..fe764a70 100644
--- a/netsync/manager.go
+++ b/netsync/manager.go
@@ -167,6 +167,9 @@ type SyncManager struct {
 	headerList       *list.List
 	startHeader      *list.Element
 	nextCheckpoint   *chaincfg.Checkpoint
+
+	// An optional fee estimator.
+	feeEstimator *mempool.FeeEstimator
 }
 
 // resetHeaderState sets the headers-first mode state to values appropriate for
@@ -1249,6 +1252,20 @@ func (sm *SyncManager) handleBlockchainNotification(notification *blockchain.Not
 			sm.peerNotifier.AnnounceNewTransactions(acceptedTxs)
 		}
 
+		// Register block with the fee estimator, if it exists.
+		if sm.feeEstimator != nil {
+			err := sm.feeEstimator.RegisterBlock(block)
+
+			// If an error is somehow generated then the fee estimator
+			// has entered an invalid state. Since it doesn't know how
+			// to recover, create a new one.
+			if err != nil {
+				sm.feeEstimator = mempool.NewFeeEstimator(
+					mempool.DefaultEstimateFeeMaxRollback,
+					mempool.DefaultEstimateFeeMinRegisteredBlocks)
+			}
+		}
+
 	// A block has been disconnected from the main block chain.
 	case blockchain.NTBlockDisconnected:
 		block, ok := notification.Data.(*btcutil.Block)
@@ -1269,6 +1286,11 @@ func (sm *SyncManager) handleBlockchainNotification(notification *blockchain.Not
 				sm.txMemPool.RemoveTransaction(tx, true)
 			}
 		}
+
+		// Rollback previous block recorded by the fee estimator.
+		if sm.feeEstimator != nil {
+			sm.feeEstimator.Rollback(block.Hash())
+		}
 	}
 }
 
@@ -1417,6 +1439,7 @@ func New(config *Config) (*SyncManager, error) {
 		msgChan:         make(chan interface{}, config.MaxPeers*3),
 		headerList:      list.New(),
 		quit:            make(chan struct{}),
+		feeEstimator:    config.FeeEstimator,
 	}
 
 	best := sm.chain.BestSnapshot()
diff --git a/server.go b/server.go
index b1249148..c9132471 100644
--- a/server.go
+++ b/server.go
@@ -2411,6 +2411,10 @@ func newServer(listenAddrs []string, db database.DB, chainParams *chaincfg.Param
 		return nil, err
 	}
 
+	feeEstimator := mempool.NewFeeEstimator(
+		mempool.DefaultEstimateFeeMaxRollback,
+		mempool.DefaultEstimateFeeMinRegisteredBlocks)
+
 	txC := mempool.Config{
 		Policy: mempool.Policy{
 			DisableRelayPriority: cfg.NoRelayPriority,
@@ -2433,6 +2437,7 @@ func newServer(listenAddrs []string, db database.DB, chainParams *chaincfg.Param
 		SigCache:           s.sigCache,
 		HashCache:          s.hashCache,
 		AddrIndex:          s.addrIndex,
+		FeeEstimator:       feeEstimator,
 	}
 	s.txMemPool = mempool.New(&txC)