core: track state size in blockchain.go

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

core/blockchain.go

@@ -334,6 +334,8 @@ type BlockChain struct {
 	processor  Processor // Block transaction processor interface
 	logger     *tracing.Hooks
 
+	stateSizeGen *state.StateSizeGenerator // State size tracking
+
 	lastForkReadyAlert time.Time // Last time there was a fork readiness print out
 }
 
@@ -526,6 +528,11 @@ func NewBlockChain(db ethdb.Database, genesis *Genesis, engine consensus.Engine,
 	if bc.cfg.TxLookupLimit >= 0 {
 		bc.txIndexer = newTxIndexer(uint64(bc.cfg.TxLookupLimit), bc)
 	}
+
+	// Start state size tracker
+	bc.stateSizeGen = state.NewStateSizeGenerator(bc.statedb.DiskDB(), bc.triedb, head.Root)
+	log.Info("Started state size generator", "root", head.Root)
+
 	return bc, nil
 }
 
@@ -1313,6 +1320,12 @@ func (bc *BlockChain) Stop() {
 			}
 		}
 	}
+	// Stop state size generator if running
+	if bc.stateSizeGen != nil {
+		bc.stateSizeGen.Stop()
+		log.Info("Stopped state size generator")
+	}
+
 	// Allow tracers to clean-up and release resources.
 	if bc.logger != nil && bc.logger.OnClose != nil {
 		bc.logger.OnClose()
@@ -1586,10 +1599,15 @@ func (bc *BlockChain) writeBlockWithState(block *types.Block, receipts []*types.
 		log.Crit("Failed to write block into disk", "err", err)
 	}
 	// Commit all cached state changes into underlying memory database.
-	root, err := statedb.Commit(block.NumberU64(), bc.chainConfig.IsEIP158(block.Number()), bc.chainConfig.IsCancun(block.Number(), block.Time()))
+	root, stateUpdate, err := statedb.CommitWithUpdate(block.NumberU64(), bc.chainConfig.IsEIP158(block.Number()), bc.chainConfig.IsCancun(block.Number(), block.Time()))
 	if err != nil {
 		return err
 	}
+
+	// Track state size changes if generator is running
+	if bc.stateSizeGen != nil && stateUpdate != nil {
+		bc.stateSizeGen.Track(stateUpdate)
+	}
 	// If node is running in path mode, skip explicit gc operation
 	// which is unnecessary in this mode.
 	if bc.triedb.Scheme() == rawdb.PathScheme {

core/state/database.go

@@ -285,6 +285,11 @@ func (db *CachingDB) TrieDB() *triedb.Database {
 	return db.triedb
 }
 
+// DiskDB returns the underlying disk database for direct access.
+func (db *CachingDB) DiskDB() ethdb.KeyValueStore {
+	return db.disk
+}
+
 // PointCache returns the cache of evaluated curve points.
 func (db *CachingDB) PointCache() *utils.PointCache {
 	return db.pointCache

core/state/statedb.go

@@ -139,9 +139,6 @@ type StateDB struct {
 	witness      *stateless.Witness
 	witnessStats *stateless.WitnessStats
 
-	// State size tracking
-	stateSizeGen *stateSizeGenerator
-
 	// Measurements gathered during execution for debugging purposes
 	AccountReads    time.Duration
 	AccountHashes   time.Duration
@@ -190,14 +187,6 @@ func NewWithReader(root common.Hash, db Database, reader Reader) (*StateDB, erro
 		sdb.accessEvents = NewAccessEvents(db.PointCache())
 	}
 
-	// Initialize state size tracking
-	if cachingDB, ok := db.(*CachingDB); ok {
-		sdb.stateSizeGen = newStateSizeGenerator(cachingDB.disk, db, root)
-		if !db.TrieDB().IsVerkle() {
-			sdb.stateSizeGen.run()
-		}
-	}
-
 	return sdb, nil
 }
 
@@ -1362,11 +1351,6 @@ func (s *StateDB) commitAndFlush(block uint64, deleteEmptyObjects bool, noStorag
 				return nil, err
 			}
 			s.TrieDBCommits += time.Since(start)
-
-			// Update state size metrics
-			if s.stateSizeGen != nil {
-				s.stateSizeGen.updateMetrics(ret)
-			}
 		}
 	}
 	s.reader, _ = s.db.Reader(s.originalRoot)
@@ -1395,6 +1379,16 @@ func (s *StateDB) Commit(block uint64, deleteEmptyObjects bool, noStorageWiping
 	return ret.root, nil
 }
 
+// CommitWithUpdate writes the state mutations and returns both the root hash and the state update.
+// This is useful for tracking state changes at the blockchain level.
+func (s *StateDB) CommitWithUpdate(block uint64, deleteEmptyObjects bool, noStorageWiping bool) (common.Hash, *stateUpdate, error) {
+	ret, err := s.commitAndFlush(block, deleteEmptyObjects, noStorageWiping)
+	if err != nil {
+		return common.Hash{}, nil, err
+	}
+	return ret.root, ret, nil
+}
+
 // Prepare handles the preparatory steps for executing a state transition with.
 // This method must be invoked before state transition.
 //

core/state/statesize.go

@@ -27,6 +27,7 @@ import (
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/metrics"
 	"github.com/ethereum/go-ethereum/rlp"
+	"github.com/ethereum/go-ethereum/triedb"
 )
 
 // State size metrics
@@ -55,75 +56,272 @@ type StateSizeMetrics struct {
 	ContractBytes uint64
 }
 
-// stateSizeGenerator handles the initialization and tracking of state size metrics
+// StateSizeGenerator handles the initialization and tracking of state size metrics
-type stateSizeGenerator struct {
+type StateSizeGenerator struct {
-	db  ethdb.KeyValueStore
+	db     ethdb.KeyValueStore
-	sdb Database
+	triedb *triedb.Database
 
 	// Generator state
 	running bool
 	abort   chan struct{}
 	done    chan struct{}
 
+	// Async message channel for updates
+	updateChan chan *stateUpdate
+
 	// Metrics state
-	metrics     *StateSizeMetrics
+	metrics  *StateSizeMetrics
-	metricsLock sync.RWMutex
+	buffered *StateSizeMetrics
 }
 
-// newStateSizeGenerator creates a new state size generator
+// NewStateSizeGenerator creates a new state size generator and starts it automatically
-func newStateSizeGenerator(db ethdb.KeyValueStore, sdb Database, root common.Hash) *stateSizeGenerator {
+func NewStateSizeGenerator(db ethdb.KeyValueStore, triedb *triedb.Database, root common.Hash) *StateSizeGenerator {
-	return &stateSizeGenerator{
+	g := &StateSizeGenerator{
-		db:      db,
+		db:         db,
-		sdb:     sdb,
+		triedb:     triedb,
-		abort:   make(chan struct{}),
+		abort:      make(chan struct{}),
-		done:    make(chan struct{}),
+		done:       make(chan struct{}),
-		metrics: &StateSizeMetrics{Root: root},
+		updateChan: make(chan *stateUpdate, 1000), // Buffered channel for updates
+		metrics:    &StateSizeMetrics{Root: root},
+		buffered:   &StateSizeMetrics{Root: root},
 	}
-}
 
-// run starts the state size initialization in the background
+	// Start the generator automatically
-func (g *stateSizeGenerator) run() {
-	if g.running {
-		g.stop()
-		log.Warn("Paused the leftover state size generation cycle")
-	}
 	g.running = true
 	go g.generate()
+
+	return g
 }
 
 // stop terminates the background generation
-func (g *stateSizeGenerator) stop() {
+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
+	g.persistMetrics()
 }
 
-// generate performs the state size initialization
+// isRunning returns true if the generator is currently running
-func (g *stateSizeGenerator) generate() {
+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)
-	start := time.Now()
+
+	var inited bool
+	var initDone chan struct{}
 
 	if g.hasExistingMetrics() {
 		log.Info("State size metrics already initialized")
-		return
+		inited = true
 	}
 
 	// Wait for snapshot generator to complete
-	if db := g.sdb.TrieDB(); db != nil {
+	snapDone := make(chan struct{})
-		for !db.SnapshotCompleted() {
+	go func() {
-			time.Sleep(5 * time.Second)
+		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)
+			}
+		}
+	}()
+
+	for {
+		select {
+		case update := <-g.updateChan:
+			g.handleUpdate(update, inited)
+		case <-g.abort:
+			log.Info("State size generation aborted")
+			// Wait for initialization goroutine to finish if it's running
+			if initDone != nil {
+				select {
+				case <-initDone:
+				case <-time.After(5 * time.Second):
+					log.Warn("Initialization goroutine did not finish in time")
+				}
+			}
+			return
+		case <-snapDone:
+			if !inited {
+				initDone = make(chan struct{})
+				go func() {
+					defer close(initDone)
+					start := time.Now()
+					log.Info("Starting state size initialization")
+					g.initializeMetrics()
+					log.Info("Completed state size initialization", "elapsed", time.Since(start))
+					inited = true
+				}()
+			}
+		case <-initDone:
+			// Initialization completed, merge buffered metrics if needed
+			if g.buffered != nil && g.buffered.Root != (common.Hash{}) {
+				log.Info("Merging buffered metrics into main metrics")
+				g.metrics.AccountCount += g.buffered.AccountCount
+				g.metrics.AccountBytes += g.buffered.AccountBytes
+				g.metrics.StorageCount += g.buffered.StorageCount
+				g.metrics.StorageBytes += g.buffered.StorageBytes
+				g.metrics.TrieNodeCount += g.buffered.TrieNodeCount
+				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
+
+	// Calculate the diff
+	diff := g.calculateUpdateDiff(update)
+
+	var m *StateSizeMetrics
+	if inited {
+		m = g.metrics
+	} else {
+		m = 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
+
+	// Fire the metrics only if the initialization is done
+	if inited {
+		g.updateMetrics()
+		g.persistMetrics()
+	}
+}
+
+// calculateUpdateDiff calculates the diff for a state update
+func (g *StateSizeGenerator) calculateUpdateDiff(update *stateUpdate) StateSizeMetrics {
+	var diff StateSizeMetrics
+
+	// Calculate account changes
+	for addr, oldValue := range update.accountsOrigin {
+		addrHash := crypto.Keccak256Hash(addr.Bytes())
+		newValue, exists := update.accounts[addrHash]
+		if !exists {
+			log.Warn("State update missing account", "address", addr)
+			continue
+		}
+		if len(newValue) == 0 {
+			diff.AccountCount -= 1
+			diff.AccountBytes -= common.HashLength
+		}
+		if len(oldValue) == 0 {
+			diff.AccountCount += 1
+			diff.AccountBytes += common.HashLength
+		}
+		diff.AccountBytes += uint64(len(newValue) - len(oldValue))
+	}
+
+	// Calculate storage changes
+	for addr, slots := range update.storagesOrigin {
+		addrHash := crypto.Keccak256Hash(addr.Bytes())
+		subset, exists := update.storages[addrHash]
+		if !exists {
+			log.Warn("State update missing storage", "address", addr)
+			continue
+		}
+		for key, oldValue := range slots {
+			var (
+				exists   bool
+				newValue []byte
+			)
+			if update.rawStorageKey {
+				newValue, exists = subset[crypto.Keccak256Hash(key.Bytes())]
+			} else {
+				newValue, exists = subset[key]
+			}
+			if !exists {
+				log.Warn("State update missing storage slot", "address", addr, "key", key)
+				continue
+			}
+			if len(newValue) == 0 {
+				diff.StorageCount -= 1
+				diff.StorageBytes -= common.HashLength
+			}
+			if len(oldValue) == 0 {
+				diff.StorageCount += 1
+				diff.StorageBytes += common.HashLength
+			}
+			diff.StorageBytes += uint64(len(newValue) - len(oldValue))
 		}
 	}
 
-	log.Info("Starting state size initialization")
+	// Calculate trie node changes
-	g.initializeMetrics()
+	for _, subset := range update.nodes.Sets {
-	log.Info("Completed state size initialization", "elapsed", time.Since(start))
+		for path, n := range subset.Nodes {
+			if len(n.Blob) == 0 {
+				diff.TrieNodeCount -= 1
+				diff.TrieNodeBytes -= uint64(len(path) + common.HashLength)
+			}
+			prev, ok := subset.Origins[path]
+			if ok {
+				diff.TrieNodeCount += 1
+				diff.TrieNodeBytes += uint64(len(path) + common.HashLength)
+			}
+			diff.TrieNodeBytes += uint64(len(n.Blob) - len(prev))
+		}
+	}
+
+	// Calculate code changes
+	for _, code := range update.codes {
+		diff.ContractCount += 1
+		diff.ContractBytes += uint64(len(code.blob) + common.HashLength)
+	}
+
+	return diff
+}
+
+// Track is an async method used to send the state update to the generator.
+// It ignores empty updates (where no state changes occurred).
+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
-func (g *stateSizeGenerator) hasExistingMetrics() bool {
+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
@@ -131,7 +329,7 @@ func (g *stateSizeGenerator) hasExistingMetrics() bool {
 }
 
 // initializeMetrics performs the actual metrics initialization
-func (g *stateSizeGenerator) initializeMetrics() {
+func (g *StateSizeGenerator) initializeMetrics() {
 	var (
 		wg                                         sync.WaitGroup
 		accountCount, accountBytes                 uint64
@@ -155,6 +353,7 @@ func (g *stateSizeGenerator) initializeMetrics() {
 			count++
 			bytes += uint64(len(iter.Key()) + len(iter.Value()))
 
+			// Check for abort
 			select {
 			case <-g.abort:
 				return
@@ -183,7 +382,6 @@ func (g *stateSizeGenerator) initializeMetrics() {
 	wg.Wait()
 
 	// Update metrics
-	g.metricsLock.Lock()
 	g.metrics.AccountCount = accountCount
 	g.metrics.AccountBytes = accountBytes
 	g.metrics.StorageCount = storageCount
@@ -192,29 +390,26 @@ func (g *stateSizeGenerator) initializeMetrics() {
 	g.metrics.TrieNodeBytes = trieAccountNodeBytes + trieStorageNodeBytes
 	g.metrics.ContractCount = contractCount
 	g.metrics.ContractBytes = contractBytes
-	g.metricsLock.Unlock()
 
-	// Update metrics in database
+	g.updateMetrics()
 	g.persistMetrics()
+}
 
+func (g *StateSizeGenerator) updateMetrics() {
 	// Update global metrics
-	stateSizeAccountsCountMeter.Mark(int64(accountCount))
+	stateSizeAccountsCountMeter.Mark(int64(g.metrics.AccountCount))
-	stateSizeAccountsBytesMeter.Mark(int64(accountBytes))
+	stateSizeAccountsBytesMeter.Mark(int64(g.metrics.AccountBytes))
-	stateSizeStorageCountMeter.Mark(int64(storageCount))
+	stateSizeStorageCountMeter.Mark(int64(g.metrics.StorageCount))
-	stateSizeStorageBytesMeter.Mark(int64(storageBytes))
+	stateSizeStorageBytesMeter.Mark(int64(g.metrics.StorageBytes))
-	stateSizeTrieNodesCountMeter.Mark(int64(trieAccountNodeCount + trieStorageNodeCount))
+	stateSizeTrieNodesCountMeter.Mark(int64(g.metrics.TrieNodeCount))
-	stateSizeTrieNodesBytesMeter.Mark(int64(trieStorageNodeBytes + trieStorageNodeBytes))
+	stateSizeTrieNodesBytesMeter.Mark(int64(g.metrics.TrieNodeBytes))
-	stateSizeContractsCountMeter.Mark(int64(contractCount))
+	stateSizeContractsCountMeter.Mark(int64(g.metrics.ContractCount))
-	stateSizeContractsBytesMeter.Mark(int64(contractBytes))
+	stateSizeContractsBytesMeter.Mark(int64(g.metrics.ContractBytes))
 }
 
 // persistMetrics saves the current metrics to the database
-func (g *stateSizeGenerator) persistMetrics() {
+func (g *StateSizeGenerator) persistMetrics() {
-	g.metricsLock.RLock()
+	data, err := rlp.EncodeToBytes(*g.metrics)
-	metrics := *g.metrics
-	g.metricsLock.RUnlock()
-
-	data, err := rlp.EncodeToBytes(metrics)
 	if err != nil {
 		log.Error("Failed to encode state size metrics", "err", err)
 		return
@@ -226,93 +421,3 @@ func (g *stateSizeGenerator) persistMetrics() {
 		log.Error("Failed to persist state size metrics", "err", err)
 	}
 }
-
-// updateMetrics updates metrics based on state changes
-func (g *stateSizeGenerator) updateMetrics(update *stateUpdate) {
-	var (
-		accountBytes, storageBytes, nodeBytes, codeBytes int
-		accountCount, storageCount, nodeCount, codeCount int
-	)
-
-	for addr, oldValue := range update.accountsOrigin {
-		addrHash := crypto.Keccak256Hash(addr.Bytes())
-		newValue, exists := update.accounts[addrHash]
-		if !exists {
-			log.Warn("State update missing account", "address", addr)
-			continue
-		}
-		if len(newValue) == 0 {
-			accountCount -= 1
-			accountBytes -= common.HashLength
-		}
-		if len(oldValue) == 0 {
-			accountCount += 1
-			accountBytes += common.HashLength
-		}
-		accountBytes += len(newValue) - len(oldValue)
-	}
-
-	for addr, slots := range update.storagesOrigin {
-		addrHash := crypto.Keccak256Hash(addr.Bytes())
-		subset, exists := update.storages[addrHash]
-		if !exists {
-			log.Warn("State update missing storage", "address", addr)
-			continue
-		}
-		for key, oldValue := range slots {
-			var (
-				exists   bool
-				newValue []byte
-			)
-			if update.rawStorageKey {
-				newValue, exists = subset[crypto.Keccak256Hash(key.Bytes())]
-			} else {
-				newValue, exists = subset[key]
-			}
-			if !exists {
-				log.Warn("State update missing storage slot", "address", addr, "key", key)
-				continue
-			}
-			if len(newValue) == 0 {
-				storageCount -= 1
-				storageBytes -= common.HashLength
-			}
-			if len(oldValue) == 0 {
-				storageCount += 1
-				storageBytes += common.HashLength
-			}
-			storageBytes += len(newValue) - len(oldValue)
-		}
-	}
-	for _, subset := range update.nodes.Sets {
-		for path, n := range subset.Nodes {
-			if len(n.Blob) == 0 {
-				nodeCount -= 1
-				nodeBytes -= len(path) + common.HashLength
-			}
-			prev, ok := subset.Origins[path]
-			if ok {
-				nodeCount += 1
-				nodeBytes += len(path) + common.HashLength
-			}
-			nodeBytes += len(n.Blob) - len(prev)
-		}
-	}
-	for _, code := range update.codes {
-		codeCount += 1
-		codeBytes += len(code.blob) + common.HashLength // no deduplication
-	}
-
-	// Update local metrics
-	g.metricsLock.Lock()
-	g.metrics.Root = update.root
-	g.metrics.AccountCount += uint64(accountCount)
-	g.metrics.AccountBytes += uint64(accountBytes)
-	g.metrics.StorageCount += uint64(storageCount)
-	g.metrics.StorageBytes += uint64(storageBytes)
-	g.metrics.TrieNodeCount += uint64(nodeCount)
-	g.metrics.TrieNodeBytes += uint64(nodeBytes)
-	g.metrics.ContractCount += uint64(codeCount)
-	g.metrics.ContractBytes += uint64(codeBytes)
-	g.metricsLock.Unlock()
-}