core/state: improve test

core/state/statedb_test.go

@@ -974,23 +974,21 @@ func TestMissingTrieNodes(t *testing.T) {
 func testMissingTrieNodes(t *testing.T, scheme string) {
 	// Create an initial state with a few accounts
 	var (
-		tdb    *triedb.Database
+		tdb   *triedb.Database
-		memDb  = rawdb.NewMemoryDatabase()
+		memDb = rawdb.NewMemoryDatabase()
-		openDb = func() *triedb.Database {
-			if scheme == rawdb.PathScheme {
-				return triedb.NewDatabase(memDb, &triedb.Config{PathDB: &pathdb.Config{
-					TrieCleanSize:   0,
-					StateCleanSize:  0,
-					WriteBufferSize: 0,
-				}}) // disable caching
-			} else {
-				return triedb.NewDatabase(memDb, &triedb.Config{HashDB: &hashdb.Config{
-					CleanCacheSize: 0,
-				}}) // disable caching
-			}
-		}
 	)
-	tdb = openDb()
+	if scheme == rawdb.PathScheme {
+		tdb = triedb.NewDatabase(memDb, &triedb.Config{PathDB: &pathdb.Config{
+			TrieCleanSize:   0,
+			StateCleanSize:  0,
+			WriteBufferSize: 0,
+			NoAsyncFlush:    true,
+		}}) // disable caching
+	} else {
+		tdb = triedb.NewDatabase(memDb, &triedb.Config{HashDB: &hashdb.Config{
+			CleanCacheSize: 0,
+		}}) // disable caching
+	}
 	db := NewDatabase(tdb, nil)
 
 	var root common.Hash
@@ -1007,7 +1005,6 @@ func testMissingTrieNodes(t *testing.T, scheme string) {
 		// force-flush
 		tdb.Commit(root, false)
 	}
-	// Create a new state on the old root
 	// Now we clear out the memdb
 	it := memDb.NewIterator(nil, nil)
 	for it.Next() {
@@ -1026,8 +1023,6 @@ func testMissingTrieNodes(t *testing.T, scheme string) {
 			}
 		}
 	}
-	tdb = openDb()
-	db = NewDatabase(tdb, nil)
 	state, _ = New(root, db)
 	balance := state.GetBalance(addr)
 	// The removed elem should lead to it returning zero balance

triedb/pathdb/database.go

@@ -121,8 +121,9 @@ type Config struct {
 	ReadOnly        bool   // Flag whether the database is opened in read only mode
 
 	// Testing configurations
-	SnapshotNoBuild bool // Flag Whether the background generation is allowed
+	SnapshotNoBuild   bool // Flag Whether the state generation is allowed
-	NoAsyncFlush    bool // Flag whether the background generation is allowed
+	NoAsyncFlush      bool // Flag whether the background buffer flushing is allowed
+	NoAsyncGeneration bool // Flag whether the background generation is allowed
 }
 
 // sanitize checks the provided user configurations and changes anything that's
@@ -369,6 +370,12 @@ func (db *Database) setStateGenerator() error {
 	}
 	stats.log("Starting snapshot generation", root, generator.Marker)
 	dl.generator.run(root)
+
+	// Block until the generation completes. It's the feature used in
+	// unit tests.
+	if db.config.NoAsyncGeneration {
+		<-dl.generator.done
+	}
 	return nil
 }
 
@@ -667,19 +674,6 @@ func (db *Database) HistoryRange() (uint64, uint64, error) {
 	return historyRange(db.freezer)
 }
 
-// waitGeneration waits until the background generation is finished. It assumes
-// that the generation is permitted; otherwise, it will block indefinitely.
-func (db *Database) waitGeneration() {
-	db.lock.RLock()
-	defer db.lock.RUnlock()
-
-	gen := db.tree.bottom().generator
-	if gen == nil || gen.completed() {
-		return
-	}
-	<-gen.done
-}
-
 // AccountIterator creates a new account iterator for the specified root hash and
 // seeks to a starting account hash.
 func (db *Database) AccountIterator(root common.Hash, seek common.Hash) (AccountIterator, error) {

triedb/pathdb/disklayer.go

@@ -581,6 +581,17 @@ func (dl *diskLayer) genComplete() bool {
 	return dl.genMarker() == nil
 }
 
+// waitFlush blocks until the background buffer flush is completed.
+func (dl *diskLayer) waitFlush() error {
+	dl.lock.RLock()
+	defer dl.lock.RUnlock()
+
+	if dl.frozen == nil {
+		return nil
+	}
+	return dl.frozen.waitFlush()
+}
+
 // terminate releases the frozen buffer if it's not nil and terminates the
 // background state generator.
 func (dl *diskLayer) terminate() error {

triedb/pathdb/iterator_binary.go

@@ -45,6 +45,10 @@ type binaryIterator struct {
 // accounts in a slow, but easily verifiable way. Note this function is used
 // for initialization, use `newBinaryAccountIterator` as the API.
 func (dl *diskLayer) initBinaryAccountIterator(seek common.Hash) *binaryIterator {
+	// Block until the frozen buffer flushing is completed.
+	if err := dl.waitFlush(); err != nil {
+		panic(err)
+	}
 	// The state set in the disk layer is mutable, hold the lock before obtaining
 	// the account list to prevent concurrent map iteration and write.
 	dl.lock.RLock()
@@ -113,6 +117,10 @@ func (dl *diffLayer) initBinaryAccountIterator(seek common.Hash) *binaryIterator
 // storage slots in a slow, but easily verifiable way. Note this function is used
 // for initialization, use `newBinaryStorageIterator` as the API.
 func (dl *diskLayer) initBinaryStorageIterator(account common.Hash, seek common.Hash) *binaryIterator {
+	// Block until the frozen buffer flushing is completed.
+	if err := dl.waitFlush(); err != nil {
+		panic(err)
+	}
 	// The state set in the disk layer is mutable, hold the lock before obtaining
 	// the storage list to prevent concurrent map iteration and write.
 	dl.lock.RLock()

triedb/pathdb/iterator_fast.go

@@ -76,6 +76,11 @@ func newFastIterator(db *Database, root common.Hash, account common.Hash, seek c
 		if accountIterator {
 			switch dl := current.(type) {
 			case *diskLayer:
+				// Ensure no active background buffer flush is in progress, otherwise,
+				// part of the state data may become invisible.
+				if err := dl.waitFlush(); err != nil {
+					return nil, err
+				}
 				// The state set in the disk layer is mutable, hold the lock before obtaining
 				// the account list to prevent concurrent map iteration and write.
 				dl.lock.RLock()
@@ -113,6 +118,11 @@ func newFastIterator(db *Database, root common.Hash, account common.Hash, seek c
 		} else {
 			switch dl := current.(type) {
 			case *diskLayer:
+				// Ensure no active background buffer flush is in progress, otherwise,
+				// part of the state data may become invisible.
+				if err := dl.waitFlush(); err != nil {
+					return nil, err
+				}
 				// The state set in the disk layer is mutable, hold the lock before obtaining
 				// the storage list to prevent concurrent map iteration and write.
 				dl.lock.RLock()

triedb/pathdb/iterator_test.go

@@ -254,11 +254,9 @@ func TestFastIteratorBasics(t *testing.T) {
 // TestAccountIteratorTraversal tests some simple multi-layer iteration.
 func TestAccountIteratorTraversal(t *testing.T) {
 	config := &Config{
-		WriteBufferSize: 0,
+		NoAsyncGeneration: true,
-		NoAsyncFlush:    true,
 	}
 	db := New(rawdb.NewMemoryDatabase(), config, false)
-	db.waitGeneration()
 
 	// Stack three diff layers on top with various overlaps
 	db.Update(common.HexToHash("0x02"), types.EmptyRootHash, 0, trienode.NewMergedNodeSet(),
@@ -299,11 +297,9 @@ func TestAccountIteratorTraversal(t *testing.T) {
 
 func TestStorageIteratorTraversal(t *testing.T) {
 	config := &Config{
-		WriteBufferSize: 0,
+		NoAsyncGeneration: true,
-		NoAsyncFlush:    true,
 	}
 	db := New(rawdb.NewMemoryDatabase(), config, false)
-	db.waitGeneration()
 
 	// Stack three diff layers on top with various overlaps
 	db.Update(common.HexToHash("0x02"), types.EmptyRootHash, 0, trienode.NewMergedNodeSet(),
@@ -313,14 +309,14 @@ func TestStorageIteratorTraversal(t *testing.T) {
 		NewStateSetWithOrigin(randomAccountSet("0xaa"), randomStorageSet([]string{"0xaa"}, [][]string{{"0x04", "0x05", "0x06"}}, nil), nil, nil, false))
 
 	db.Update(common.HexToHash("0x04"), common.HexToHash("0x03"), 0, trienode.NewMergedNodeSet(),
-		NewStateSetWithOrigin(randomAccountSet("0xaa"), randomStorageSet([]string{"0xaa"}, [][]string{{"0x01", "0x02", "0x03"}}, nil), nil, nil, false))
+		NewStateSetWithOrigin(randomAccountSet("0xaa"), randomStorageSet([]string{"0xaa"}, [][]string{{"0x01", "0x02"}}, nil), nil, nil, false))
 
 	// Verify the single and multi-layer iterators
 	head := db.tree.get(common.HexToHash("0x04"))
 
 	// singleLayer: 0x1, 0x2, 0x3
 	diffIter := newDiffStorageIterator(common.HexToHash("0xaa"), common.Hash{}, head.(*diffLayer).states.stateSet.storageList(common.HexToHash("0xaa")), nil)
-	verifyIterator(t, 3, diffIter, verifyNothing)
+	verifyIterator(t, 2, diffIter, verifyNothing)
 
 	// binaryIterator: 0x1, 0x2, 0x3, 0x4, 0x5, 0x6
 	verifyIterator(t, 6, head.(*diffLayer).newBinaryStorageIterator(common.HexToHash("0xaa"), common.Hash{}), verifyStorage)
@@ -344,11 +340,9 @@ func TestStorageIteratorTraversal(t *testing.T) {
 // also expect the correct values to show up.
 func TestAccountIteratorTraversalValues(t *testing.T) {
 	config := &Config{
-		WriteBufferSize: 0,
+		NoAsyncGeneration: true,
-		NoAsyncFlush:    true,
 	}
 	db := New(rawdb.NewMemoryDatabase(), config, false)
-	db.waitGeneration()
 
 	// Create a batch of account sets to seed subsequent layers with
 	var (
@@ -461,11 +455,9 @@ func TestAccountIteratorTraversalValues(t *testing.T) {
 
 func TestStorageIteratorTraversalValues(t *testing.T) {
 	config := &Config{
-		WriteBufferSize: 0,
+		NoAsyncGeneration: true,
-		NoAsyncFlush:    true,
 	}
 	db := New(rawdb.NewMemoryDatabase(), config, false)
-	db.waitGeneration()
 
 	wrapStorage := func(storage map[common.Hash][]byte) map[common.Hash]map[common.Hash][]byte {
 		return map[common.Hash]map[common.Hash][]byte{
@@ -595,11 +587,9 @@ func TestAccountIteratorLargeTraversal(t *testing.T) {
 	}
 	// Build up a large stack of snapshots
 	config := &Config{
-		WriteBufferSize: 0,
+		NoAsyncGeneration: true,
-		NoAsyncFlush:    true,
 	}
 	db := New(rawdb.NewMemoryDatabase(), config, false)
-	db.waitGeneration()
 
 	for i := 1; i < 128; i++ {
 		parent := types.EmptyRootHash
@@ -635,10 +625,10 @@ func TestAccountIteratorLargeTraversal(t *testing.T) {
 // - continues iterating
 func TestAccountIteratorFlattening(t *testing.T) {
 	config := &Config{
-		WriteBufferSize: 10 * 1024,
+		WriteBufferSize:   10 * 1024,
+		NoAsyncGeneration: true,
 	}
 	db := New(rawdb.NewMemoryDatabase(), config, false)
-	db.waitGeneration()
 
 	// Create a stack of diffs on top
 	db.Update(common.HexToHash("0x02"), types.EmptyRootHash, 1, trienode.NewMergedNodeSet(),
@@ -667,12 +657,24 @@ func TestAccountIteratorFlattening(t *testing.T) {
 }
 
 func TestAccountIteratorSeek(t *testing.T) {
+	t.Run("fast", func(t *testing.T) {
+		testAccountIteratorSeek(t, func(db *Database, root, seek common.Hash) AccountIterator {
+			it, _ := db.AccountIterator(root, seek)
+			return it
+		})
+	})
+	t.Run("binary", func(t *testing.T) {
+		testAccountIteratorSeek(t, func(db *Database, root, seek common.Hash) AccountIterator {
+			return db.tree.get(root).(*diffLayer).newBinaryAccountIterator(seek)
+		})
+	})
+}
+
+func testAccountIteratorSeek(t *testing.T, newIterator func(db *Database, root, seek common.Hash) AccountIterator) {
 	config := &Config{
-		WriteBufferSize: 0,
+		NoAsyncGeneration: true,
-		NoAsyncFlush:    true,
 	}
 	db := New(rawdb.NewMemoryDatabase(), config, false)
-	db.waitGeneration()
 
 	db.Update(common.HexToHash("0x02"), types.EmptyRootHash, 1, trienode.NewMergedNodeSet(),
 		NewStateSetWithOrigin(randomAccountSet("0xaa", "0xee", "0xff", "0xf0"), nil, nil, nil, false))
@@ -688,39 +690,39 @@ func TestAccountIteratorSeek(t *testing.T) {
 	// 03: aa, bb, dd, ee, f0 (, f0), ff
 	// 04: aa, bb, cc, dd, ee, f0 (, f0), ff (, ff)
 	// Construct various iterators and ensure their traversal is correct
-	it, _ := db.AccountIterator(common.HexToHash("0x02"), common.HexToHash("0xdd"))
+	it := newIterator(db, common.HexToHash("0x02"), common.HexToHash("0xdd"))
 	defer it.Release()
 	verifyIterator(t, 3, it, verifyAccount) // expected: ee, f0, ff
 
-	it, _ = db.AccountIterator(common.HexToHash("0x02"), common.HexToHash("0xaa"))
+	it = newIterator(db, common.HexToHash("0x02"), common.HexToHash("0xaa"))
 	defer it.Release()
 	verifyIterator(t, 4, it, verifyAccount) // expected: aa, ee, f0, ff
 
-	it, _ = db.AccountIterator(common.HexToHash("0x02"), common.HexToHash("0xff"))
+	it = newIterator(db, common.HexToHash("0x02"), common.HexToHash("0xff"))
 	defer it.Release()
 	verifyIterator(t, 1, it, verifyAccount) // expected: ff
 
-	it, _ = db.AccountIterator(common.HexToHash("0x02"), common.HexToHash("0xff1"))
+	it = newIterator(db, common.HexToHash("0x02"), common.HexToHash("0xff1"))
 	defer it.Release()
 	verifyIterator(t, 0, it, verifyAccount) // expected: nothing
 
-	it, _ = db.AccountIterator(common.HexToHash("0x04"), common.HexToHash("0xbb"))
+	it = newIterator(db, common.HexToHash("0x04"), common.HexToHash("0xbb"))
 	defer it.Release()
 	verifyIterator(t, 6, it, verifyAccount) // expected: bb, cc, dd, ee, f0, ff
 
-	it, _ = db.AccountIterator(common.HexToHash("0x04"), common.HexToHash("0xef"))
+	it = newIterator(db, common.HexToHash("0x04"), common.HexToHash("0xef"))
 	defer it.Release()
 	verifyIterator(t, 2, it, verifyAccount) // expected: f0, ff
 
-	it, _ = db.AccountIterator(common.HexToHash("0x04"), common.HexToHash("0xf0"))
+	it = newIterator(db, common.HexToHash("0x04"), common.HexToHash("0xf0"))
 	defer it.Release()
 	verifyIterator(t, 2, it, verifyAccount) // expected: f0, ff
 
-	it, _ = db.AccountIterator(common.HexToHash("0x04"), common.HexToHash("0xff"))
+	it = newIterator(db, common.HexToHash("0x04"), common.HexToHash("0xff"))
 	defer it.Release()
 	verifyIterator(t, 1, it, verifyAccount) // expected: ff
 
-	it, _ = db.AccountIterator(common.HexToHash("0x04"), common.HexToHash("0xff1"))
+	it = newIterator(db, common.HexToHash("0x04"), common.HexToHash("0xff1"))
 	defer it.Release()
 	verifyIterator(t, 0, it, verifyAccount) // expected: nothing
 }
@@ -741,11 +743,9 @@ func TestStorageIteratorSeek(t *testing.T) {
 
 func testStorageIteratorSeek(t *testing.T, newIterator func(db *Database, root, account, seek common.Hash) StorageIterator) {
 	config := &Config{
-		WriteBufferSize: 0,
+		NoAsyncGeneration: true,
-		NoAsyncFlush:    true,
 	}
 	db := New(rawdb.NewMemoryDatabase(), config, false)
-	db.waitGeneration()
 
 	// Stack three diff layers on top with various overlaps
 	db.Update(common.HexToHash("0x02"), types.EmptyRootHash, 1, trienode.NewMergedNodeSet(),
@@ -814,11 +814,9 @@ func TestAccountIteratorDeletions(t *testing.T) {
 
 func testAccountIteratorDeletions(t *testing.T, newIterator func(db *Database, root, seek common.Hash) AccountIterator) {
 	config := &Config{
-		WriteBufferSize: 0,
+		NoAsyncGeneration: true,
-		NoAsyncFlush:    true,
 	}
 	db := New(rawdb.NewMemoryDatabase(), config, false)
-	db.waitGeneration()
 
 	// Stack three diff layers on top with various overlaps
 	db.Update(common.HexToHash("0x02"), types.EmptyRootHash, 1, trienode.NewMergedNodeSet(),
@@ -855,11 +853,9 @@ func testAccountIteratorDeletions(t *testing.T, newIterator func(db *Database, r
 
 func TestStorageIteratorDeletions(t *testing.T) {
 	config := &Config{
-		WriteBufferSize: 0,
+		NoAsyncGeneration: true,
-		NoAsyncFlush:    true,
 	}
 	db := New(rawdb.NewMemoryDatabase(), config, false)
-	db.waitGeneration()
 
 	// Stack three diff layers on top with various overlaps
 	db.Update(common.HexToHash("0x02"), types.EmptyRootHash, 1, trienode.NewMergedNodeSet(),
@@ -924,11 +920,10 @@ func TestStaleIterator(t *testing.T) {
 
 func testStaleIterator(t *testing.T, newIter func(db *Database, hash common.Hash) Iterator) {
 	config := &Config{
-		WriteBufferSize: 16 * 1024 * 1024,
+		WriteBufferSize:   16 * 1024 * 1024,
-		NoAsyncFlush:    true,
+		NoAsyncGeneration: true,
 	}
 	db := New(rawdb.NewMemoryDatabase(), config, false)
-	db.waitGeneration()
 
 	// [02 (disk), 03]
 	db.Update(common.HexToHash("0x02"), types.EmptyRootHash, 1, trienode.NewMergedNodeSet(),
@@ -980,11 +975,9 @@ func BenchmarkAccountIteratorTraversal(b *testing.B) {
 		return accounts
 	}
 	config := &Config{
-		WriteBufferSize: 0,
+		NoAsyncGeneration: true,
-		NoAsyncFlush:    true,
 	}
 	db := New(rawdb.NewMemoryDatabase(), config, false)
-	db.waitGeneration()
 
 	for i := 1; i <= 100; i++ {
 		parent := types.EmptyRootHash
@@ -1076,11 +1069,9 @@ func BenchmarkAccountIteratorLargeBaselayer(b *testing.B) {
 		return accounts
 	}
 	config := &Config{
-		WriteBufferSize: 0,
+		NoAsyncGeneration: true,
-		NoAsyncFlush:    true,
 	}
 	db := New(rawdb.NewMemoryDatabase(), config, false)
-	db.waitGeneration()
 
 	db.Update(common.HexToHash("0x02"), types.EmptyRootHash, 1, trienode.NewMergedNodeSet(), NewStateSetWithOrigin(makeAccounts(2000), nil, nil, nil, false))
 	for i := 2; i <= 100; i++ {