Utilize Read Only Head State (#11932)

beacon-chain/blockchain/chain_info.go

@@ -54,6 +54,7 @@ type HeadFetcher interface {
 	HeadRoot(ctx context.Context) ([]byte, error)
 	HeadBlock(ctx context.Context) (interfaces.SignedBeaconBlock, error)
 	HeadState(ctx context.Context) (state.BeaconState, error)
+	HeadStateReadOnly(ctx context.Context) (state.ReadOnlyBeaconState, error)
 	HeadValidatorsIndices(ctx context.Context, epoch primitives.Epoch) ([]primitives.ValidatorIndex, error)
 	HeadGenesisValidatorsRoot() [32]byte
 	HeadETH1Data() *ethpb.Eth1Data
@@ -187,6 +188,28 @@ func (s *Service) HeadState(ctx context.Context) (state.BeaconState, error) {
 	return s.cfg.StateGen.StateByRoot(ctx, s.headRoot())
 }
 
+// HeadStateReadOnly returns the read only head state of the chain.
+// If the head is nil from service struct, it will attempt to get the
+// head state from DB. Any callers of this method MUST only use the
+// state instance to read fields from the state. Any type assertions back
+// to the concrete type and subsequent use of it could lead to corruption
+// of the state.
+func (s *Service) HeadStateReadOnly(ctx context.Context) (state.ReadOnlyBeaconState, error) {
+	ctx, span := trace.StartSpan(ctx, "blockChain.HeadStateReadOnly")
+	defer span.End()
+	s.headLock.RLock()
+	defer s.headLock.RUnlock()
+
+	ok := s.hasHeadState()
+	span.AddAttributes(trace.BoolAttribute("cache_hit", ok))
+
+	if ok {
+		return s.headStateReadOnly(ctx), nil
+	}
+
+	return s.cfg.StateGen.StateByRoot(ctx, s.headRoot())
+}
+
 // HeadValidatorsIndices returns a list of active validator indices from the head view of a given epoch.
 func (s *Service) HeadValidatorsIndices(ctx context.Context, epoch primitives.Epoch) ([]primitives.ValidatorIndex, error) {
 	s.headLock.RLock()

beacon-chain/blockchain/head.go

@@ -276,6 +276,16 @@ func (s *Service) headState(ctx context.Context) state.BeaconState {
 	return s.head.state.Copy()
 }
 
+// This returns a read only version of the head state.
+// It does not perform a copy of the head state.
+// This is a lock free version.
+func (s *Service) headStateReadOnly(ctx context.Context) state.ReadOnlyBeaconState {
+	ctx, span := trace.StartSpan(ctx, "blockChain.headStateReadOnly")
+	defer span.End()
+
+	return s.head.state
+}
+
 // This returns the genesis validators root of the head state.
 // This is a lock free version.
 func (s *Service) headGenesisValidatorsRoot() [32]byte {

beacon-chain/blockchain/head_test.go

@@ -232,6 +232,48 @@ func Test_notifyNewHeadEvent(t *testing.T) {
 	})
 }
 
+func TestRetrieveHead_ReadOnly(t *testing.T) {
+	ctx := context.Background()
+	beaconDB := testDB.SetupDB(t)
+	service := setupBeaconChain(t, beaconDB)
+
+	oldBlock := util.SaveBlock(t, context.Background(), service.cfg.BeaconDB, util.NewBeaconBlock())
+	oldRoot, err := oldBlock.Block().HashTreeRoot()
+	require.NoError(t, err)
+	service.head = &head{
+		root:  oldRoot,
+		block: oldBlock,
+	}
+
+	newHeadSignedBlock := util.NewBeaconBlock()
+	newHeadSignedBlock.Block.Slot = 1
+	newHeadBlock := newHeadSignedBlock.Block
+	ojc := &ethpb.Checkpoint{Root: params.BeaconConfig().ZeroHash[:]}
+	ofc := &ethpb.Checkpoint{Root: params.BeaconConfig().ZeroHash[:]}
+
+	wsb := util.SaveBlock(t, context.Background(), service.cfg.BeaconDB, newHeadSignedBlock)
+	newRoot, err := newHeadBlock.HashTreeRoot()
+	require.NoError(t, err)
+	state, blkRoot, err := prepareForkchoiceState(ctx, wsb.Block().Slot()-1, wsb.Block().ParentRoot(), service.cfg.ForkChoiceStore.CachedHeadRoot(), [32]byte{}, ojc, ofc)
+	require.NoError(t, err)
+	require.NoError(t, service.cfg.ForkChoiceStore.InsertNode(ctx, state, blkRoot))
+
+	state, blkRoot, err = prepareForkchoiceState(ctx, wsb.Block().Slot(), newRoot, wsb.Block().ParentRoot(), [32]byte{}, ojc, ofc)
+	require.NoError(t, err)
+	require.NoError(t, service.cfg.ForkChoiceStore.InsertNode(ctx, state, blkRoot))
+	headState, err := util.NewBeaconState()
+	require.NoError(t, err)
+	require.NoError(t, headState.SetSlot(1))
+	require.NoError(t, service.cfg.BeaconDB.SaveStateSummary(context.Background(), &ethpb.StateSummary{Slot: 1, Root: newRoot[:]}))
+	require.NoError(t, service.cfg.BeaconDB.SaveState(context.Background(), headState, newRoot))
+	require.NoError(t, service.saveHead(context.Background(), newRoot, wsb, headState))
+
+	rOnlyState, err := service.HeadStateReadOnly(ctx)
+	require.NoError(t, err)
+
+	assert.Equal(t, rOnlyState, service.head.state, "Head is not the same object")
+}
+
 func TestSaveOrphanedAtts(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)

beacon-chain/blockchain/testing/mock.go

@@ -284,6 +284,11 @@ func (s *ChainService) HeadState(context.Context) (state.BeaconState, error) {
 	return s.State, nil
 }
 
+// HeadStateReadOnly mocks HeadStateReadOnly method in chain service.
+func (s *ChainService) HeadStateReadOnly(context.Context) (state.ReadOnlyBeaconState, error) {
+	return s.State, nil
+}
+
 // CurrentFork mocks HeadState method in chain service.
 func (s *ChainService) CurrentFork() *ethpb.Fork {
 	return s.Fork

beacon-chain/sync/validate_bls_to_execution_change.go

@@ -41,7 +41,7 @@ func (s *Service) validateBlsToExecutionChange(ctx context.Context, pid peer.ID,
 	if s.cfg.blsToExecPool.ValidatorExists(blsChange.Message.ValidatorIndex) {
 		return pubsub.ValidationIgnore, nil
 	}
-	st, err := s.cfg.chain.HeadState(ctx)
+	st, err := s.cfg.chain.HeadStateReadOnly(ctx)
 	if err != nil {
 		return pubsub.ValidationIgnore, err
 	}