core/rawdb: add test for PruneTransactionIndex and fix bug

core/rawdb/chain_iterator.go

@@ -366,8 +366,8 @@ func unindexTransactionsForTesting(db ethdb.Database, from uint64, to uint64, in
 
 // PruneTransactionIndex removes all tx index entries below a certain block number.
 func PruneTransactionIndex(db ethdb.Database, pruneBlock uint64) {
-	index := NewTable(db, string(txLookupPrefix))
+	db = NewTable(db, string(txLookupPrefix))
-	iter := index.NewIterator(nil, nil)
+	iter := db.NewIterator(nil, nil)
 	defer iter.Release()
 
 	for iter.Next() {
@@ -378,8 +378,7 @@ func PruneTransactionIndex(db ethdb.Database, pruneBlock uint64) {
 		}
 		bn := decodeNumber(v)
 		if bn < pruneBlock {
-			fmt.Println("found entry", iter.Key())
+			db.Delete(iter.Key())
-			// db.Delete(iter.Key())
 		}
 	}
 }

core/rawdb/chain_iterator_test.go

@@ -25,6 +25,7 @@ import (
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/ethdb"
 )
 
 func TestChainIterator(t *testing.T) {
@@ -102,19 +103,18 @@ func TestChainIterator(t *testing.T) {
 	}
 }
 
-func TestIndexTransactions(t *testing.T) {
+func initDatabaseWithTransactions(db ethdb.Database) ([]*types.Block, []*types.Transaction) {
-	// Construct test chain db
+	var blocks []*types.Block
-	chainDb := NewMemoryDatabase()
-
-	var block *types.Block
 	var txs []*types.Transaction
 	to := common.BytesToAddress([]byte{0x11})
 
 	// Write empty genesis block
-	block = types.NewBlock(&types.Header{Number: big.NewInt(int64(0))}, nil, nil, newTestHasher())
+	block := types.NewBlock(&types.Header{Number: big.NewInt(int64(0))}, nil, nil, newTestHasher())
-	WriteBlock(chainDb, block)
+	WriteBlock(db, block)
-	WriteCanonicalHash(chainDb, block.Hash(), block.NumberU64())
+	WriteCanonicalHash(db, block.Hash(), block.NumberU64())
+	blocks = append(blocks, block)
 
+	// Create transactions.
 	for i := uint64(1); i <= 10; i++ {
 		var tx *types.Transaction
 		if i%2 == 0 {
@@ -138,10 +138,21 @@ func TestIndexTransactions(t *testing.T) {
 			})
 		}
 		txs = append(txs, tx)
-		block = types.NewBlock(&types.Header{Number: big.NewInt(int64(i))}, &types.Body{Transactions: types.Transactions{tx}}, nil, newTestHasher())
+		block := types.NewBlock(&types.Header{Number: big.NewInt(int64(i))}, &types.Body{Transactions: types.Transactions{tx}}, nil, newTestHasher())
-		WriteBlock(chainDb, block)
+		WriteBlock(db, block)
-		WriteCanonicalHash(chainDb, block.Hash(), block.NumberU64())
+		WriteCanonicalHash(db, block.Hash(), block.NumberU64())
+		blocks = append(blocks, block)
 	}
+
+	return blocks, txs
+}
+
+func TestIndexTransactions(t *testing.T) {
+	// Construct test chain db
+	chainDB := NewMemoryDatabase()
+
+	_, txs := initDatabaseWithTransactions(chainDB)
+
 	// verify checks whether the tx indices in the range [from, to)
 	// is expected.
 	verify := func(from, to int, exist bool, tail uint64) {
@@ -149,7 +160,7 @@ func TestIndexTransactions(t *testing.T) {
 			if i == 0 {
 				continue
 			}
-			number := ReadTxLookupEntry(chainDb, txs[i-1].Hash())
+			number := ReadTxLookupEntry(chainDB, txs[i-1].Hash())
 			if exist && number == nil {
 				t.Fatalf("Transaction index %d missing", i)
 			}
@@ -157,29 +168,29 @@ func TestIndexTransactions(t *testing.T) {
 				t.Fatalf("Transaction index %d is not deleted", i)
 			}
 		}
-		number := ReadTxIndexTail(chainDb)
+		number := ReadTxIndexTail(chainDB)
 		if number == nil || *number != tail {
 			t.Fatalf("Transaction tail mismatch")
 		}
 	}
-	IndexTransactions(chainDb, 5, 11, nil, false)
+	IndexTransactions(chainDB, 5, 11, nil, false)
 	verify(5, 11, true, 5)
 	verify(0, 5, false, 5)
 
-	IndexTransactions(chainDb, 0, 5, nil, false)
+	IndexTransactions(chainDB, 0, 5, nil, false)
 	verify(0, 11, true, 0)
 
-	UnindexTransactions(chainDb, 0, 5, nil, false)
+	UnindexTransactions(chainDB, 0, 5, nil, false)
 	verify(5, 11, true, 5)
 	verify(0, 5, false, 5)
 
-	UnindexTransactions(chainDb, 5, 11, nil, false)
+	UnindexTransactions(chainDB, 5, 11, nil, false)
 	verify(0, 11, false, 11)
 
 	// Testing corner cases
 	signal := make(chan struct{})
 	var once sync.Once
-	indexTransactionsForTesting(chainDb, 5, 11, signal, func(n uint64) bool {
+	indexTransactionsForTesting(chainDB, 5, 11, signal, func(n uint64) bool {
 		if n <= 8 {
 			once.Do(func() {
 				close(signal)
@@ -190,11 +201,11 @@ func TestIndexTransactions(t *testing.T) {
 	})
 	verify(9, 11, true, 9)
 	verify(0, 9, false, 9)
-	IndexTransactions(chainDb, 0, 9, nil, false)
+	IndexTransactions(chainDB, 0, 9, nil, false)
 
 	signal = make(chan struct{})
 	var once2 sync.Once
-	unindexTransactionsForTesting(chainDb, 0, 11, signal, func(n uint64) bool {
+	unindexTransactionsForTesting(chainDB, 0, 11, signal, func(n uint64) bool {
 		if n >= 8 {
 			once2.Do(func() {
 				close(signal)
@@ -206,3 +217,36 @@ func TestIndexTransactions(t *testing.T) {
 	verify(8, 11, true, 8)
 	verify(0, 8, false, 8)
 }
+
+func TestPruneTransactionIndex(t *testing.T) {
+	chainDB := NewMemoryDatabase()
+	blocks, _ := initDatabaseWithTransactions(chainDB)
+	lastBlock := blocks[len(blocks)-1].NumberU64()
+	pruneBlock := lastBlock - 3
+
+	IndexTransactions(chainDB, 0, lastBlock+1, nil, false)
+
+	// Check all transactions are in index.
+	for _, block := range blocks {
+		for _, tx := range block.Transactions() {
+			num := ReadTxLookupEntry(chainDB, tx.Hash())
+			if num == nil || *num != block.NumberU64() {
+				t.Fatalf("wrong TxLookup entry: %x -> %v", tx.Hash(), num)
+			}
+		}
+	}
+
+	PruneTransactionIndex(chainDB, pruneBlock)
+
+	// Check transactions from old blocks not included.
+	for _, block := range blocks {
+		for _, tx := range block.Transactions() {
+			num := ReadTxLookupEntry(chainDB, tx.Hash())
+			if block.NumberU64() < pruneBlock {
+				if num != nil {
+					t.Fatalf("TxLookup entry not removed: %x -> %v", tx.Hash(), num)
+				}
+			}
+		}
+	}
+}