From 6ef1a712c27ed4a4f7befabeb80f3400714e5809 Mon Sep 17 00:00:00 2001
From: terence tsao <terence@prysmaticlabs.com>
Date: Tue, 14 Jan 2020 08:05:22 -0800
Subject: [PATCH] OnBlockCacheFilteredTree (#4541)

---
 .../blockchain/forkchoice/process_block.go    | 30 ++++++++++++-------
 beacon-chain/blockchain/forkchoice/service.go |  1 +
 beacon-chain/blockchain/receive_block.go      |  2 +-
 beacon-chain/blockchain/service_test.go       |  4 +++
 beacon-chain/rpc/beacon/validators.go         |  4 +--
 5 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/beacon-chain/blockchain/forkchoice/process_block.go b/beacon-chain/blockchain/forkchoice/process_block.go
index ca3b758d97f..a626d36aabe 100644
--- a/beacon-chain/blockchain/forkchoice/process_block.go
+++ b/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
diff --git a/beacon-chain/blockchain/forkchoice/service.go b/beacon-chain/blockchain/forkchoice/service.go
index 1b13e871ad8..d56a935ec49 100644
--- a/beacon-chain/blockchain/forkchoice/service.go
+++ b/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
diff --git a/beacon-chain/blockchain/receive_block.go b/beacon-chain/blockchain/receive_block.go
index b094ea7b75b..c4d0c451b8d 100644
--- a/beacon-chain/blockchain/receive_block.go
+++ b/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
diff --git a/beacon-chain/blockchain/service_test.go b/beacon-chain/blockchain/service_test.go
index 638b51329d1..c82bb05dbbd 100644
--- a/beacon-chain/blockchain/service_test.go
+++ b/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
 }
diff --git a/beacon-chain/rpc/beacon/validators.go b/beacon-chain/rpc/beacon/validators.go
index dd87cb31d30..a18cdc167ad 100644
--- a/beacon-chain/rpc/beacon/validators.go
+++ b/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.