[EIP-4399 2/3] Added total difficulty stage and transition logic

cmd/integration/commands/stages.go

@@ -1083,7 +1083,7 @@ func newSync(ctx context.Context, db kv.RwDB, miningConfig *params.MiningConfig)
 		cfg.Miner = *miningConfig
 	}
 
-	sync, err := stages2.NewStagedSync(context.Background(), logger, db, p2p.Config{}, cfg, downloadServer, tmpdir, txPool, txPoolP2PServer, nil)
+	sync, err := stages2.NewStagedSync(context.Background(), logger, db, p2p.Config{}, cfg, *chainConfig, downloadServer, tmpdir, txPool, txPoolP2PServer, nil)
 	if err != nil {
 		panic(err)
 	}

core/rawdb/accessors_chain.go

@@ -976,3 +976,34 @@ func WritePendingEpoch(tx kv.RwTx, blockNum uint64, blockHash common.Hash, trans
 	copy(k[8:], blockHash[:])
 	return tx.Put(kv.PendingEpoch, k, transitionProof)
 }
+
+// Transitioned returns true if the block number comes after POS transition
+func Transitioned(db kv.Getter, blockNum uint64) (trans bool, err error) {
+	data, err := db.GetOne(kv.HeaderTD, dbutils.HeaderKey(math.MaxUint64, common.Hash{}))
+	if err != nil {
+		return false, fmt.Errorf("failed ReadTd: %w", err)
+	}
+	if len(data) == 0 {
+		return false, nil
+	}
+	return blockNum >= binary.BigEndian.Uint64(data), nil
+}
+
+// MarkTreansition sets transition to proof-of-stake from the block number
+func MarkTransition(db kv.StatelessRwTx, blockNum uint64) error {
+	data := make([]byte, 8)
+	binary.BigEndian.PutUint64(data, blockNum)
+	// If we already transitioned then we do not update the transition
+	marked, err := db.Has(kv.HeaderTD, dbutils.HeaderKey(math.MaxUint64, common.Hash{}))
+	if err != nil {
+		return err
+	}
+
+	if marked {
+		return nil
+	}
+	if err := db.Put(kv.HeaderTD, dbutils.HeaderKey(math.MaxUint64, common.Hash{}), data); err != nil {
+		return fmt.Errorf("failed to store block total difficulty: %w", err)
+	}
+	return nil
+}

core/rawdb/accessors_chain_test.go

@@ -416,6 +416,29 @@ func TestBlockReceiptStorage(t *testing.T) {
 	}
 }
 
+// Tests that transitions is handled correctly
+func TestTransition(t *testing.T) {
+	_, tx := memdb.NewTestTx(t)
+	require := require.New(t)
+	transitionBlock := uint64(1000)
+
+	isTrans, err := Transitioned(tx, 1500)
+	require.NoError(err)
+	require.False(isTrans)
+	isTrans, err = Transitioned(tx, 20)
+	require.NoError(err)
+	require.False(isTrans)
+
+	require.NoError(MarkTransition(tx, transitionBlock))
+
+	isTrans, err = Transitioned(tx, 1500)
+	require.NoError(err)
+	require.True(isTrans)
+	isTrans, err = Transitioned(tx, 20)
+	require.NoError(err)
+	require.False(isTrans)
+}
+
 func checkReceiptsRLP(have, want types.Receipts) error {
 	if len(have) != len(want) {
 		return fmt.Errorf("receipts sizes mismatch: have %d, want %d", len(have), len(want))

eth/backend.go

@@ -463,7 +463,7 @@ func New(stack *node.Node, config *ethconfig.Config, logger log.Logger) (*Ethere
 		return nil, err
 	}
 
-	backend.stagedSync, err = stages2.NewStagedSync(backend.downloadCtx, backend.logger, backend.chainDB, stack.Config().P2P, *config, backend.downloadServer, tmpdir, backend.txPool, backend.txPoolP2PServer, backend.notifications.Accumulator)
+	backend.stagedSync, err = stages2.NewStagedSync(backend.downloadCtx, backend.logger, backend.chainDB, stack.Config().P2P, *config, *chainConfig, backend.downloadServer, tmpdir, backend.txPool, backend.txPoolP2PServer, backend.notifications.Accumulator)
 	if err != nil {
 		return nil, err
 	}

eth/stagedsync/default_stages.go

@@ -13,6 +13,7 @@ func DefaultStages(ctx context.Context,
 	headers HeadersCfg,
 	blockHashCfg BlockHashesCfg,
 	bodies BodiesCfg,
+	difficulty DifficultyCfg,
 	senders SendersCfg,
 	exec ExecuteBlockCfg,
 	trans TranspileCfg,
@@ -69,6 +70,19 @@ func DefaultStages(ctx context.Context,
 				return PruneBodiesStage(p, tx, bodies, ctx)
 			},
 		},
+		{
+			ID:          stages.TotalDifficulty,
+			Description: "Compute total difficulty",
+			Forward: func(firstCycle bool, badBlockUnwind bool, s *StageState, u Unwinder, tx kv.RwTx) error {
+				return SpawnDifficultyStage(s, tx, difficulty, ctx)
+			},
+			Unwind: func(firstCycle bool, u *UnwindState, s *StageState, tx kv.RwTx) error {
+				return UnwindDifficultyStage(u, tx, ctx)
+			},
+			Prune: func(firstCycle bool, p *PruneState, tx kv.RwTx) error {
+				return PruneDifficultyStage(p, tx, ctx)
+			},
+		},
 		{
 			ID:          stages.Senders,
 			Description: "Recover senders from tx signatures",

eth/stagedsync/stage_difficulty.go

@@ -0,0 +1,150 @@
+package stagedsync
+
+import (
+	"bytes"
+	"context"
+	"encoding/binary"
+	"fmt"
+	"math/big"
+
+	"github.com/ledgerwatch/erigon-lib/etl"
+	"github.com/ledgerwatch/erigon-lib/kv"
+	"github.com/ledgerwatch/erigon/core/rawdb"
+	"github.com/ledgerwatch/erigon/core/types"
+	"github.com/ledgerwatch/erigon/eth/stagedsync/stages"
+	"github.com/ledgerwatch/erigon/rlp"
+)
+
+type DifficultyCfg struct {
+	tmpDir                  string
+	terminalTotalDifficulty *big.Int
+	db                      kv.RwDB
+}
+
+func StageDifficultyCfg(db kv.RwDB, tmpDir string, terminalTotalDifficulty *big.Int) DifficultyCfg {
+	return DifficultyCfg{
+		db:                      db,
+		tmpDir:                  tmpDir,
+		terminalTotalDifficulty: terminalTotalDifficulty,
+	}
+}
+
+func SpawnDifficultyStage(s *StageState, tx kv.RwTx, cfg DifficultyCfg, ctx context.Context) (err error) {
+	useExternalTx := tx != nil
+
+	if !useExternalTx {
+		var err error
+		tx, err = cfg.db.BeginRw(context.Background())
+		if err != nil {
+			return err
+		}
+		defer tx.Rollback()
+	}
+
+	quit := ctx.Done()
+	headNumber, err := stages.GetStageProgress(tx, stages.Headers)
+	if err != nil {
+		return fmt.Errorf("getting headers progress: %w", err)
+	}
+
+	td := big.NewInt(0)
+	if s.BlockNumber > 0 {
+		td, err = rawdb.ReadTd(tx, rawdb.ReadHeaderByNumber(tx, s.BlockNumber).Hash(), s.BlockNumber)
+		if err != nil {
+			return err
+		}
+	}
+	// If the chain does not have a proof of stake config or has reached terminalTotalDifficulty then we can skip this stage
+	if cfg.terminalTotalDifficulty == nil || td.Cmp(cfg.terminalTotalDifficulty) >= 0 {
+		if err = s.Update(tx, headNumber); err != nil {
+			return err
+		}
+		if !useExternalTx {
+			if err = tx.Commit(); err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+
+	startKey := make([]byte, 8)
+	binary.BigEndian.PutUint64(startKey, s.BlockNumber)
+
+	header := new(types.Header)
+	if err := etl.Transform(
+		s.LogPrefix(),
+		tx,
+		kv.Headers,
+		kv.HeaderTD,
+		cfg.tmpDir,
+		func(k []byte, v []byte, next etl.ExtractNextFunc) error {
+			if len(k) != 40 {
+				return nil
+			}
+
+			blockNum := binary.BigEndian.Uint64(k)
+			canonical, err := rawdb.ReadCanonicalHash(tx, blockNum)
+			if err != nil {
+				return err
+			}
+
+			if !bytes.Equal(k[8:], canonical[:]) {
+				return nil
+			}
+			if err := rlp.Decode(bytes.NewReader(v), header); err != nil {
+				return err
+			}
+
+			td.Add(td, header.Difficulty)
+			if td.Cmp(cfg.terminalTotalDifficulty) > 0 {
+				return rawdb.MarkTransition(tx, blockNum)
+			}
+			data, err := rlp.EncodeToBytes(td)
+			if err != nil {
+				return fmt.Errorf("failed to RLP encode block total difficulty: %w", err)
+			}
+			return next(k, k, data)
+		},
+		etl.IdentityLoadFunc,
+		etl.TransformArgs{
+			ExtractStartKey: startKey,
+			Quit:            quit,
+		},
+	); err != nil {
+		return err
+	}
+	if err = s.Update(tx, headNumber); err != nil {
+		return err
+	}
+	if !useExternalTx {
+		if err = tx.Commit(); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func UnwindDifficultyStage(u *UnwindState, tx kv.RwTx, ctx context.Context) (err error) {
+	useExternalTx := tx != nil
+
+	if err = u.Done(tx); err != nil {
+		return fmt.Errorf(" reset: %w", err)
+	}
+	if !useExternalTx {
+		if err = tx.Commit(); err != nil {
+			return fmt.Errorf("failed to write db commit: %w", err)
+		}
+	}
+	return nil
+}
+
+func PruneDifficultyStage(p *PruneState, tx kv.RwTx, ctx context.Context) (err error) {
+	useExternalTx := tx != nil
+
+	if !useExternalTx {
+		if err = tx.Commit(); err != nil {
+			return err
+		}
+	}
+	return nil
+}