feeEstimator changed to FeeEstimator. A number of optimizations and improvements.

Rollback takes a block hash rather than a BlockStamp.

Increase rounds in TestEstimateFeeRollback to test dropping txs that have been in the mempool too long.

mempool/estimatefee.go

@@ -33,20 +33,27 @@ const (
 	// estimateFeeMaxReplacements is the max number of replacements that
 	// can be made by the txs found in a given block.
 	estimateFeeMaxReplacements = 10
+
+	bytePerKb = 1024
+
+	btcPerSatoshi = 1E-8
 )
 
 // SatoshiPerByte is number with units of satoshis per byte.
 type SatoshiPerByte float64
 
-// ToSatoshiPerKb returns a float value that represents the given
+// BtcPerKilobyte is number with units of bitcoins per kilobyte.
+type BtcPerKilobyte float64
+
+// ToBtcPerKb returns a float value that represents the given
 // SatoshiPerByte converted to satoshis per kb.
-func (rate SatoshiPerByte) ToSatoshiPerKb() float64 {
+func (rate SatoshiPerByte) ToBtcPerKb() BtcPerKilobyte {
 	// If our rate is the error value, return that.
 	if rate == SatoshiPerByte(-1.0) {
 		return -1.0
 	}
 
-	return float64(rate) * 1024
+	return BtcPerKilobyte(float64(rate) * bytePerKb * btcPerSatoshi)
 }
 
 // Fee returns the fee for a transaction of a given size for
@@ -88,7 +95,7 @@ type observedTransaction struct {
 // is used if Rollback is called to reverse the effect of registering
 // a block.
 type registeredBlock struct {
-	hash         *chainhash.Hash
+	hash         chainhash.Hash
 	transactions []*observedTransaction
 }
 
@@ -96,11 +103,11 @@ type registeredBlock struct {
 // fee estimations. It is safe for concurrent access.
 type FeeEstimator struct {
 	maxRollback uint32
-	binSize     int
+	binSize     int32
 
 	// The maximum number of replacements that can be made in a single
 	// bin per block. Default is estimateFeeMaxReplacements
-	maxReplacements int
+	maxReplacements int32
 
 	// The minimum number of blocks that can be registered with the fee
 	// estimator before it will provide answers.
@@ -109,17 +116,19 @@ type FeeEstimator struct {
 	// The last known height.
 	lastKnownHeight int32
 
-	sync.RWMutex
+	// The number of blocks that have been registered.
-	observed            map[chainhash.Hash]observedTransaction
+	numBlocksRegistered uint32
-	bin                 [estimateFeeDepth][]*observedTransaction
+
-	numBlocksRegistered uint32 // The number of blocks that have been registered.
+	mtx      sync.RWMutex
+	observed map[chainhash.Hash]*observedTransaction
+	bin      [estimateFeeDepth][]*observedTransaction
 
 	// The cached estimates.
 	cached []SatoshiPerByte
 
 	// Transactions that have been removed from the bins. This allows us to
 	// revert in case of an orphaned block.
-	dropped []registeredBlock
+	dropped []*registeredBlock
 }
 
 // NewFeeEstimator creates a FeeEstimator for which at most maxRollback blocks
@@ -132,21 +141,27 @@ func NewFeeEstimator(maxRollback, minRegisteredBlocks uint32) *FeeEstimator {
 		lastKnownHeight:     mining.UnminedHeight,
 		binSize:             estimateFeeBinSize,
 		maxReplacements:     estimateFeeMaxReplacements,
-		observed:            make(map[chainhash.Hash]observedTransaction),
+		observed:            make(map[chainhash.Hash]*observedTransaction),
-		dropped:             make([]registeredBlock, 0, maxRollback),
+		dropped:             make([]*registeredBlock, 0, maxRollback),
 	}
 }
 
 // ObserveTransaction is called when a new transaction is observed in the mempool.
 func (ef *FeeEstimator) ObserveTransaction(t *TxDesc) {
-	ef.Lock()
+	ef.mtx.Lock()
-	defer ef.Unlock()
+	defer ef.mtx.Unlock()
+
+	// If we haven't seen a block yet we don't know when this one arrived,
+	// so we ignore it.
+	if ef.lastKnownHeight == mining.UnminedHeight {
+		return
+	}
 
 	hash := *t.Tx.Hash()
 	if _, ok := ef.observed[hash]; !ok {
 		size := uint32(t.Tx.MsgTx().SerializeSize())
 
-		ef.observed[hash] = observedTransaction{
+		ef.observed[hash] = &observedTransaction{
 			hash:     hash,
 			feeRate:  NewSatoshiPerByte(btcutil.Amount(t.Fee), size),
 			observed: t.Height,
@@ -157,8 +172,8 @@ func (ef *FeeEstimator) ObserveTransaction(t *TxDesc) {
 
 // RegisterBlock informs the fee estimator of a new block to take into account.
 func (ef *FeeEstimator) RegisterBlock(block *btcutil.Block) error {
-	ef.Lock()
+	ef.mtx.Lock()
-	defer ef.Unlock()
+	defer ef.mtx.Unlock()
 
 	// The previous sorted list is invalid, so delete it.
 	ef.cached = nil
@@ -184,8 +199,8 @@ func (ef *FeeEstimator) RegisterBlock(block *btcutil.Block) error {
 	var replacementCounts [estimateFeeDepth]int
 
 	// Keep track of which txs were dropped in case of an orphan block.
-	dropped := registeredBlock{
+	dropped := &registeredBlock{
-		hash:         block.Hash(),
+		hash:         *block.Hash(),
 		transactions: make([]*observedTransaction, 0, 100),
 	}
 
@@ -200,41 +215,50 @@ func (ef *FeeEstimator) RegisterBlock(block *btcutil.Block) error {
 		}
 
 		// Put the observed tx in the oppropriate bin.
-		o.mined = height
-
 		blocksToConfirm := height - o.observed - 1
 
+		// This shouldn't happen if the fee estimator works correctly,
+		// but return an error if it does.
+		if o.mined != mining.UnminedHeight {
+			log.Error("Estimate fee: transaction ", hash.String(), " has already been mined")
+			return errors.New("Transaction has already been mined")
+		}
+
 		// This shouldn't happen but check just in case to avoid
-		// a panic later.
+		// an out-of-bounds array index later.
 		if blocksToConfirm >= estimateFeeDepth {
 			continue
 		}
 
 		// Make sure we do not replace too many transactions per min.
-		if replacementCounts[blocksToConfirm] == ef.maxReplacements {
+		if replacementCounts[blocksToConfirm] == int(ef.maxReplacements) {
 			continue
 		}
 
+		o.mined = height
+
 		replacementCounts[blocksToConfirm]++
 
 		bin := ef.bin[blocksToConfirm]
 
 		// Remove a random element and replace it with this new tx.
-		if len(bin) == ef.binSize {
+		if len(bin) == int(ef.binSize) {
-			l := ef.binSize - replacementCounts[blocksToConfirm]
+			// Don't drop transactions we have just added from this same block.
+			l := int(ef.binSize) - replacementCounts[blocksToConfirm]
 			drop := rand.Intn(l)
 			dropped.transactions = append(dropped.transactions, bin[drop])
 
 			bin[drop] = bin[l-1]
-			bin[l-1] = &o
+			bin[l-1] = o
 		} else {
-			ef.bin[blocksToConfirm] = append(bin, &o)
+			bin = append(bin, o)
 		}
+		ef.bin[blocksToConfirm] = bin
 	}
 
 	// Go through the mempool for txs that have been in too long.
 	for hash, o := range ef.observed {
-		if height-o.observed >= estimateFeeDepth {
+		if o.mined == mining.UnminedHeight && height-o.observed >= estimateFeeDepth {
 			delete(ef.observed, hash)
 		}
 	}
@@ -253,6 +277,14 @@ func (ef *FeeEstimator) RegisterBlock(block *btcutil.Block) error {
 	return nil
 }
 
+// LastKnownHeight returns the height of the last block which was registered.
+func (ef *FeeEstimator) LastKnownHeight() int32 {
+	ef.mtx.Lock()
+	defer ef.mtx.Unlock()
+
+	return ef.lastKnownHeight
+}
+
 // Rollback unregisters a recently registered block from the FeeEstimator.
 // This can be used to reverse the effect of an orphaned block on the fee
 // estimator. The maximum number of rollbacks allowed is given by
@@ -262,29 +294,24 @@ func (ef *FeeEstimator) RegisterBlock(block *btcutil.Block) error {
 // deleted if they have been observed too long ago. That means the result
 // of Rollback won't always be exactly the same as if the last block had not
 // happened, but it should be close enough.
-func (ef *FeeEstimator) Rollback(block *btcutil.Block) error {
+func (ef *FeeEstimator) Rollback(hash *chainhash.Hash) error {
-	ef.Lock()
+	ef.mtx.Lock()
-	defer ef.Unlock()
+	defer ef.mtx.Unlock()
-
-	hash := block.Hash()
 
 	// Find this block in the stack of recent registered blocks.
 	var n int
-	for n = 1; n < len(ef.dropped); n++ {
+	for n = 1; n <= len(ef.dropped); n++ {
 		if ef.dropped[len(ef.dropped)-n].hash.IsEqual(hash) {
 			break
 		}
 	}
 
-	if n == len(ef.dropped) {
+	if n > len(ef.dropped) {
 		return errors.New("no such block was recently registered")
 	}
 
 	for i := 0; i < n; i++ {
-		err := ef.rollback()
+		ef.rollback()
-		if err != nil {
-			return err
-		}
 	}
 
 	return nil
@@ -292,28 +319,24 @@ func (ef *FeeEstimator) Rollback(block *btcutil.Block) error {
 
 // rollback rolls back the effect of the last block in the stack
 // of registered blocks.
-func (ef *FeeEstimator) rollback() error {
+func (ef *FeeEstimator) rollback() {
-
 	// The previous sorted list is invalid, so delete it.
 	ef.cached = nil
 
 	// pop the last list of dropped txs from the stack.
 	last := len(ef.dropped) - 1
 	if last == -1 {
-		// Return if we cannot rollback.
+		// Cannot really happen because the exported calling function
-		return errors.New("max rollbacks reached")
+		// only rolls back a block already known to be in the list
+		// of dropped transactions.
+		return
 	}
 
-	ef.numBlocksRegistered--
-
 	dropped := ef.dropped[last]
-	ef.dropped = ef.dropped[0:last]
 
 	// where we are in each bin as we replace txs?
 	var replacementCounters [estimateFeeDepth]int
 
-	var err error
-
 	// Go through the txs in the dropped block.
 	for _, o := range dropped.transactions {
 		// Which bin was this tx in?
@@ -326,9 +349,8 @@ func (ef *FeeEstimator) rollback() error {
 		// Continue to go through that bin where we left off.
 		for {
 			if counter >= len(bin) {
-				// Create an error but keep going in case we can roll back
+				// Panic, as we have entered an unrecoverable invalid state.
-				// more transactions successfully.
+				panic(errors.New("illegal state: cannot rollback dropped transaction"))
-				err = errors.New("illegal state: cannot rollback dropped transaction")
 			}
 
 			prev := bin[counter]
@@ -344,6 +366,8 @@ func (ef *FeeEstimator) rollback() error {
 
 			counter++
 		}
+
+		replacementCounters[blocksToConfirm] = counter
 	}
 
 	// Continue going through bins to find other txs to remove
@@ -373,9 +397,11 @@ func (ef *FeeEstimator) rollback() error {
 		}
 	}
 
-	ef.lastKnownHeight--
+	ef.dropped = ef.dropped[0:last]
 
-	return err
+	// The number of blocks the fee estimator has seen is decrimented.
+	ef.numBlocksRegistered--
+	ef.lastKnownHeight--
 }
 
 // estimateFeeSet is a set of txs that can that is sorted
@@ -407,19 +433,26 @@ func (b *estimateFeeSet) estimateFee(confirmations int) SatoshiPerByte {
 		return 0
 	}
 
-	var min, max uint32 = 0, 0
-	for i := 0; i < confirmations-1; i++ {
-		min += b.bin[i]
-	}
-
-	max = min + b.bin[confirmations-1]
-
 	// We don't have any transactions!
-	if min == 0 && max == 0 {
+	if len(b.feeRate) == 0 {
 		return 0
 	}
 
-	return b.feeRate[(min+max-1)/2] * 1E-8
+	var min, max int = 0, 0
+	for i := 0; i < confirmations-1; i++ {
+		min += int(b.bin[i])
+	}
+
+	max = min + int(b.bin[confirmations-1]) - 1
+	if max < min {
+		max = min
+	}
+	feeIndex := (min + max) / 2
+	if feeIndex >= len(b.feeRate) {
+		feeIndex = len(b.feeRate) - 1
+	}
+
+	return b.feeRate[feeIndex]
 }
 
 // newEstimateFeeSet creates a temporary data structure that
@@ -464,9 +497,9 @@ func (ef *FeeEstimator) estimates() []SatoshiPerByte {
 
 // EstimateFee estimates the fee per byte to have a tx confirmed a given
 // number of blocks from now.
-func (ef *FeeEstimator) EstimateFee(numBlocks uint32) (SatoshiPerByte, error) {
+func (ef *FeeEstimator) EstimateFee(numBlocks uint32) (BtcPerKilobyte, error) {
-	ef.Lock()
+	ef.mtx.Lock()
-	defer ef.Unlock()
+	defer ef.mtx.Unlock()
 
 	// If the number of registered blocks is below the minimum, return
 	// an error.
@@ -489,5 +522,5 @@ func (ef *FeeEstimator) EstimateFee(numBlocks uint32) (SatoshiPerByte, error) {
 		ef.cached = ef.estimates()
 	}
 
-	return ef.cached[int(numBlocks)-1], nil
+	return ef.cached[int(numBlocks)-1].ToBtcPerKb(), nil
 }

mempool/estimatefee_test.go

@@ -19,19 +19,30 @@ import (
 func newTestFeeEstimator(binSize, maxReplacements, maxRollback uint32) *FeeEstimator {
 	return &FeeEstimator{
 		maxRollback:         maxRollback,
-		lastKnownHeight:     mining.UnminedHeight,
+		lastKnownHeight:     0,
-		binSize:             int(binSize),
+		binSize:             int32(binSize),
 		minRegisteredBlocks: 0,
-		maxReplacements:     int(maxReplacements),
+		maxReplacements:     int32(maxReplacements),
-		observed:            make(map[chainhash.Hash]observedTransaction),
+		observed:            make(map[chainhash.Hash]*observedTransaction),
-		dropped:             make([]registeredBlock, 0, maxRollback),
+		dropped:             make([]*registeredBlock, 0, maxRollback),
 	}
 }
 
+// lastBlock is a linked list of the block hashes which have been
+// processed by the test FeeEstimator.
+type lastBlock struct {
+	hash *chainhash.Hash
+	prev *lastBlock
+}
+
+// estimateFeeTester interacts with the FeeEstimator to keep track
+// of its expected state.
 type estimateFeeTester struct {
+	ef      *FeeEstimator
 	t       *testing.T
 	version int32
 	height  int32
+	last    *lastBlock
 }
 
 func (eft *estimateFeeTester) testTx(fee btcutil.Amount) *TxDesc {
@@ -48,30 +59,48 @@ func (eft *estimateFeeTester) testTx(fee btcutil.Amount) *TxDesc {
 	}
 }
 
-func expectedFeePerByte(t *TxDesc) SatoshiPerByte {
+func expectedFeePerKilobyte(t *TxDesc) BtcPerKilobyte {
 	size := float64(t.TxDesc.Tx.MsgTx().SerializeSize())
 	fee := float64(t.TxDesc.Fee)
 
-	return SatoshiPerByte(fee / size * 1E-8)
+	return SatoshiPerByte(fee / size).ToBtcPerKb()
 }
 
-func (eft *estimateFeeTester) testBlock(txs []*wire.MsgTx) *btcutil.Block {
+func (eft *estimateFeeTester) newBlock(txs []*wire.MsgTx) {
-
 	eft.height++
+
 	block := btcutil.NewBlock(&wire.MsgBlock{
 		Transactions: txs,
 	})
 	block.SetHeight(eft.height)
 
-	return block
+	eft.last = &lastBlock{block.Hash(), eft.last}
+
+	eft.ef.RegisterBlock(block)
 }
 
+func (eft *estimateFeeTester) rollback() {
+	if eft.last == nil {
+		return
+	}
+
+	err := eft.ef.Rollback(eft.last.hash)
+
+	if err != nil {
+		eft.t.Errorf("Could not rollback: %v", err)
+	}
+
+	eft.height--
+	eft.last = eft.last.prev
+}
+
+// TestEstimateFee tests basic functionality in the FeeEstimator.
 func TestEstimateFee(t *testing.T) {
-	ef := newTestFeeEstimator(5, 3, 0)
+	ef := newTestFeeEstimator(5, 3, 1)
-	eft := estimateFeeTester{t: t}
+	eft := estimateFeeTester{ef: ef, t: t}
 
 	// Try with no txs and get zero for all queries.
-	expected := SatoshiPerByte(0.0)
+	expected := BtcPerKilobyte(0.0)
 	for i := uint32(1); i <= estimateFeeDepth; i++ {
 		estimated, _ := ef.EstimateFee(i)
 
@@ -85,7 +114,7 @@ func TestEstimateFee(t *testing.T) {
 	ef.ObserveTransaction(tx)
 
 	// Expected should still be zero because this is still in the mempool.
-	expected = SatoshiPerByte(0.0)
+	expected = BtcPerKilobyte(0.0)
 	for i := uint32(1); i <= estimateFeeDepth; i++ {
 		estimated, _ := ef.EstimateFee(i)
 
@@ -97,7 +126,7 @@ func TestEstimateFee(t *testing.T) {
 	// Change minRegisteredBlocks to make sure that works. Error return
 	// value expected.
 	ef.minRegisteredBlocks = 1
-	expected = SatoshiPerByte(-1.0)
+	expected = BtcPerKilobyte(-1.0)
 	for i := uint32(1); i <= estimateFeeDepth; i++ {
 		estimated, _ := ef.EstimateFee(i)
 
@@ -106,9 +135,35 @@ func TestEstimateFee(t *testing.T) {
 		}
 	}
 
-	// Record a block.
+	// Record a block with the new tx.
-	ef.RegisterBlock(eft.testBlock([]*wire.MsgTx{tx.Tx.MsgTx()}))
+	eft.newBlock([]*wire.MsgTx{tx.Tx.MsgTx()})
-	expected = expectedFeePerByte(tx)
+	expected = expectedFeePerKilobyte(tx)
+	for i := uint32(1); i <= estimateFeeDepth; i++ {
+		estimated, _ := ef.EstimateFee(i)
+
+		if estimated != expected {
+			t.Errorf("Estimate fee error: expected %f when one tx is binned; got %f", expected, estimated)
+		}
+	}
+
+	// Roll back the last block; this was an orphan block.
+	ef.minRegisteredBlocks = 0
+	eft.rollback()
+	expected = BtcPerKilobyte(0.0)
+	for i := uint32(1); i <= estimateFeeDepth; i++ {
+		estimated, _ := ef.EstimateFee(i)
+
+		if estimated != expected {
+			t.Errorf("Estimate fee error: expected %f after rolling back block; got %f", expected, estimated)
+		}
+	}
+
+	// Record an empty block and then a block with the new tx.
+	// This test was made because of a bug that only appeared when there
+	// were no transactions in the first bin.
+	eft.newBlock([]*wire.MsgTx{})
+	eft.newBlock([]*wire.MsgTx{tx.Tx.MsgTx()})
+	expected = expectedFeePerKilobyte(tx)
 	for i := uint32(1); i <= estimateFeeDepth; i++ {
 		estimated, _ := ef.EstimateFee(i)
 
@@ -125,22 +180,22 @@ func TestEstimateFee(t *testing.T) {
 	ef.ObserveTransaction(txB)
 	ef.ObserveTransaction(txC)
 
-	// Record 8 empty blocks.
+	// Record 7 empty blocks.
-	for i := 0; i < 8; i++ {
+	for i := 0; i < 7; i++ {
-		ef.RegisterBlock(eft.testBlock([]*wire.MsgTx{}))
+		eft.newBlock([]*wire.MsgTx{})
 	}
 
 	// Mine the first tx.
-	ef.RegisterBlock(eft.testBlock([]*wire.MsgTx{txA.Tx.MsgTx()}))
+	eft.newBlock([]*wire.MsgTx{txA.Tx.MsgTx()})
 
 	// Now the estimated amount should depend on the value
 	// of the argument to estimate fee.
 	for i := uint32(1); i <= estimateFeeDepth; i++ {
 		estimated, _ := ef.EstimateFee(i)
-		if i > 8 {
+		if i > 2 {
-			expected = expectedFeePerByte(txA)
+			expected = expectedFeePerKilobyte(txA)
 		} else {
-			expected = expectedFeePerByte(tx)
+			expected = expectedFeePerKilobyte(tx)
 		}
 		if estimated != expected {
 			t.Errorf("Estimate fee error: expected %f on round %d; got %f", expected, i, estimated)
@@ -149,22 +204,22 @@ func TestEstimateFee(t *testing.T) {
 
 	// Record 5 more empty blocks.
 	for i := 0; i < 5; i++ {
-		ef.RegisterBlock(eft.testBlock([]*wire.MsgTx{}))
+		eft.newBlock([]*wire.MsgTx{})
 	}
 
 	// Mine the next tx.
-	ef.RegisterBlock(eft.testBlock([]*wire.MsgTx{txB.Tx.MsgTx()}))
+	eft.newBlock([]*wire.MsgTx{txB.Tx.MsgTx()})
 
 	// Now the estimated amount should depend on the value
 	// of the argument to estimate fee.
 	for i := uint32(1); i <= estimateFeeDepth; i++ {
 		estimated, _ := ef.EstimateFee(i)
-		if i <= 8 {
+		if i <= 2 {
-			expected = expectedFeePerByte(txB)
+			expected = expectedFeePerKilobyte(txB)
-		} else if i <= 8+6 {
+		} else if i <= 8 {
-			expected = expectedFeePerByte(tx)
+			expected = expectedFeePerKilobyte(tx)
 		} else {
-			expected = expectedFeePerByte(txA)
+			expected = expectedFeePerKilobyte(txA)
 		}
 
 		if estimated != expected {
@@ -174,22 +229,24 @@ func TestEstimateFee(t *testing.T) {
 
 	// Record 9 more empty blocks.
 	for i := 0; i < 10; i++ {
-		ef.RegisterBlock(eft.testBlock([]*wire.MsgTx{}))
+		eft.newBlock([]*wire.MsgTx{})
 	}
 
 	// Mine txC.
-	ef.RegisterBlock(eft.testBlock([]*wire.MsgTx{txC.Tx.MsgTx()}))
+	eft.newBlock([]*wire.MsgTx{txC.Tx.MsgTx()})
 
 	// This should have no effect on the outcome because too
 	// many blocks have been mined for txC to be recorded.
 	for i := uint32(1); i <= estimateFeeDepth; i++ {
 		estimated, _ := ef.EstimateFee(i)
-		if i <= 8 {
+		if i <= 2 {
-			expected = expectedFeePerByte(txB)
+			expected = expectedFeePerKilobyte(txC)
+		} else if i <= 8 {
+			expected = expectedFeePerKilobyte(txB)
 		} else if i <= 8+6 {
-			expected = expectedFeePerByte(tx)
+			expected = expectedFeePerKilobyte(tx)
 		} else {
-			expected = expectedFeePerByte(txA)
+			expected = expectedFeePerKilobyte(txA)
 		}
 
 		if estimated != expected {
@@ -198,133 +255,112 @@ func TestEstimateFee(t *testing.T) {
 	}
 }
 
-func (eft *estimateFeeTester) estimates(ef *FeeEstimator) [estimateFeeDepth]SatoshiPerByte {
+func (eft *estimateFeeTester) estimates() [estimateFeeDepth]BtcPerKilobyte {
 
 	// Generate estimates
-	var estimates [estimateFeeDepth]SatoshiPerByte
+	var estimates [estimateFeeDepth]BtcPerKilobyte
 	for i := 0; i < estimateFeeDepth; i++ {
-		estimates[i], _ = ef.EstimateFee(1)
+		estimates[i], _ = eft.ef.EstimateFee(uint32(i + 1))
 	}
 
 	// Check that all estimated fee results go in descending order.
 	for i := 1; i < estimateFeeDepth; i++ {
 		if estimates[i] > estimates[i-1] {
-			eft.t.Error("Estimates not in descending order.")
+			eft.t.Error("Estimates not in descending order; got ",
+				estimates[i], " for estimate ", i, " and ", estimates[i-1], " for ", (i - 1))
+			panic("invalid state.")
 		}
 	}
 
 	return estimates
 }
 
-func (eft *estimateFeeTester) round(ef *FeeEstimator,
+func (eft *estimateFeeTester) round(txHistory [][]*TxDesc,
-	txHistory [][]*TxDesc, blockHistory []*btcutil.Block,
+	estimateHistory [][estimateFeeDepth]BtcPerKilobyte,
-	estimateHistory [][estimateFeeDepth]SatoshiPerByte,
+	txPerRound, txPerBlock uint32) ([][]*TxDesc, [][estimateFeeDepth]BtcPerKilobyte) {
-	txPerRound, txPerBlock, maxRollback uint32) ([][]*TxDesc,
-	[]*btcutil.Block, [][estimateFeeDepth]SatoshiPerByte) {
 
 	// generate new txs.
 	var newTxs []*TxDesc
 	for i := uint32(0); i < txPerRound; i++ {
 		newTx := eft.testTx(btcutil.Amount(rand.Intn(1000000)))
-		ef.ObserveTransaction(newTx)
+		eft.ef.ObserveTransaction(newTx)
 		newTxs = append(newTxs, newTx)
 	}
 
-	// Construct new tx history.
+	// Generate mempool.
-	txHistory = append(txHistory, newTxs)
+	mempool := make(map[*observedTransaction]*TxDesc)
-	if len(txHistory) > estimateFeeDepth {
+	for _, h := range txHistory {
-		txHistory = txHistory[1 : estimateFeeDepth+1]
+		for _, t := range h {
+			if o, exists := eft.ef.observed[*t.Tx.Hash()]; exists && o.mined == mining.UnminedHeight {
+				mempool[o] = t
+			}
+		}
 	}
 
 	// generate new block, with no duplicates.
-	newBlockTxs := make(map[chainhash.Hash]*wire.MsgTx)
 	i := uint32(0)
-	for i < txPerBlock {
+	newBlockList := make([]*wire.MsgTx, 0, txPerBlock)
-		n := rand.Intn(len(txHistory))
+	for _, t := range mempool {
-		m := rand.Intn(int(txPerRound))
+		newBlockList = append(newBlockList, t.TxDesc.Tx.MsgTx())
-
-		tx := txHistory[n][m]
-		hash := *tx.Tx.Hash()
-
-		if _, ok := newBlockTxs[hash]; ok {
-			continue
-		}
-
-		newBlockTxs[hash] = tx.Tx.MsgTx()
 		i++
+
+		if i == txPerBlock {
+			break
+		}
 	}
 
-	var newBlockList []*wire.MsgTx
+	// Register a new block.
-	for _, tx := range newBlockTxs {
+	eft.newBlock(newBlockList)
-		newBlockList = append(newBlockList, tx)
-	}
-
-	newBlock := eft.testBlock(newBlockList)
-	ef.RegisterBlock(newBlock)
 
 	// return results.
-	estimates := eft.estimates(ef)
+	estimates := eft.estimates()
 
 	// Return results
-	blockHistory = append(blockHistory, newBlock)
+	return append(txHistory, newTxs), append(estimateHistory, estimates)
-	if len(blockHistory) > int(maxRollback) {
-		blockHistory = blockHistory[1 : maxRollback+1]
-	}
-
-	return txHistory, blockHistory, append(estimateHistory, estimates)
 }
 
+// TestEstimateFeeRollback tests the rollback function, which undoes the
+// effect of a adding a new block.
 func TestEstimateFeeRollback(t *testing.T) {
-	txPerRound := uint32(20)
+	txPerRound := uint32(7)
-	txPerBlock := uint32(10)
+	txPerBlock := uint32(5)
-	binSize := uint32(5)
+	binSize := uint32(6)
-	maxReplacements := uint32(3)
+	maxReplacements := uint32(4)
 	stepsBack := 2
 	rounds := 30
 
-	ef := newTestFeeEstimator(binSize, maxReplacements, uint32(stepsBack))
+	eft := estimateFeeTester{ef: newTestFeeEstimator(binSize, maxReplacements, uint32(stepsBack)), t: t}
-	eft := estimateFeeTester{t: t}
 	var txHistory [][]*TxDesc
-	var blockHistory []*btcutil.Block
+	estimateHistory := [][estimateFeeDepth]BtcPerKilobyte{eft.estimates()}
-	estimateHistory := [][estimateFeeDepth]SatoshiPerByte{eft.estimates(ef)}
-
-	// Make some initial rounds so that we have room to step back.
-	for round := 0; round < stepsBack-1; round++ {
-		txHistory, blockHistory, estimateHistory =
-			eft.round(ef, txHistory, blockHistory, estimateHistory,
-				txPerRound, txPerBlock, uint32(stepsBack))
-	}
 
 	for round := 0; round < rounds; round++ {
-		txHistory, blockHistory, estimateHistory =
+		// Go forward a few rounds.
-			eft.round(ef, txHistory, blockHistory, estimateHistory,
+		for step := 0; step <= stepsBack; step++ {
-				txPerRound, txPerBlock, uint32(stepsBack))
+			txHistory, estimateHistory =
+				eft.round(txHistory, estimateHistory, txPerRound, txPerBlock)
+		}
 
+		// Now go back.
 		for step := 0; step < stepsBack; step++ {
-			err := ef.rollback()
+			eft.rollback()
-			if err != nil {
-				t.Fatal("Could not rollback: ", err)
-			}
 
+			// After rolling back, we should have the same estimated
+			// fees as before.
 			expected := estimateHistory[len(estimateHistory)-step-2]
-			estimates := eft.estimates(ef)
+			estimates := eft.estimates()
 
 			// Ensure that these are both the same.
 			for i := 0; i < estimateFeeDepth; i++ {
 				if expected[i] != estimates[i] {
 					t.Errorf("Rollback value mismatch. Expected %f, got %f. ",
 						expected[i], estimates[i])
+					return
 				}
 			}
 		}
 
-		// Remove last estries from estimateHistory
+		// Erase history.
+		txHistory = txHistory[0 : len(txHistory)-stepsBack]
 		estimateHistory = estimateHistory[0 : len(estimateHistory)-stepsBack]
-
-		// replay the previous blocks.
-		for b := 0; b < stepsBack; b++ {
-			ef.RegisterBlock(blockHistory[b])
-			estimateHistory = append(estimateHistory, eft.estimates(ef))
-		}
 	}
 }