Feature flag: update head timely (#8412)

* Feature flag: update head timely

* Move finalized imply justified up

* Fix resolve error

* Use invert feature flag

* Make diff easier to review

* Line

* Typo

* Fix condition

Co-authored-by: Raul Jordan <raul@prysmaticlabs.com>

beacon-chain/blockchain/process_block.go

@@ -125,9 +125,24 @@ func (s *Service) onBlock(ctx context.Context, signed *ethpb.SignedBeaconBlock,
 			return err
 		}
 	}
+	var newFinalized bool
+	if featureconfig.Get().UpdateHeadTimely {
+		if postState.FinalizedCheckpointEpoch() > s.finalizedCheckpt.Epoch {
+			if err := s.finalizedImpliesNewJustified(ctx, postState); err != nil {
+				return errors.Wrap(err, "could not save new justified")
+			}
+			s.prevFinalizedCheckpt = s.finalizedCheckpt
+			s.finalizedCheckpt = postState.FinalizedCheckpoint()
+			newFinalized = true
+		}
+
+		if err := s.updateHead(ctx, s.getJustifiedBalances()); err != nil {
+			log.WithError(err).Warn("Could not update head")
+		}
+	}
 
 	// Update finalized check point.
-	if postState.FinalizedCheckpointEpoch() > s.finalizedCheckpt.Epoch {
+	if (postState.FinalizedCheckpointEpoch() > s.finalizedCheckpt.Epoch) || (featureconfig.Get().UpdateHeadTimely && newFinalized) {
 		if err := s.beaconDB.SaveBlocks(ctx, s.getInitSyncBlocks()); err != nil {
 			return err
 		}
@@ -141,9 +156,10 @@ func (s *Service) onBlock(ctx context.Context, signed *ethpb.SignedBeaconBlock,
 		if err := s.forkChoiceStore.Prune(ctx, fRoot); err != nil {
 			return errors.Wrap(err, "could not prune proto array fork choice nodes")
 		}
-
-		if err := s.finalizedImpliesNewJustified(ctx, postState); err != nil {
-			return errors.Wrap(err, "could not save new justified")
+		if !featureconfig.Get().UpdateHeadTimely {
+			if err := s.finalizedImpliesNewJustified(ctx, postState); err != nil {
+				return errors.Wrap(err, "could not save new justified")
+			}
 		}
 
 		// Update deposit cache.

beacon-chain/blockchain/process_block_helpers.go

@@ -11,6 +11,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"
 	"github.com/prysmaticlabs/prysm/shared/traceutil"
 	"go.opencensus.io/trace"
@@ -230,9 +231,10 @@ func (s *Service) updateFinalized(ctx context.Context, cp *ethpb.Checkpoint) err
 	if err := s.beaconDB.SaveFinalizedCheckpoint(ctx, cp); err != nil {
 		return err
 	}
-
-	s.prevFinalizedCheckpt = s.finalizedCheckpt
-	s.finalizedCheckpt = cp
+	if !featureconfig.Get().UpdateHeadTimely {
+		s.prevFinalizedCheckpt = s.finalizedCheckpt
+		s.finalizedCheckpt = cp
+	}
 
 	fRoot := bytesutil.ToBytes32(cp.Root)
 	if err := s.stateGen.MigrateToCold(ctx, fRoot); err != nil {

beacon-chain/blockchain/receive_block.go

@@ -10,6 +10,7 @@ import (
 	statefeed "github.com/prysmaticlabs/prysm/beacon-chain/core/feed/state"
 	"github.com/prysmaticlabs/prysm/beacon-chain/core/helpers"
 	stateTrie "github.com/prysmaticlabs/prysm/beacon-chain/state"
+	"github.com/prysmaticlabs/prysm/shared/featureconfig"
 	"github.com/prysmaticlabs/prysm/shared/traceutil"
 	"github.com/sirupsen/logrus"
 	"go.opencensus.io/trace"
@@ -44,8 +45,10 @@ func (s *Service) ReceiveBlock(ctx context.Context, block *ethpb.SignedBeaconBlo
 	}
 
 	// Update and save head block after fork choice.
-	if err := s.updateHead(ctx, s.getJustifiedBalances()); err != nil {
-		log.WithError(err).Warn("Could not update head")
+	if !featureconfig.Get().UpdateHeadTimely {
+		if err := s.updateHead(ctx, s.getJustifiedBalances()); err != nil {
+			log.WithError(err).Warn("Could not update head")
+		}
 	}
 
 	// Send notification of the processed block to the state feed.

shared/featureconfig/config.go

@@ -47,6 +47,7 @@ type Flags struct {
 	EnableLargerGossipHistory          bool // EnableLargerGossipHistory increases the gossip history we store in our caches.
 	WriteWalletPasswordOnWebOnboarding bool // WriteWalletPasswordOnWebOnboarding writes the password to disk after Prysm web signup.
 	DisableAttestingHistoryDBCache     bool // DisableAttestingHistoryDBCache for the validator client increases disk reads/writes.
+	UpdateHeadTimely                   bool // UpdateHeadTimely updates head right after state transition.
 
 	// Logging related toggles.
 	DisableGRPCConnectionLogs bool // Disables logging when a new grpc client has connected.
@@ -191,6 +192,10 @@ func ConfigureBeaconChain(ctx *cli.Context) {
 		log.Warn("Enabling next slot state cache")
 		cfg.EnableNextSlotStateCache = true
 	}
+	if ctx.Bool(updateHeadTimely.Name) {
+		log.Warn("Enabling update head timely feature")
+		cfg.UpdateHeadTimely = true
+	}
 	Init(cfg)
 }
 

shared/featureconfig/flags.go

@@ -114,13 +114,18 @@ var (
 		Name:  "enable-next-slot-state-cache",
 		Usage: "Improves attesting and proposing efficiency by caching the next slot state at the end of the current slot",
 	}
+	updateHeadTimely = &cli.BoolFlag{
+		Name:  "update-head-timely",
+		Usage: "Improves update head time by updating head right after state transition",
+	}
 )
 
 // devModeFlags holds list of flags that are set when development mode is on.
 var devModeFlags = []cli.Flag{
 	enableLargerGossipHistory,
 	enableNextSlotStateCache,
 	forceOptMaxCoverAggregationStategy,
+	updateHeadTimely,
 }
 
 // ValidatorFlags contains a list of all the feature flags that apply to the validator client.
@@ -167,6 +172,7 @@ var BeaconChainFlags = append(deprecatedFlags, []cli.Flag{
 	disableBroadcastSlashingFlag,
 	enableNextSlotStateCache,
 	forceOptMaxCoverAggregationStategy,
+	updateHeadTimely,
 }...)
 
 // E2EBeaconChainFlags contains a list of the beacon chain feature flags to be tested in E2E.