Merge pull request #2058 from 0xPolygon/feat/force-pmt-regen-startup

force sequencer to regenerate pmt on startup flag

Author: elliothllm

cmd/utils/flags.go

@@ -965,6 +965,11 @@ var (
 		Usage: "Only use SMT v2 for state changes",
 		Value: true,
 	}
+	ForcePMTInterhashesRegenOnRestart = cli.BoolFlag{
+		Name:  "zkevm.force-pmt-interhashes-regen-on-restart",
+		Usage: "Force regeneration of PMT interhashes on node restart",
+		Value: false,
+	}
 	SequencerBlockGasLimit = cli.Uint64Flag{
 		Name:  "zkevm.sequencer-block-gas-limit",
 		Usage: "The gas limit of the sequencer block.  Default (0) means no limit.",

core/state/intra_block_state_zkevm.go

@@ -2,6 +2,7 @@ package state
 
 import (
 	"errors"
+	"github.com/erigontech/erigon-lib/log/v3"
 
 	"github.com/erigontech/erigon-lib/chain"
 	libcommon "github.com/erigontech/erigon-lib/common"
@@ -135,6 +136,7 @@ func (sdb *IntraBlockState) SyncerPreExecuteStateSet(
 		//save ger with l1blockhash - but only in the case that the l1 info tree index hasn't been
 		// re-used.  If it has been re-used we never write this to the contract storage
 		if !reUsedL1InfoTreeIndex && blockGer != nil && *blockGer != emptyHash {
+			log.Info("[Pre-Execute] [SR-DEBUG] Writing Global Exit Root L1 block hash to state DB", "ger", blockGer.String(), "l1BlockHash", l1BlockHash.String())
 			sdb.WriteGerManagerL1BlockHash(*blockGer, *l1BlockHash)
 		}
 	} else {
@@ -148,6 +150,7 @@ func (sdb *IntraBlockState) SyncerPreExecuteStateSet(
 
 		for _, ger := range *gerUpdates {
 			//save ger
+			log.Info("[Pre-Execute] [SR-DEBUG] Writing Global Exit Root timestamp to state DB", "ger", ger.GlobalExitRoot.String(), "timestamp", ger.Timestamp)
 			sdb.WriteGlobalExitRootTimestamp(ger.GlobalExitRoot, ger.Timestamp)
 		}
 

eth/ethconfig/config_zkevm.go

@@ -134,10 +134,11 @@ type Zk struct {
 	InjectGers      bool
 	HonourChainspec bool `yaml:"zkevm.honour-chainspec"`
 
-	SkipSmt                bool
-	OnlySmtV2              bool
-	SequencerBlockGasLimit uint64
-	PessimisticForkNumber  uint64
+	SkipSmt                           bool
+	OnlySmtV2                         bool
+	ForcePMTInterhashesRegenOnRestart bool
+	SequencerBlockGasLimit            uint64
+	PessimisticForkNumber             uint64
 }
 
 type Hardfork string

eth/stagedsync/stage_interhashes.go

@@ -6,6 +6,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"math/bits"
+	"runtime"
 	"slices"
 	"sync/atomic"
 
@@ -596,6 +597,11 @@ func IncrementIntermediateHashes(logPrefix string, s *StageState, db kv.RwTx, to
 			return trie.EmptyRoot, err
 		}
 	}
+
+	var memStats runtime.MemStats
+	runtime.ReadMemStats(&memStats)
+	logger.Info(fmt.Sprintf("[%s] [SR-DEBUG] PMT Increment intermediate hashes finished promotion", logPrefix), "retainListSize", rl.Len(), "allocMB", memStats.Alloc/1024/1024, "headAllocMB", memStats.HeapAlloc/1024/1024, "heapObjects", memStats.HeapObjects, "numGC", memStats.NumGC)
+
 	accTrieCollector := etl.NewCollector(logPrefix, cfg.tmpDir, etl.NewSortableBuffer(etl.BufferOptimalSize), logger)
 	defer accTrieCollector.Close()
 	accTrieCollectorFunc := accountTrieCollector(accTrieCollector)
@@ -605,6 +611,7 @@ func IncrementIntermediateHashes(logPrefix string, s *StageState, db kv.RwTx, to
 	stTrieCollectorFunc := storageTrieCollector(stTrieCollector)
 
 	loader := trie.NewFlatDBTrieLoader(logPrefix, rl, accTrieCollectorFunc, stTrieCollectorFunc, false)
+
 	hash, err := loader.CalcTrieRoot(db, quit)
 	if err != nil {
 		return trie.EmptyRoot, err

turbo/cli/default_flags.go

@@ -342,6 +342,7 @@ var DefaultFlags = []cli.Flag{
 	&utils.InjectGers,
 	&utils.SkipSmt,
 	&utils.OnlySmtV2,
+	&utils.ForcePMTInterhashesRegenOnRestart,
 	&utils.SequencerBlockGasLimit,
 	&utils.PessimisticForkNumber,
 }

turbo/cli/flags_zkevm.go

@@ -336,6 +336,7 @@ func ApplyFlagsForZkConfig(ctx *cli.Context, cfg *ethconfig.Config) {
 		InjectGers:                             ctx.Bool(utils.InjectGers.Name),
 		SkipSmt:                                ctx.Bool(utils.SkipSmt.Name),
 		OnlySmtV2:                              ctx.Bool(utils.OnlySmtV2.Name),
+		ForcePMTInterhashesRegenOnRestart:      ctx.Bool(utils.ForcePMTInterhashesRegenOnRestart.Name),
 		SequencerBlockGasLimit:                 ctx.Uint64(utils.SequencerBlockGasLimit.Name),
 		PessimisticForkNumber:                  ctx.Uint64(utils.PessimisticForkNumber.Name),
 	}

zk/stages/stage_sequence_execute.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/erigontech/erigon-lib/common"
@@ -32,6 +33,8 @@ type TxYielder interface {
 	BeginYielding()
 }
 
+var regenPmtOnce sync.Once
+
 func SpawnSequencingStage(
 	s *stagedsync.StageState,
 	u stagedsync.Unwinder,
@@ -114,6 +117,19 @@ func sequencingBatchStep(
 	}
 	defer sdb.tx.Rollback()
 
+	if cfg.zk.ForcePMTInterhashesRegenOnRestart {
+		if cfg.zk.UsingPMT() {
+			log.Info(fmt.Sprintf("[%s] [SR-DEBUG] Forcing PMT interhashes regeneration as per configuration", logPrefix))
+			regenPmtOnce.Do(func() {
+				if err = sequencerRegentIntermediateHashesPMT(ctx, s, sdb.tx, cfg); err != nil {
+					panic(fmt.Sprintf("failed to regen PMT interhashes: %v", err))
+				}
+			})
+		} else {
+			log.Warn(fmt.Sprintf("[%s] [SR-DEBUG] Not regenerating PMT as zkevm.force-pmt-interhashes-regen-on-restart is set but PMT is not being used", logPrefix))
+		}
+	}
+
 	if err := cfg.infoTreeUpdater.WarmUp(sdb.tx); err != nil {
 		return err
 	}

zk/stages/stage_sequence_execute_blocks.go

@@ -61,6 +61,7 @@ func handleStateForNewBlockStarting(
 			l1BlockHash := ibs.ReadGerManagerL1BlockHash(l1info.GER)
 			if l1BlockHash == (common.Hash{}) {
 				// not in the contract so let's write it!
+				log.Info(fmt.Sprintf("[%s] Writing GER manager L1 block hash to Intra Block State", batchContext.s.LogPrefix()), "ger", l1info.GER.String(), "l1BlockHash", l1info.ParentHash.String())
 				ibs.WriteGerManagerL1BlockHash(l1info.GER, l1info.ParentHash)
 				if err := hermezDb.WriteLatestUsedGer(blockNumber, l1info.GER); err != nil {
 					return err
@@ -201,15 +202,19 @@ func finaliseBlock(
 	if batchContext.cfg.zk.UsingPMT() {
 		commitment = "pmt"
 		logger := log.New()
+		log.Info(fmt.Sprintf("[%s] [SR-DEBUG] Hashing State for the PMT", batchContext.s.LogPrefix()), "startingBlock", newHeader.Number.Uint64()-1, "endingBlock", newHeader.Number.Uint64())
 		if err = stagedsync.HashStateFromTo(batchContext.s.LogPrefix(), batchContext.sdb.tx, batchContext.cfg.hashStateCfg, newHeader.Number.Uint64()-1, newHeader.Number.Uint64(), batchContext.ctx, logger); err != nil {
 			return nil, err
 		}
 
+		log.Info(fmt.Sprintf("[%s] [SR-DEBUG] IncrementIntermediateHashes for the PMT", batchContext.s.LogPrefix()), "startingBlock", batchContext.s.BlockNumber, "endingBlock", thisBlockNumber)
 		newRoot, err = stagedsync.IncrementIntermediateHashes(batchContext.s.LogPrefix(), batchContext.s, batchContext.sdb.tx, thisBlockNumber, trieConfigSequencer(batchContext.cfg.intersCfg), common.Hash{}, quit, logger)
 	} else {
+		log.Info(fmt.Sprintf("[%s] [SR-DEBUG] IncrementIntermediateHashes for the SMT", batchContext.s.LogPrefix()), "startingBlock", newHeader.Number.Uint64()-1, "endingBlock", newHeader.Number.Uint64())
 		newRoot, err = zkIncrementIntermediateHashes(batchContext.ctx, batchContext.s.LogPrefix(), batchContext.s, batchContext.sdb.tx, batchContext.sdb.eridb, batchContext.sdb.smt, newHeader.Number.Uint64()-1, newHeader.Number.Uint64())
 	}
 	if err != nil {
+		log.Error(fmt.Sprintf("[%s] [SR-DEBUG] IncrementIntermediateHashes failed", batchContext.s.LogPrefix()), "err", err)
 		batchContext.sdb.eridb.RollbackBatch()
 		return nil, err
 	}

zk/stages/stage_sequence_execute_utils.go

@@ -694,3 +694,18 @@ func sequencerRegenIntermediateHashes(ctx context.Context, s *stagedsync.StageSt
 	}
 	return tx.Commit()
 }
+
+// sequencerRegentIntermediateHashesPMT will build the intermediate hashes PMT for the given block number.
+// This will run when flag `zkevm.force-pmt-interhashes-regen-on-restart` is enabled to forcefully regenerate the PMT if there is any issues there.
+func sequencerRegentIntermediateHashesPMT(ctx context.Context, s *stagedsync.StageState, tx kv.RwTx, cfg SequenceBlockCfg) error {
+	timeStart := time.Now()
+	_, err := stagedsync.RegenerateIntermediateHashes(s.LogPrefix(), tx, trieConfigSequencer(cfg.intersCfg), common.Hash{}, ctx, log.New())
+	if err != nil {
+		log.Error(fmt.Sprintf("[%s] [SR-DEBUG] Failed to regenerate intermediate hashes PMT: %v", s.LogPrefix(), err))
+		return err
+	}
+	timeFinish := time.Now()
+	log.Info(fmt.Sprintf("[%s] [SR-DEBUG] Regenerated intermediate hashes PMT in %s", s.LogPrefix(), timeFinish.Sub(timeStart)))
+
+	return nil
+}

zk/utils/global_exit_root.go

@@ -3,6 +3,7 @@ package utils
 import (
 	"encoding/binary"
 	"errors"
+	"github.com/erigontech/erigon-lib/log/v3"
 	"math/big"
 
 	"github.com/erigontech/erigon-lib/common"
@@ -57,6 +58,7 @@ func WriteGlobalExitRoot(stateReader state.StateReader, stateWriter state.Writer
 	}
 
 	// write global exit root to state
+	log.Info("[SR-DEBUG] Writing Global Exit Root to state", "ger", ger.String(), "timestamp", timestamp)
 	if err := stateWriter.WriteAccountStorage(addr, uint64(1), &gerp, emptyUint256, headerTime); err != nil {
 		return err
 	}