core: remove transactions from MemPool when they're stored in block

Extend MemPool with Remove() and use it. Fixes #446.

pkg/core/blockchain.go

@@ -467,6 +467,9 @@ func (bc *Blockchain) storeBlock(block *Block) error {
 	}
 
 	atomic.StoreUint32(&bc.blockHeight, block.Index)
+	for _, tx := range block.Transactions {
+		bc.memPool.Remove(tx.Hash())
+	}
 	return nil
 }
 

pkg/core/mem_pool.go

@@ -132,6 +132,39 @@ func (mp MemPool) TryAdd(hash util.Uint256, pItem *PoolItem) bool {
 	return ok
 }
 
+// Remove removes an item from the mempool, if it exists there (and does
+// nothing if it doesn't).
+func (mp *MemPool) Remove(hash util.Uint256) {
+	var mapAndPools = []struct {
+		unsortedMap map[util.Uint256]*PoolItem
+		sortedPools []*PoolItems
+	}{
+		{unsortedMap: mp.unsortedTxn, sortedPools: []*PoolItems{&mp.sortedHighPrioTxn, &mp.sortedLowPrioTxn}},
+		{unsortedMap: mp.unverifiedTxn, sortedPools: []*PoolItems{&mp.unverifiedSortedHighPrioTxn, &mp.unverifiedSortedLowPrioTxn}},
+	}
+	mp.lock.Lock()
+	for _, mapAndPool := range mapAndPools {
+		if _, ok := mapAndPool.unsortedMap[hash]; ok {
+			delete(mapAndPool.unsortedMap, hash)
+			for _, pool := range mapAndPool.sortedPools {
+				var num int
+				var item *PoolItem
+				for num, item = range *pool {
+					if hash.Equals(item.txn.Hash()) {
+						break
+					}
+				}
+				if num < len(*pool)-1 {
+					*pool = append((*pool)[:num], (*pool)[num+1:]...)
+				} else if num == len(*pool)-1 {
+					*pool = (*pool)[:num]
+				}
+			}
+		}
+	}
+	mp.lock.Unlock()
+}
+
 // RemoveOverCapacity removes transactions with lowest fees until the total number of transactions
 // in the MemPool is within the capacity of the MemPool.
 func (mp *MemPool) RemoveOverCapacity() {

pkg/core/mem_pool_test.go

@@ -0,0 +1,64 @@
+package core
+
+import (
+	"testing"
+
+	"github.com/CityOfZion/neo-go/pkg/core/transaction"
+	"github.com/CityOfZion/neo-go/pkg/util"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type FeerStub struct {
+	lowPriority bool
+	sysFee      util.Fixed8
+	netFee      util.Fixed8
+	perByteFee  util.Fixed8
+}
+
+func (fs *FeerStub) NetworkFee(*transaction.Transaction) util.Fixed8 {
+	return fs.netFee
+}
+
+func (fs *FeerStub) IsLowPriority(*transaction.Transaction) bool {
+	return fs.lowPriority
+}
+
+func (fs *FeerStub) FeePerByte(*transaction.Transaction) util.Fixed8 {
+	return fs.perByteFee
+}
+
+func (fs *FeerStub) SystemFee(*transaction.Transaction) util.Fixed8 {
+	return fs.sysFee
+}
+
+func testMemPoolAddRemoveWithFeer(t *testing.T, fs Feer) {
+	mp := NewMemPool(10)
+	tx := newMinerTX()
+	item := NewPoolItem(tx, fs)
+	_, ok := mp.TryGetValue(tx.Hash())
+	require.Equal(t, false, ok)
+	require.Equal(t, true, mp.TryAdd(tx.Hash(), item))
+	// Re-adding should fail.
+	require.Equal(t, false, mp.TryAdd(tx.Hash(), item))
+	tx2, ok := mp.TryGetValue(tx.Hash())
+	require.Equal(t, true, ok)
+	require.Equal(t, tx, tx2)
+	mp.Remove(tx.Hash())
+	_, ok = mp.TryGetValue(tx.Hash())
+	require.Equal(t, false, ok)
+	// Make sure nothing left in the mempool after removal.
+	assert.Equal(t, 0, len(mp.unsortedTxn))
+	assert.Equal(t, 0, len(mp.unverifiedTxn))
+	assert.Equal(t, 0, len(mp.sortedHighPrioTxn))
+	assert.Equal(t, 0, len(mp.sortedLowPrioTxn))
+	assert.Equal(t, 0, len(mp.unverifiedSortedHighPrioTxn))
+	assert.Equal(t, 0, len(mp.unverifiedSortedLowPrioTxn))
+}
+
+func TestMemPoolAddRemove(t *testing.T) {
+	var fs = &FeerStub{lowPriority: false}
+	t.Run("low priority", func(t *testing.T) { testMemPoolAddRemoveWithFeer(t, fs) })
+	fs.lowPriority = true
+	t.Run("high priority", func(t *testing.T) { testMemPoolAddRemoveWithFeer(t, fs) })
+}