set state root

Signed-off-by: jsvisa <delweng@gmail.com>

core/state/statesize.go

@@ -17,7 +17,8 @@
 package state
 
 import (
-	"sync"
+	"context"
+	"sync/atomic"
 	"time"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -28,6 +29,7 @@ import (
 	"github.com/ethereum/go-ethereum/metrics"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/triedb"
+	"golang.org/x/sync/errgroup"
 )
 
 // State size metrics
@@ -62,16 +64,18 @@ type StateSizeGenerator struct {
 	triedb *triedb.Database
 
 	// Generator state
-	running bool
 	abort chan struct{}
 	done  chan struct{}
 
 	// Async message channel for updates
 	updateChan chan *stateUpdate
 
-	// Metrics state
+	// Metrics state (only modified by generate() goroutine)
 	metrics  *StateSizeMetrics
 	buffered *StateSizeMetrics
+
+	// Initialization state
+	initialized atomic.Bool
 }
 
 // NewStateSizeGenerator creates a new state size generator and starts it automatically
@@ -87,98 +91,105 @@ func NewStateSizeGenerator(db ethdb.KeyValueStore, triedb *triedb.Database, root
 	}
 
 	// Start the generator automatically
-	g.running = true
 	go g.generate()
 
 	return g
 }
 
-// stop terminates the background generation
+// Stop terminates the background generation
 func (g *StateSizeGenerator) Stop() {
-	if !g.running {
-		return
-	}
-
-	// Signal the goroutine to stop
 	close(g.abort)
 
-	// Wait for the goroutine to actually finish
 	<-g.done
 
-	// Now it's safe to persist metrics since the goroutine has stopped
-	g.running = false
+	// Persist metrics after all the goroutines were stopped
 	g.persistMetrics()
 }
 
-// isRunning returns true if the generator is currently running
-func (g *StateSizeGenerator) IsRunning() bool {
-	return g.running
-}
-
-// waitForCompletion waits for the generator to complete (useful for testing or graceful shutdown)
-func (g *StateSizeGenerator) WaitForCompletion() {
-	if g.running {
-		<-g.done
-	}
-}
-
 // generate performs the state size initialization and handles updates
 func (g *StateSizeGenerator) generate() {
 	defer close(g.done)
 
-	var inited bool
-	var initDone chan struct{}
-
-	if g.hasExistingMetrics() {
-		log.Info("State size metrics already initialized")
-		inited = true
-	}
-
-	// Wait for snapshot generator to complete
-	snapDone := make(chan struct{})
-	go func() {
-		defer close(snapDone)
-
-		for !g.triedb.SnapshotCompleted() {
-			select {
-			case <-g.abort:
-				log.Info("State size generation aborted during snapshot")
-				return
-			default:
-				time.Sleep(10 * time.Second)
-			}
-		}
-	}()
+	initDone := g.initialize()
 
 	for {
 		select {
 		case update := <-g.updateChan:
-			g.handleUpdate(update, inited)
+			g.handleUpdate(update, g.initialized.Load())
+
 		case <-g.abort:
 			log.Info("State size generation aborted")
-			// Wait for initialization goroutine to finish if it's running
+
+			// Wait for initialization to complete with timeout
 			if initDone != nil {
 				select {
 				case <-initDone:
+					log.Debug("Initialization completed before abort")
 				case <-time.After(5 * time.Second):
-					log.Warn("Initialization goroutine did not finish in time")
+					log.Warn("Initialization did not finish in time during abort")
 				}
 			}
 			return
-		case <-snapDone:
-			if !inited {
-				initDone = make(chan struct{})
+
+		case <-initDone:
+			// Initialization completed, merge buffered metrics
+			g.mergeBufferedMetrics()
+			initDone = nil // Clear the channel to prevent future selects
+		}
+	}
+}
+
+// initialize starts the initialization process if not already initialized
+func (g *StateSizeGenerator) initialize() chan struct{} {
+	initDone := make(chan struct{})
+
+	// Check if we already have existing metrics
+	if g.hasExistingMetrics() {
+		log.Info("State size metrics already initialized")
+		g.initialized.Store(true)
+		close(initDone)
+		return initDone
+	}
+
+	// Wait for snapshot completion and then initialize
 	go func() {
 		defer close(initDone)
+
+	LOOP:
+		// Wait for snapshot generator to complete first
+		for {
+			select {
+			case <-g.abort:
+				log.Info("State size initialization aborted during snapshot wait")
+				return
+			default:
+				root, done := g.triedb.SnapshotCompleted()
+				if done {
+					g.metrics.Root = root
+					break LOOP
+				}
+				time.Sleep(10 * time.Second)
+			}
+		}
+
+		// Start actual initialization
 		start := time.Now()
 		log.Info("Starting state size initialization")
-					g.initializeMetrics()
-					log.Info("Completed state size initialization", "elapsed", time.Since(start))
-					inited = true
-				}()
+		if err := g.initializeMetrics(); err != nil {
+			log.Error("Failed to initialize state size metrics", "err", err)
+			return
 		}
-		case <-initDone:
-			// Initialization completed, merge buffered metrics if needed
+
+		g.initialized.Store(true)
+
+		log.Info("Completed state size initialization", "elapsed", time.Since(start))
+	}()
+
+	return initDone
+}
+
+// mergeBufferedMetrics merges buffered metrics into main metrics
+func (g *StateSizeGenerator) mergeBufferedMetrics() {
 	if g.buffered != nil && g.buffered.Root != (common.Hash{}) {
 		log.Info("Merging buffered metrics into main metrics")
 		g.metrics.AccountCount += g.buffered.AccountCount
@@ -189,41 +200,47 @@ func (g *StateSizeGenerator) generate() {
 		g.metrics.TrieNodeBytes += g.buffered.TrieNodeBytes
 		g.metrics.ContractCount += g.buffered.ContractCount
 		g.metrics.ContractBytes += g.buffered.ContractBytes
+
 		// Reset buffered metrics
 		g.buffered = &StateSizeMetrics{Root: g.metrics.Root}
 	}
-			initDone = nil // Clear the channel
-		}
-	}
 }
 
 // handleUpdate processes a single update with proper root continuity checking
-func (g *StateSizeGenerator) handleUpdate(update *stateUpdate, inited bool) {
-	// TODO: Check if the update root matches the current metrics root
-
+func (g *StateSizeGenerator) handleUpdate(update *stateUpdate, initialized bool) {
 	// Calculate the diff
 	diff := g.calculateUpdateDiff(update)
 
-	var m *StateSizeMetrics
-	if inited {
-		m = g.metrics
+	var targetMetrics *StateSizeMetrics
+	if initialized {
+		targetMetrics = g.metrics
 	} else {
-		m = g.buffered
+		targetMetrics = g.buffered
 	}
 
-	// TODO: When to merge the buffered metrics into the main metrics
-	m.Root = update.root
-	m.AccountCount += diff.AccountCount
-	m.AccountBytes += diff.AccountBytes
-	m.StorageCount += diff.StorageCount
-	m.StorageBytes += diff.StorageBytes
-	m.TrieNodeCount += diff.TrieNodeCount
-	m.TrieNodeBytes += diff.TrieNodeBytes
-	m.ContractCount += diff.ContractCount
-	m.ContractBytes += diff.ContractBytes
+	// Check root continuity - the update should build on our current state
+	if targetMetrics.Root != (common.Hash{}) && targetMetrics.Root != update.originRoot {
+		log.Warn("State update root discontinuity detected",
+			"current", targetMetrics.Root,
+			"updateOrigin", update.originRoot,
+			"updateNew", update.root)
+		// For now, we accept the discontinuity but log it
+		// In production, you might want to reset metrics or handle differently
+	}
 
-	// Fire the metrics only if the initialization is done
-	if inited {
+	// Update to the new state root
+	targetMetrics.Root = update.root
+	targetMetrics.AccountCount += diff.AccountCount
+	targetMetrics.AccountBytes += diff.AccountBytes
+	targetMetrics.StorageCount += diff.StorageCount
+	targetMetrics.StorageBytes += diff.StorageBytes
+	targetMetrics.TrieNodeCount += diff.TrieNodeCount
+	targetMetrics.TrieNodeBytes += diff.TrieNodeBytes
+	targetMetrics.ContractCount += diff.ContractCount
+	targetMetrics.ContractBytes += diff.ContractBytes
+
+	// Fire the metrics and persist only if initialization is done
+	if initialized {
 		g.updateMetrics()
 		g.persistMetrics()
 	}
@@ -313,25 +330,60 @@ func (g *StateSizeGenerator) calculateUpdateDiff(update *stateUpdate) StateSizeM
 
 // Track is an async method used to send the state update to the generator.
 // It ignores empty updates (where no state changes occurred).
+// If the channel is full, it drops the update to avoid blocking.
 func (g *StateSizeGenerator) Track(update *stateUpdate) {
 	if update == nil || update.empty() {
 		return
 	}
+
 	g.updateChan <- update
 }
 
 // hasExistingMetrics checks if state size metrics already exist in the database
+// and if they are continuous with the current root
 func (g *StateSizeGenerator) hasExistingMetrics() bool {
-	// Check for existing metrics by looking for a marker key
-	marker := rawdb.ReadStateSizeMetrics(g.db)
-	// TODO: check if the marker's root is the same as the current root
-	return marker != nil
+	data := rawdb.ReadStateSizeMetrics(g.db)
+	if data == nil {
+		return false
+	}
+
+	var existed StateSizeMetrics
+	if err := rlp.DecodeBytes(data, &existed); err != nil {
+		log.Warn("Failed to decode existing state size metrics", "err", err)
+		return false
+	}
+
+	// Check if the existing metrics root matches our current root
+	if (g.metrics.Root != common.Hash{}) && existed.Root != g.metrics.Root {
+		log.Info("Existing state size metrics found but root mismatch", "existing", existed.Root, "current", g.metrics.Root)
+		return false
+	}
+
+	// Root matches - load the existing metrics
+	log.Info("Loading existing state size metrics", "root", existed.Root)
+	g.metrics = &existed
+	return true
 }
 
-// initializeMetrics performs the actual metrics initialization
-func (g *StateSizeGenerator) initializeMetrics() {
+// initializeMetrics performs the actual metrics initialization using errgroup
+func (g *StateSizeGenerator) initializeMetrics() error {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	go func() {
+		select {
+		case <-g.abort:
+			cancel() // Cancel context when abort is signaled
+		case <-ctx.Done():
+			// Context already cancelled
+		}
+	}()
+
+	// Create errgroup with context
+	group, ctx := errgroup.WithContext(ctx)
+
+	// Metrics will be directly updated by each goroutine
 	var (
-		wg                                         sync.WaitGroup
 		accountCount, accountBytes                 uint64
 		storageCount, storageBytes                 uint64
 		trieAccountNodeCount, trieAccountNodeBytes uint64
@@ -339,49 +391,58 @@ func (g *StateSizeGenerator) initializeMetrics() {
 		contractCount, contractBytes               uint64
 	)
 
-	iterate := func(prefix []byte, name string, count, bytes uint64) {
-		defer wg.Done()
-
-		log.Info("Iterating over state size", "table", name)
-		defer func(st time.Time) {
-			log.Info("Finished iterating over state size", "table", name, "count", count, "bytes", bytes, "elapsed", common.PrettyDuration(time.Since(st)))
-		}(time.Now())
-
-		iter := g.db.NewIterator(prefix, nil)
-		defer iter.Release()
-		for iter.Next() {
-			count++
-			bytes += uint64(len(iter.Key()) + len(iter.Value()))
-
-			// Check for abort
-			select {
-			case <-g.abort:
-				return
-			default:
+	// Start all table iterations concurrently with direct metric updates
+	group.Go(func() error {
+		count, bytes, err := g.iterateTable(ctx, rawdb.SnapshotAccountPrefix, "account")
+		if err != nil {
+			return err
 		}
+		accountCount, accountBytes = count, bytes
+		return nil
+	})
+
+	group.Go(func() error {
+		count, bytes, err := g.iterateTable(ctx, rawdb.SnapshotStoragePrefix, "storage")
+		if err != nil {
+			return err
 		}
+		storageCount, storageBytes = count, bytes
+		return nil
+	})
+
+	group.Go(func() error {
+		count, bytes, err := g.iterateTable(ctx, rawdb.TrieNodeAccountPrefix, "trie account node")
+		if err != nil {
+			return err
+		}
+		trieAccountNodeCount, trieAccountNodeBytes = count, bytes
+		return nil
+	})
+
+	group.Go(func() error {
+		count, bytes, err := g.iterateTable(ctx, rawdb.TrieNodeStoragePrefix, "trie storage node")
+		if err != nil {
+			return err
+		}
+		trieStorageNodeCount, trieStorageNodeBytes = count, bytes
+		return nil
+	})
+
+	group.Go(func() error {
+		count, bytes, err := g.iterateTable(ctx, rawdb.CodePrefix, "contract code")
+		if err != nil {
+			return err
+		}
+		contractCount, contractBytes = count, bytes
+		return nil
+	})
+
+	// Wait for all goroutines to complete
+	if err := group.Wait(); err != nil {
+		return err
 	}
 
-	tables := []struct {
-		prefix []byte
-		name   string
-		count  *uint64
-		bytes  *uint64
-	}{
-		{rawdb.SnapshotAccountPrefix, "account", &accountCount, &accountBytes},
-		{rawdb.SnapshotStoragePrefix, "storage", &storageCount, &storageBytes},
-		{rawdb.TrieNodeAccountPrefix, "trie account node", &trieAccountNodeCount, &trieAccountNodeBytes},
-		{rawdb.TrieNodeStoragePrefix, "trie storage node", &trieStorageNodeCount, &trieStorageNodeBytes},
-		{rawdb.CodePrefix, "contract code", &contractCount, &contractBytes},
-	}
-	wg.Add(len(tables))
-	for _, table := range tables {
-		go iterate(table.prefix, table.name, *table.count, *table.bytes)
-	}
-
-	wg.Wait()
-
-	// Update metrics
+	// Update metrics (safe since we're in the single writer goroutine)
 	g.metrics.AccountCount = accountCount
 	g.metrics.AccountBytes = accountBytes
 	g.metrics.StorageCount = storageCount
@@ -393,10 +454,46 @@ func (g *StateSizeGenerator) initializeMetrics() {
 
 	g.updateMetrics()
 	g.persistMetrics()
+
+	return nil
+}
+
+// iterateTable performs iteration over a specific table and returns the results
+func (g *StateSizeGenerator) iterateTable(ctx context.Context, prefix []byte, name string) (uint64, uint64, error) {
+	log.Info("Iterating over state size", "table", name)
+	start := time.Now()
+
+	var count, bytes uint64
+	iter := g.db.NewIterator(prefix, nil)
+	defer iter.Release()
+
+	for iter.Next() {
+		count++
+		bytes += uint64(len(iter.Key()) + len(iter.Value()))
+
+		// Check for cancellation periodically for performance
+		if count%10000 == 0 {
+			select {
+			case <-ctx.Done():
+				log.Info("State size iteration cancelled", "table", name, "count", count)
+				return 0, 0, ctx.Err()
+			default:
+			}
+		}
+	}
+
+	// Check for iterator errors
+	if err := iter.Error(); err != nil {
+		log.Error("Iterator error during state size calculation", "table", name, "err", err)
+		return 0, 0, err
+	}
+
+	log.Info("Finished iterating over state size", "table", name, "count", count, "bytes", bytes, "elapsed", common.PrettyDuration(time.Since(start)))
+
+	return count, bytes, nil
 }
 
 func (g *StateSizeGenerator) updateMetrics() {
-	// Update global metrics
 	stateSizeAccountsCountMeter.Mark(int64(g.metrics.AccountCount))
 	stateSizeAccountsBytesMeter.Mark(int64(g.metrics.AccountBytes))
 	stateSizeStorageCountMeter.Mark(int64(g.metrics.StorageCount))

triedb/database.go

@@ -377,10 +377,10 @@ func (db *Database) Disk() ethdb.Database {
 }
 
 // SnapshotCompleted returns the indicator if the snapshot is completed.
-func (db *Database) SnapshotCompleted() bool {
+func (db *Database) SnapshotCompleted() (common.Hash, bool) {
 	pdb, ok := db.backend.(*pathdb.Database)
 	if !ok {
-		return false
+		return common.Hash{}, false
 	}
 	return pdb.SnapshotCompleted()
 }

triedb/pathdb/database.go

@@ -682,7 +682,14 @@ func (db *Database) StorageIterator(root common.Hash, account common.Hash, seek
 	return newFastStorageIterator(db, root, account, seek)
 }
 
-// SnapshotCompleted returns the indicator if the snapshot generation is completed.
-func (db *Database) SnapshotCompleted() bool {
-	return !db.waitSync && db.tree.bottom().genComplete()
+// SnapshotCompleted returns the snapshot root if the snapshot generation is completed.
+func (db *Database) SnapshotCompleted() (common.Hash, bool) {
+	if db.waitSync {
+		return common.Hash{}, false
+	}
+	dl := db.tree.bottom()
+	if dl.genComplete() {
+		return dl.rootHash(), true
+	}
+	return common.Hash{}, false
 }