Get the right head state when proposing a failed reorg (#13579)

* Get the right head state when proposing a failed reorg * add unit test * split logic
prysmaticlabs · Feb 5, 2024 · e2e7e84 · e2e7e84
1 parent 91b0a93
commit e2e7e84
Show file tree

Hide file tree

Showing 2 changed files with 113 additions and 23 deletions.
diff --git a/beacon-chain/rpc/prysm/v1alpha1/validator/proposer.go b/beacon-chain/rpc/prysm/v1alpha1/validator/proposer.go
@@ -62,40 +62,21 @@ func (vs *Server) GetBeaconBlock(ctx context.Context, req *ethpb.BlockRequest) (
 	if vs.SyncChecker.Syncing() {
 		return nil, status.Error(codes.Unavailable, "Syncing to latest head, not ready to respond")
 	}
-
-	// process attestations and update head in forkchoice
-	vs.ForkchoiceFetcher.UpdateHead(ctx, vs.TimeFetcher.CurrentSlot())
-	headRoot := vs.ForkchoiceFetcher.CachedHeadRoot()
-	parentRoot := vs.ForkchoiceFetcher.GetProposerHead()
-	if parentRoot != headRoot {
-		blockchain.LateBlockAttemptedReorgCount.Inc()
-		log.WithFields(logrus.Fields{
-			"slot":       req.Slot,
-			"parentRoot": fmt.Sprintf("%#x", parentRoot),
-			"headRoot":   fmt.Sprintf("%#x", headRoot),
-		}).Warn("late block attempted reorg failed")
-	}
-
 	// An optimistic validator MUST NOT produce a block (i.e., sign across the DOMAIN_BEACON_PROPOSER domain).
 	if slots.ToEpoch(req.Slot) >= params.BeaconConfig().BellatrixForkEpoch {
 		if err := vs.optimisticStatus(ctx); err != nil {
 			return nil, status.Errorf(codes.Unavailable, "Validator is not ready to propose: %v", err)
 		}
 	}
 
-	sBlk, err := getEmptyBlock(req.Slot)
+	head, parentRoot, err := vs.getParentState(ctx, req.Slot)
 	if err != nil {
-		return nil, status.Errorf(codes.Internal, "Could not prepare block: %v", err)
+		return nil, err
 	}
-	head, err := vs.HeadFetcher.HeadState(ctx)
-	if err != nil {
-		return nil, status.Errorf(codes.Internal, "Could not get head state: %v", err)
-	}
-	head, err = transition.ProcessSlotsUsingNextSlotCache(ctx, head, parentRoot[:], req.Slot)
+	sBlk, err := getEmptyBlock(req.Slot)
 	if err != nil {
-		return nil, status.Errorf(codes.Internal, "Could not process slots up to %d: %v", req.Slot, err)
+		return nil, status.Errorf(codes.Internal, "Could not prepare block: %v", err)
 	}
-
 	// Set slot, graffiti, randao reveal, and parent root.
 	sBlk.SetSlot(req.Slot)
 	sBlk.SetGraffiti(req.Graffiti)
@@ -134,6 +115,67 @@ func (vs *Server) GetBeaconBlock(ctx context.Context, req *ethpb.BlockRequest) (
 	return vs.constructGenericBeaconBlock(sBlk, bundleCache.get(req.Slot))
 }
 
+func (vs *Server) handleFailedReorgAttempt(ctx context.Context, slot primitives.Slot, parentRoot, headRoot [32]byte) (state.BeaconState, error) {
+	blockchain.LateBlockAttemptedReorgCount.Inc()
+	log.WithFields(logrus.Fields{
+		"slot":       slot,
+		"parentRoot": fmt.Sprintf("%#x", parentRoot),
+		"headRoot":   fmt.Sprintf("%#x", headRoot),
+	}).Warn("late block attempted reorg failed")
+	// Try to get the state from the NSC
+	head := transition.NextSlotState(parentRoot[:], slot)
+	if head != nil {
+		return head, nil
+	}
+	// cache miss
+	head, err := vs.StateGen.StateByRoot(ctx, parentRoot)
+	if err != nil {
+		return nil, status.Error(codes.Unavailable, "could not obtain head state")
+	}
+	return head, nil
+}
+
+func (vs *Server) getHeadNoFailedReorg(ctx context.Context, slot primitives.Slot, parentRoot [32]byte) (state.BeaconState, error) {
+	// Try to get the state from the NSC
+	head := transition.NextSlotState(parentRoot[:], slot)
+	if head != nil {
+		return head, nil
+	}
+	head, err := vs.HeadFetcher.HeadState(ctx)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "Could not get head state: %v", err)
+	}
+	return head, nil
+}
+
+func (vs *Server) getParentStateFromReorgData(ctx context.Context, slot primitives.Slot, parentRoot, headRoot [32]byte) (head state.BeaconState, err error) {
+	if parentRoot != headRoot {
+		head, err = vs.handleFailedReorgAttempt(ctx, slot, parentRoot, headRoot)
+	} else {
+		head, err = vs.getHeadNoFailedReorg(ctx, slot, parentRoot)
+	}
+	if err != nil {
+		return nil, err
+	}
+	if head.Slot() >= slot {
+		return head, nil
+	}
+	head, err = transition.ProcessSlotsUsingNextSlotCache(ctx, head, parentRoot[:], slot)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "Could not process slots up to %d: %v", slot, err)
+	}
+	return head, nil
+}
+
+func (vs *Server) getParentState(ctx context.Context, slot primitives.Slot) (state.BeaconState, [32]byte, error) {
+	// process attestations and update head in forkchoice
+	vs.ForkchoiceFetcher.UpdateHead(ctx, vs.TimeFetcher.CurrentSlot())
+	headRoot := vs.ForkchoiceFetcher.CachedHeadRoot()
+	parentRoot := vs.ForkchoiceFetcher.GetProposerHead()
+	head, err := vs.getParentStateFromReorgData(ctx, slot, parentRoot, headRoot)
+	return head, parentRoot, err
+}
+
 func (vs *Server) BuildBlockParallel(ctx context.Context, sBlk interfaces.SignedBeaconBlock, head state.BeaconState, skipMevBoost bool, builderBoostFactor uint64) error {
 	// Build consensus fields in background
 	var wg sync.WaitGroup

diff --git a/beacon-chain/rpc/prysm/v1alpha1/validator/proposer_test.go b/beacon-chain/rpc/prysm/v1alpha1/validator/proposer_test.go
@@ -2859,3 +2859,51 @@ func TestProposer_GetFeeRecipientByPubKey(t *testing.T) {
 
 	require.Equal(t, common.HexToAddress("0x055Fb65722E7b2455012BFEBf6177F1D2e9728D8").Hex(), common.BytesToAddress(resp.FeeRecipient).Hex())
 }
+
+func TestProposer_GetParentHeadState(t *testing.T) {
+	db := dbutil.SetupDB(t)
+	ctx := context.Background()
+
+	parentState, parentRoot, _ := util.DeterministicGenesisStateWithGenesisBlock(t, ctx, db, 100)
+	headState, headRoot, _ := util.DeterministicGenesisStateWithGenesisBlock(t, ctx, db, 50)
+	require.NoError(t, transition.UpdateNextSlotCache(ctx, parentRoot[:], parentState))
+
+	proposerServer := &Server{
+		ChainStartFetcher: &mockExecution.Chain{},
+		Eth1InfoFetcher:   &mockExecution.Chain{},
+		Eth1BlockFetcher:  &mockExecution.Chain{},
+		StateGen:          stategen.New(db, doublylinkedtree.New()),
+	}
+	t.Run("failed reorg", func(tt *testing.T) {
+		head, err := proposerServer.getParentStateFromReorgData(ctx, 1, parentRoot, headRoot)
+		require.NoError(t, err)
+		st := parentState.Copy()
+		st, err = transition.ProcessSlots(ctx, st, st.Slot()+1)
+		require.NoError(t, err)
+		str, err := st.StateRootAtIndex(0)
+		require.NoError(t, err)
+		headStr, err := head.StateRootAtIndex(0)
+		require.NoError(t, err)
+		genesisStr, err := headState.StateRootAtIndex(0)
+		require.NoError(t, err)
+		require.Equal(t, [32]byte(str), [32]byte(headStr))
+		require.NotEqual(t, [32]byte(str), [32]byte(genesisStr))
+	})
+
+	t.Run("no reorg", func(tt *testing.T) {
+		require.NoError(t, transition.UpdateNextSlotCache(ctx, headRoot[:], headState))
+		head, err := proposerServer.getParentStateFromReorgData(ctx, 1, headRoot, headRoot)
+		require.NoError(t, err)
+		st := headState.Copy()
+		st, err = transition.ProcessSlots(ctx, st, st.Slot()+1)
+		require.NoError(t, err)
+		str, err := st.StateRootAtIndex(0)
+		require.NoError(t, err)
+		headStr, err := head.StateRootAtIndex(0)
+		require.NoError(t, err)
+		genesisStr, err := parentState.StateRootAtIndex(0)
+		require.NoError(t, err)
+		require.Equal(t, [32]byte(str), [32]byte(headStr))
+		require.NotEqual(t, [32]byte(str), [32]byte(genesisStr))
+	})
+}