Add Flag for Checking HeadState (#4974)

* gate feature
* imports
* add flag
* Merge branch 'master' into gateFeature

beacon-chain/blockchain/process_attestation_helpers.go

@@ -15,6 +15,7 @@ import (
 	stateTrie "github.com/prysmaticlabs/prysm/beacon-chain/state"
 	"github.com/prysmaticlabs/prysm/shared/attestationutil"
 	"github.com/prysmaticlabs/prysm/shared/bytesutil"
+	"github.com/prysmaticlabs/prysm/shared/featureconfig"
 	"github.com/prysmaticlabs/prysm/shared/params"
 )
 
@@ -29,22 +30,24 @@ func (s *Service) getAttPreState(ctx context.Context, c *ethpb.Checkpoint) (*sta
 	if cachedState != nil {
 		return cachedState, nil
 	}
-	headRoot, err := s.HeadRoot(ctx)
-	if err != nil {
-		return nil, errors.Wrapf(err, "could not get head root")
-	}
-	if bytes.Equal(headRoot, c.Root) {
-		st, err := s.HeadState(ctx)
+	if featureconfig.Get().CheckHeadState {
+		headRoot, err := s.HeadRoot(ctx)
 		if err != nil {
-			return nil, errors.Wrapf(err, "could not get head state")
+			return nil, errors.Wrapf(err, "could not get head root")
 		}
-		if err := s.checkpointState.AddCheckpointState(&cache.CheckpointState{
-			Checkpoint: c,
-			State:      st.Copy(),
-		}); err != nil {
-			return nil, errors.Wrap(err, "could not saved checkpoint state to cache")
+		if bytes.Equal(headRoot, c.Root) {
+			st, err := s.HeadState(ctx)
+			if err != nil {
+				return nil, errors.Wrapf(err, "could not get head state")
+			}
+			if err := s.checkpointState.AddCheckpointState(&cache.CheckpointState{
+				Checkpoint: c,
+				State:      st.Copy(),
+			}); err != nil {
+				return nil, errors.Wrap(err, "could not saved checkpoint state to cache")
+			}
+			return st, nil
 		}
-		return st, nil
 	}
 
 	baseState, err := s.beaconDB.State(ctx, bytesutil.ToBytes32(c.Root))

beacon-chain/blockchain/process_block_helpers.go

@@ -13,6 +13,7 @@ import (
 	"github.com/prysmaticlabs/prysm/beacon-chain/db/filters"
 	stateTrie "github.com/prysmaticlabs/prysm/beacon-chain/state"
 	"github.com/prysmaticlabs/prysm/shared/bytesutil"
+	"github.com/prysmaticlabs/prysm/shared/featureconfig"
 	"github.com/prysmaticlabs/prysm/shared/params"
 	"github.com/prysmaticlabs/prysm/shared/traceutil"
 	"github.com/sirupsen/logrus"
@@ -58,13 +59,16 @@ func (s *Service) getBlockPreState(ctx context.Context, b *ethpb.BeaconBlock) (*
 // verifyBlkPreState validates input block has a valid pre-state.
 func (s *Service) verifyBlkPreState(ctx context.Context, b *ethpb.BeaconBlock) (*stateTrie.BeaconState, error) {
 	preState := s.initSyncState[bytesutil.ToBytes32(b.ParentRoot)]
+	var err error
 	if preState == nil {
-		headRoot, err := s.HeadRoot(ctx)
-		if err != nil {
-			return nil, errors.Wrapf(err, "could not get head root")
-		}
-		if bytes.Equal(headRoot, b.ParentRoot) {
-			return s.HeadState(ctx)
+		if featureconfig.Get().CheckHeadState {
+			headRoot, err := s.HeadRoot(ctx)
+			if err != nil {
+				return nil, errors.Wrapf(err, "could not get head root")
+			}
+			if bytes.Equal(headRoot, b.ParentRoot) {
+				return s.HeadState(ctx)
+			}
 		}
 		preState, err = s.beaconDB.State(ctx, bytesutil.ToBytes32(b.ParentRoot))
 		if err != nil {

shared/featureconfig/config.go

@@ -44,6 +44,7 @@ type Flags struct {
 	EnableByteMempool                          bool   // EnaableByteMempool memory management.
 	EnableDomainDataCache                      bool   // EnableDomainDataCache caches validator calls to DomainData per epoch.
 	EnableStateGenSigVerify                    bool   // EnableStateGenSigVerify verifies proposer and randao signatures during state gen.
+	CheckHeadState                             bool   // CheckHeadState checks the current headstate before retrieving the desired state from the db.
 
 	// DisableForkChoice disables using LMD-GHOST fork choice to update
 	// the head of the chain based on attestations and instead accepts any valid received block
@@ -149,6 +150,10 @@ func ConfigureBeaconChain(ctx *cli.Context) {
 		log.Warn("Enabling sig verify for state gen")
 		cfg.EnableStateGenSigVerify = true
 	}
+	if ctx.GlobalBool(checkHeadState.Name) {
+		log.Warn("Enabling check head state for chainservice")
+		cfg.CheckHeadState = true
+	}
 	Init(cfg)
 }
 

shared/featureconfig/flags.go

@@ -102,6 +102,10 @@ var (
 		Usage: "Enable signature verification for state gen. This feature increases the cost to generate a historical state," +
 			"the resulting state is signature verified.",
 	}
+	checkHeadState = cli.BoolFlag{
+		Name:  "check-head-state",
+		Usage: "Enables the checking of head state in chainservice first before retrieving the desired state from the db.",
+	}
 )
 
 // Deprecated flags list.
@@ -274,6 +278,7 @@ var BeaconChainFlags = append(deprecatedFlags, []cli.Flag{
 	disableUpdateHeadPerAttestation,
 	enableByteMempool,
 	enableStateGenSigVerify,
+	checkHeadState,
 }...)
 
 // E2EBeaconChainFlags contains a list of the beacon chain feature flags to be tested in E2E.
@@ -286,4 +291,5 @@ var E2EBeaconChainFlags = []string{
 	"--proto-array-forkchoice",
 	"--enable-byte-mempool",
 	"--enable-state-gen-sig-verify",
+	"--check-head-state",
 }