Handle nil head block in cache (#4888)

* Nil check

* Fixed tests

* Covered rest of the codebase

* Race tests

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

beacon-chain/blockchain/chain_info.go

@@ -32,7 +32,7 @@ type TimeFetcher interface {
 type HeadFetcher interface {
 	HeadSlot() uint64
 	HeadRoot(ctx context.Context) ([]byte, error)
-	HeadBlock() *ethpb.SignedBeaconBlock
+	HeadBlock(ctx context.Context) (*ethpb.SignedBeaconBlock, error)
 	HeadState(ctx context.Context) (*state.BeaconState, error)
 	HeadValidatorsIndices(epoch uint64) ([]uint64, error)
 	HeadSeed(epoch uint64) ([32]byte, error)
@@ -135,23 +135,25 @@ func (s *Service) HeadRoot(ctx context.Context) ([]byte, error) {
 }
 
 // HeadBlock returns the head block of the chain.
-func (s *Service) HeadBlock() *ethpb.SignedBeaconBlock {
-	return s.headBlock()
+// If the head state is nil from service struct,
+// it will attempt to get the head block from DB.
+func (s *Service) HeadBlock(ctx context.Context) (*ethpb.SignedBeaconBlock, error) {
+	if s.hasHeadState() {
+		return s.headBlock(), nil
+	}
+
+	return s.beaconDB.HeadBlock(ctx)
 }
 
 // HeadState returns the head state of the chain.
 // If the head state is nil from service struct,
-// it will attempt to get from DB and error if nil again.
+// it will attempt to get the head state from DB.
 func (s *Service) HeadState(ctx context.Context) (*state.BeaconState, error) {
 	if s.hasHeadState() {
 		return s.headState(), nil
 	}
 
-	headState, err := s.beaconDB.HeadState(ctx)
-	if err != nil {
-		return nil, err
-	}
-	return headState, nil
+	return s.beaconDB.HeadState(ctx)
 }
 
 // HeadValidatorsIndices returns a list of active validator indices from the head view of a given epoch.

beacon-chain/blockchain/chain_info_norace_test.go

@@ -54,7 +54,7 @@ func TestHeadBlock_DataRace(t *testing.T) {
 			[32]byte{},
 		)
 	}()
-	s.HeadBlock()
+	s.HeadBlock(context.Background())
 }
 
 func TestHeadState_DataRace(t *testing.T) {

beacon-chain/blockchain/chain_info_test.go

@@ -145,10 +145,16 @@ func TestHeadRoot_CanRetrieve(t *testing.T) {
 
 func TestHeadBlock_CanRetrieve(t *testing.T) {
 	b := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{Slot: 1}}
+	s, _ := state.InitializeFromProto(&pb.BeaconState{})
 	c := &Service{}
-	c.head = &head{block: b}
+	c.head = &head{block: b, state: s}
 
-	if !reflect.DeepEqual(b, c.HeadBlock()) {
+	recevied, err := c.HeadBlock(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !reflect.DeepEqual(b, recevied) {
 		t.Error("incorrect head block received")
 	}
 }

beacon-chain/blockchain/service_test.go

@@ -305,7 +305,11 @@ func TestChainService_InitializeBeaconChain(t *testing.T) {
 	if _, err := bc.HeadState(ctx); err != nil {
 		t.Error(err)
 	}
-	if bc.HeadBlock() == nil {
+	headBlk, err := bc.HeadBlock(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if headBlk == nil {
 		t.Error("Head state can't be nil after initialize beacon chain")
 	}
 	if bc.headRoot() == params.BeaconConfig().ZeroHash {
@@ -356,7 +360,11 @@ func TestChainService_InitializeChainInfo(t *testing.T) {
 	if err := c.initializeChainInfo(ctx); err != nil {
 		t.Fatal(err)
 	}
-	if !reflect.DeepEqual(c.HeadBlock(), headBlock) {
+	headBlk, err := c.HeadBlock(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !reflect.DeepEqual(headBlk, headBlock) {
 		t.Error("head block incorrect")
 	}
 	s, err := c.HeadState(ctx)

beacon-chain/blockchain/testing/mock.go

@@ -160,8 +160,8 @@ func (ms *ChainService) HeadRoot(ctx context.Context) ([]byte, error) {
 }
 
 // HeadBlock mocks HeadBlock method in chain service.
-func (ms *ChainService) HeadBlock() *ethpb.SignedBeaconBlock {
-	return ms.Block
+func (ms *ChainService) HeadBlock(context.Context) (*ethpb.SignedBeaconBlock, error) {
+	return ms.Block, nil
 }
 
 // HeadState mocks HeadState method in chain service.

beacon-chain/rpc/beacon/blocks.go

@@ -237,7 +237,10 @@ func (bs *Server) StreamChainHead(_ *ptypes.Empty, stream ethpb.BeaconChain_Stre
 
 // Retrieve chain head information from the DB and the current beacon state.
 func (bs *Server) chainHeadRetrieval(ctx context.Context) (*ethpb.ChainHead, error) {
-	headBlock := bs.HeadFetcher.HeadBlock()
+	headBlock, err := bs.HeadFetcher.HeadBlock(ctx)
+	if err != nil {
+		return nil, status.Error(codes.Internal, "Could not get head block")
+	}
 	if headBlock == nil {
 		return nil, status.Error(codes.Internal, "Head block of chain was nil")
 	}

beacon-chain/rpc/validator/server.go

@@ -158,7 +158,11 @@ func (vs *Server) DomainData(ctx context.Context, request *ethpb.DomainRequest)
 // CanonicalHead of the current beacon chain. This method is requested on-demand
 // by a validator when it is their time to propose or attest.
 func (vs *Server) CanonicalHead(ctx context.Context, req *ptypes.Empty) (*ethpb.SignedBeaconBlock, error) {
-	return vs.HeadFetcher.HeadBlock(), nil
+	headBlk, err := vs.HeadFetcher.HeadBlock(ctx)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "Could not get head block: %v", err)
+	}
+	return headBlk, nil
 }
 
 // WaitForChainStart queries the logs of the Deposit Contract in order to verify the beacon chain