OnBlockCacheFilteredTree (#4541)

beacon-chain/blockchain/forkchoice/process_block.go

@@ -93,16 +93,6 @@ func (s *Store) OnBlock(ctx context.Context, signed *ethpb.SignedBeaconBlock) er
 		return errors.Wrap(err, "could not save state")
 	}
 
-	if featureconfig.Get().EnableBlockTreeCache {
-		tree, err := s.getFilterBlockTree(ctx)
-		if err != nil {
-			return errors.Wrap(err, "could not calculate filtered block tree")
-		}
-		s.filteredBlockTreeLock.Lock()
-		s.filteredBlockTree = tree
-		s.filteredBlockTreeLock.Unlock()
-	}
-
 	// Update justified check point.
 	if postState.CurrentJustifiedCheckpoint.Epoch > s.justifiedCheckpt.Epoch {
 		if err := s.updateJustified(ctx, postState); err != nil {
@@ -155,6 +145,26 @@ func (s *Store) OnBlock(ctx context.Context, signed *ethpb.SignedBeaconBlock) er
 	return nil
 }
 
+// OnBlockCacheFilteredTree calls OnBlock with additional of caching of filtered block tree
+// for efficient fork choice processing.
+func (s *Store) OnBlockCacheFilteredTree(ctx context.Context, signed *ethpb.SignedBeaconBlock) error {
+	if err := s.OnBlock(ctx, signed); err != nil {
+		return err
+	}
+
+	if featureconfig.Get().EnableBlockTreeCache {
+		tree, err := s.getFilterBlockTree(ctx)
+		if err != nil {
+			return errors.Wrap(err, "could not calculate filtered block tree")
+		}
+		s.filteredBlockTreeLock.Lock()
+		s.filteredBlockTree = tree
+		s.filteredBlockTreeLock.Unlock()
+	}
+
+	return nil
+}
+
 // OnBlockInitialSyncStateTransition is called when an initial sync block is received.
 // It runs state transition on the block and without any BLS verification. The BLS verification
 // includes proposer signature, randao and attestation's aggregated signature. It also does not save

beacon-chain/blockchain/forkchoice/service.go

@@ -28,6 +28,7 @@ import (
 type ForkChoicer interface {
 	Head(ctx context.Context) ([]byte, error)
 	OnBlock(ctx context.Context, b *ethpb.SignedBeaconBlock) error
+	OnBlockCacheFilteredTree(ctx context.Context, b *ethpb.SignedBeaconBlock) error
 	OnBlockInitialSyncStateTransition(ctx context.Context, b *ethpb.SignedBeaconBlock) error
 	OnAttestation(ctx context.Context, a *ethpb.Attestation) error
 	GenesisStore(ctx context.Context, justifiedCheckpoint *ethpb.Checkpoint, finalizedCheckpoint *ethpb.Checkpoint) error

beacon-chain/blockchain/receive_block.go

@@ -70,7 +70,7 @@ func (s *Service) ReceiveBlockNoPubsub(ctx context.Context, block *ethpb.SignedB
 	blockCopy := proto.Clone(block).(*ethpb.SignedBeaconBlock)
 
 	// Apply state transition on the new block.
-	if err := s.forkChoiceStore.OnBlock(ctx, blockCopy); err != nil {
+	if err := s.forkChoiceStore.OnBlockCacheFilteredTree(ctx, blockCopy); err != nil {
 		err := errors.Wrap(err, "could not process block from fork choice service")
 		traceutil.AnnotateError(span, err)
 		return err

beacon-chain/blockchain/service_test.go

@@ -46,6 +46,10 @@ func (s *store) OnBlock(ctx context.Context, b *ethpb.SignedBeaconBlock) error {
 	return nil
 }
 
+func (s *store) OnBlockCacheFilteredTree(ctx context.Context, b *ethpb.SignedBeaconBlock) error {
+	return nil
+}
+
 func (s *store) OnBlockInitialSyncStateTransition(ctx context.Context, b *ethpb.SignedBeaconBlock) error {
 	return nil
 }

beacon-chain/rpc/beacon/validators.go

@@ -8,13 +8,13 @@ import (
 
 	ptypes "github.com/gogo/protobuf/types"
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/status"
 	"github.com/prysmaticlabs/prysm/beacon-chain/core/helpers"
 	"github.com/prysmaticlabs/prysm/beacon-chain/core/state"
 	"github.com/prysmaticlabs/prysm/beacon-chain/core/validators"
 	"github.com/prysmaticlabs/prysm/shared/pagination"
 	"github.com/prysmaticlabs/prysm/shared/params"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 // ListValidatorBalances retrieves the validator balances for a given set of public keys.