Move IsCanonical implementation to forkchoice package

beacon-chain/blockchain/chain_info.go

@@ -255,9 +255,6 @@ func (s *Service) IsCanonical(ctx context.Context, blockRoot [32]byte) (bool, er
 		return true, nil
 	}
 
-	// If the block has not been finalized, the block must be recent. Check recent canonical roots
-	// mapping which uses proto array fork choice.
-	s.recentCanonicalBlocksLock.RLock()
-	defer s.recentCanonicalBlocksLock.RUnlock()
-	return s.recentCanonicalBlocks[blockRoot], nil
+	// If the block has not been finalized, check fork choice store to see if the block is canonical
+	return s.forkChoiceStore.IsCanonical(blockRoot), nil
 }

beacon-chain/blockchain/chain_info_test.go

@@ -181,3 +181,23 @@ func TestHeadETH1Data_CanRetrieve(t *testing.T) {
 		t.Error("Received incorrect eth1 data")
 	}
 }
+
+func TestIsCanonical_Ok(t *testing.T) {
+	ctx := context.Background()
+	db, sc := testDB.SetupDB(t)
+	c := setupBeaconChain(t, db, sc)
+
+	blk := testutil.NewBeaconBlock()
+	blk.Block.Slot = 0
+	root, err := blk.Block.HashTreeRoot()
+	require.NoError(t, err)
+	require.NoError(t, db.SaveBlock(ctx, blk))
+	require.NoError(t, db.SaveGenesisBlockRoot(ctx, root))
+	can, err := c.IsCanonical(ctx, root)
+	require.NoError(t, err)
+	assert.Equal(t, true, can)
+
+	can, err = c.IsCanonical(ctx, [32]byte{'a'})
+	require.NoError(t, err)
+	assert.Equal(t, false, can)
+}

beacon-chain/blockchain/head.go

@@ -72,10 +72,6 @@ func (s *Service) updateHead(ctx context.Context, balances []uint64) error {
 		return err
 	}
 
-	if err := s.updateRecentCanonicalBlocks(ctx, headRoot); err != nil {
-		return err
-	}
-
 	// Save head to the local service cache.
 	return s.saveHead(ctx, headRoot)
 }
@@ -252,34 +248,6 @@ func (s *Service) hasHeadState() bool {
 	return s.head != nil && s.head.state != nil
 }
 
-// This updates recent canonical block mapping. It uses input head root and retrieves
-// all the canonical block roots that are ancestor of the input head block root.
-func (s *Service) updateRecentCanonicalBlocks(ctx context.Context, headRoot [32]byte) error {
-	ctx, span := trace.StartSpan(ctx, "blockChain.updateRecentCanonicalBlocks")
-	defer span.End()
-
-	s.recentCanonicalBlocksLock.Lock()
-	defer s.recentCanonicalBlocksLock.Unlock()
-
-	s.recentCanonicalBlocks = make(map[[32]byte]bool)
-	s.recentCanonicalBlocks[headRoot] = true
-	nodes := s.forkChoiceStore.Nodes()
-	node := s.forkChoiceStore.Node(headRoot)
-	if node == nil {
-		return nil
-	}
-
-	for node.Parent() != protoarray.NonExistentNode {
-		if ctx.Err() != nil {
-			return ctx.Err()
-		}
-		node = nodes[node.Parent()]
-		s.recentCanonicalBlocks[node.Root()] = true
-	}
-
-	return nil
-}
-
 // This caches justified state balances to be used for fork choice.
 func (s *Service) cacheJustifiedStateBalances(ctx context.Context, justifiedRoot [32]byte) error {
 	if err := s.beaconDB.SaveBlocks(ctx, s.getInitSyncBlocks()); err != nil {

beacon-chain/blockchain/head_test.go

@@ -94,39 +94,6 @@ func TestSaveHead_Different_Reorg(t *testing.T) {
 	require.LogsContain(t, hook, "Chain reorg occurred")
 }
 
-func TestUpdateRecentCanonicalBlocks_CanUpdateWithoutParent(t *testing.T) {
-	db, sc := testDB.SetupDB(t)
-	service := setupBeaconChain(t, db, sc)
-
-	r := [32]byte{'a'}
-	require.NoError(t, service.updateRecentCanonicalBlocks(context.Background(), r))
-	canonical, err := service.IsCanonical(context.Background(), r)
-	require.NoError(t, err)
-	assert.Equal(t, true, canonical, "Block should be canonical")
-}
-
-func TestUpdateRecentCanonicalBlocks_CanUpdateWithParent(t *testing.T) {
-	db, sc := testDB.SetupDB(t)
-	service := setupBeaconChain(t, db, sc)
-	oldHead := [32]byte{'a'}
-	require.NoError(t, service.forkChoiceStore.ProcessBlock(context.Background(), 1, oldHead, [32]byte{'g'}, [32]byte{}, 0, 0))
-	currentHead := [32]byte{'b'}
-	require.NoError(t, service.forkChoiceStore.ProcessBlock(context.Background(), 3, currentHead, oldHead, [32]byte{}, 0, 0))
-	forkedRoot := [32]byte{'c'}
-	require.NoError(t, service.forkChoiceStore.ProcessBlock(context.Background(), 2, forkedRoot, oldHead, [32]byte{}, 0, 0))
-
-	require.NoError(t, service.updateRecentCanonicalBlocks(context.Background(), currentHead))
-	canonical, err := service.IsCanonical(context.Background(), currentHead)
-	require.NoError(t, err)
-	assert.Equal(t, true, canonical, "Block should be canonical")
-	canonical, err = service.IsCanonical(context.Background(), oldHead)
-	require.NoError(t, err)
-	assert.Equal(t, true, canonical, "Block should be canonical")
-	canonical, err = service.IsCanonical(context.Background(), forkedRoot)
-	require.NoError(t, err)
-	assert.Equal(t, false, canonical, "Block should not be canonical")
-}
-
 func TestCacheJustifiedStateBalances_CanCache(t *testing.T) {
 	db, sc := testDB.SetupDB(t)
 	service := setupBeaconChain(t, db, sc)

beacon-chain/blockchain/service.go

@@ -40,44 +40,42 @@ import (
 // Service represents a service that handles the internal
 // logic of managing the full PoS beacon chain.
 type Service struct {
-	ctx                       context.Context
-	cancel                    context.CancelFunc
-	beaconDB                  db.HeadAccessDatabase
-	depositCache              *depositcache.DepositCache
-	chainStartFetcher         powchain.ChainStartFetcher
-	attPool                   attestations.Pool
-	slashingPool              *slashings.Pool
-	exitPool                  *voluntaryexits.Pool
-	genesisTime               time.Time
-	p2p                       p2p.Broadcaster
-	maxRoutines               int
-	head                      *head
-	headLock                  sync.RWMutex
-	stateNotifier             statefeed.Notifier
-	genesisRoot               [32]byte
-	forkChoiceStore           f.ForkChoicer
-	justifiedCheckpt          *ethpb.Checkpoint
-	prevJustifiedCheckpt      *ethpb.Checkpoint
-	bestJustifiedCheckpt      *ethpb.Checkpoint
-	finalizedCheckpt          *ethpb.Checkpoint
-	prevFinalizedCheckpt      *ethpb.Checkpoint
-	nextEpochBoundarySlot     uint64
-	initSyncState             map[[32]byte]*stateTrie.BeaconState
-	boundaryRoots             [][32]byte
-	checkpointState           *cache.CheckpointStateCache
-	checkpointStateLock       sync.Mutex
-	stateGen                  *stategen.State
-	opsService                *attestations.Service
-	initSyncBlocks            map[[32]byte]*ethpb.SignedBeaconBlock
-	initSyncBlocksLock        sync.RWMutex
-	recentCanonicalBlocks     map[[32]byte]bool
-	recentCanonicalBlocksLock sync.RWMutex
-	justifiedBalances         []uint64
-	justifiedBalancesLock     sync.RWMutex
-	checkPtInfoCache          *checkPtInfoCache
-	wsEpoch                   uint64
-	wsRoot                    []byte
-	wsVerified                bool
+	ctx                   context.Context
+	cancel                context.CancelFunc
+	beaconDB              db.HeadAccessDatabase
+	depositCache          *depositcache.DepositCache
+	chainStartFetcher     powchain.ChainStartFetcher
+	attPool               attestations.Pool
+	slashingPool          *slashings.Pool
+	exitPool              *voluntaryexits.Pool
+	genesisTime           time.Time
+	p2p                   p2p.Broadcaster
+	maxRoutines           int
+	head                  *head
+	headLock              sync.RWMutex
+	stateNotifier         statefeed.Notifier
+	genesisRoot           [32]byte
+	forkChoiceStore       f.ForkChoicer
+	justifiedCheckpt      *ethpb.Checkpoint
+	prevJustifiedCheckpt  *ethpb.Checkpoint
+	bestJustifiedCheckpt  *ethpb.Checkpoint
+	finalizedCheckpt      *ethpb.Checkpoint
+	prevFinalizedCheckpt  *ethpb.Checkpoint
+	nextEpochBoundarySlot uint64
+	initSyncState         map[[32]byte]*stateTrie.BeaconState
+	boundaryRoots         [][32]byte
+	checkpointState       *cache.CheckpointStateCache
+	checkpointStateLock   sync.Mutex
+	stateGen              *stategen.State
+	opsService            *attestations.Service
+	initSyncBlocks        map[[32]byte]*ethpb.SignedBeaconBlock
+	initSyncBlocksLock    sync.RWMutex
+	justifiedBalances     []uint64
+	justifiedBalancesLock sync.RWMutex
+	checkPtInfoCache      *checkPtInfoCache
+	wsEpoch               uint64
+	wsRoot                []byte
+	wsVerified            bool
 }
 
 // Config options for the service.
@@ -104,29 +102,28 @@ type Config struct {
 func NewService(ctx context.Context, cfg *Config) (*Service, error) {
 	ctx, cancel := context.WithCancel(ctx)
 	return &Service{
-		ctx:                   ctx,
-		cancel:                cancel,
-		beaconDB:              cfg.BeaconDB,
-		depositCache:          cfg.DepositCache,
-		chainStartFetcher:     cfg.ChainStartFetcher,
-		attPool:               cfg.AttPool,
-		exitPool:              cfg.ExitPool,
-		slashingPool:          cfg.SlashingPool,
-		p2p:                   cfg.P2p,
-		maxRoutines:           cfg.MaxRoutines,
-		stateNotifier:         cfg.StateNotifier,
-		forkChoiceStore:       cfg.ForkChoiceStore,
-		initSyncState:         make(map[[32]byte]*stateTrie.BeaconState),
-		boundaryRoots:         [][32]byte{},
-		checkpointState:       cache.NewCheckpointStateCache(),
-		opsService:            cfg.OpsService,
-		stateGen:              cfg.StateGen,
-		initSyncBlocks:        make(map[[32]byte]*ethpb.SignedBeaconBlock),
-		recentCanonicalBlocks: make(map[[32]byte]bool),
-		justifiedBalances:     make([]uint64, 0),
-		checkPtInfoCache:      newCheckPointInfoCache(),
-		wsEpoch:               cfg.WspEpoch,
-		wsRoot:                cfg.WspBlockRoot,
+		ctx:               ctx,
+		cancel:            cancel,
+		beaconDB:          cfg.BeaconDB,
+		depositCache:      cfg.DepositCache,
+		chainStartFetcher: cfg.ChainStartFetcher,
+		attPool:           cfg.AttPool,
+		exitPool:          cfg.ExitPool,
+		slashingPool:      cfg.SlashingPool,
+		p2p:               cfg.P2p,
+		maxRoutines:       cfg.MaxRoutines,
+		stateNotifier:     cfg.StateNotifier,
+		forkChoiceStore:   cfg.ForkChoiceStore,
+		initSyncState:     make(map[[32]byte]*stateTrie.BeaconState),
+		boundaryRoots:     [][32]byte{},
+		checkpointState:   cache.NewCheckpointStateCache(),
+		opsService:        cfg.OpsService,
+		stateGen:          cfg.StateGen,
+		initSyncBlocks:    make(map[[32]byte]*ethpb.SignedBeaconBlock),
+		justifiedBalances: make([]uint64, 0),
+		checkPtInfoCache:  newCheckPointInfoCache(),
+		wsEpoch:           cfg.WspEpoch,
+		wsRoot:            cfg.WspBlockRoot,
 	}, nil
 }
 

beacon-chain/forkchoice/interfaces.go

@@ -43,4 +43,5 @@ type Getter interface {
 	Store() *protoarray.Store
 	HasParent(root [32]byte) bool
 	AncestorRoot(ctx context.Context, root [32]byte, slot uint64) ([]byte, error)
+	IsCanonical(root [32]byte) bool
 }

beacon-chain/forkchoice/protoarray/nodes.go

@@ -51,9 +51,47 @@ func (s *Store) head(ctx context.Context, justifiedRoot [32]byte) ([32]byte, err
 		lastHeadRoot = bestNode.root
 	}
 
+	// Update canonical mapping given the head root.
+	if err := s.updateCanonicalNodes(ctx, bestNode.root); err != nil {
+		return [32]byte{}, err
+	}
+
 	return bestNode.root, nil
 }
 
+// updateCanonicalNodes updates the canonical nodes mapping given the input block root.
+func (s *Store) updateCanonicalNodes(ctx context.Context, root [32]byte) error {
+	ctx, span := trace.StartSpan(ctx, "protoArrayForkChoice.updateCanonicalNodes")
+	defer span.End()
+
+	// Set the input node to canonical.
+	s.canonicalNodes[root] = true
+
+	// Get the input's parent node index.
+	i := s.nodesIndices[root]
+	n := s.nodes[i]
+	p := n.parent
+
+	for p != NonExistentNode {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+
+		// Get the parent node, if the node is already in canonical mapping,
+		// we can be sure rest of the ancestors are canonical. Exit early.
+		n = s.nodes[p]
+		if s.canonicalNodes[n.root] {
+			break
+		}
+
+		// Set parent node to canonical. Repeat until parent node index is undefined.
+		s.canonicalNodes[n.root] = true
+		p = n.parent
+	}
+
+	return nil
+}
+
 // insert registers a new block node to the fork choice store's node list.
 // It then updates the new node's parent with best child and descendant node.
 func (s *Store) insert(ctx context.Context,