Faster ancestor root look up via fork choice store (#6753)

* forkchoice: add `HasParent` and `AncestorRoot` getters
* Tests
* Interfaces
* process block: use getters in `ancestor`
* Update tests
* Merge branch 'master' of github.com:prysmaticlabs/prysm into faster-ancestor-block
* feedback from @rauljordan

beacon-chain/blockchain/process_attestation_test.go

@@ -358,7 +358,7 @@ func TestVerifyLMDFFGConsistent_NotOK(t *testing.T) {
 	ctx := context.Background()
 	db, _ := testDB.SetupDB(t)
 
-	cfg := &Config{BeaconDB: db}
+	cfg := &Config{BeaconDB: db, ForkChoiceStore: protoarray.New(0, 0, [32]byte{})}
 	service, err := NewService(ctx, cfg)
 	require.NoError(t, err)
 
@@ -379,7 +379,7 @@ func TestVerifyLMDFFGConsistent_OK(t *testing.T) {
 	ctx := context.Background()
 	db, _ := testDB.SetupDB(t)
 
-	cfg := &Config{BeaconDB: db}
+	cfg := &Config{BeaconDB: db, ForkChoiceStore: protoarray.New(0, 0, [32]byte{})}
 	service, err := NewService(ctx, cfg)
 	require.NoError(t, err)
 

beacon-chain/blockchain/process_block_helpers.go

@@ -255,13 +255,20 @@ func (s *Service) ancestor(ctx context.Context, root []byte, slot uint64) ([]byt
 		return nil, ctx.Err()
 	}
 
-	signed, err := s.beaconDB.Block(ctx, bytesutil.ToBytes32(root))
+	r := bytesutil.ToBytes32(root)
+	// Get ancestor root from fork choice store instead of recursively looking up blocks in DB.
+	// This is most optimal outcome.
+	if s.forkChoiceStore.HasParent(r) {
+		return s.forkChoiceStore.AncestorRoot(ctx, r, slot)
+	}
+
+	signed, err := s.beaconDB.Block(ctx, r)
 	if err != nil {
 		return nil, errors.Wrap(err, "could not get ancestor block")
 	}
 
-	if s.hasInitSyncBlock(bytesutil.ToBytes32(root)) {
-		signed = s.getInitSyncBlock(bytesutil.ToBytes32(root))
+	if s.hasInitSyncBlock(r) {
+		signed = s.getInitSyncBlock(r)
 	}
 
 	if signed == nil || signed.Block == nil {

beacon-chain/blockchain/process_block_test.go

@@ -33,8 +33,9 @@ func TestStore_OnBlock(t *testing.T) {
 	db, sc := testDB.SetupDB(t)
 
 	cfg := &Config{
-		BeaconDB: db,
-		StateGen: stategen.New(db, sc),
+		BeaconDB:        db,
+		StateGen:        stategen.New(db, sc),
+		ForkChoiceStore: protoarray.New(0, 0, [32]byte{}),
 	}
 	service, err := NewService(ctx, cfg)
 	require.NoError(t, err)
@@ -153,7 +154,7 @@ func TestRemoveStateSinceLastFinalized_EmptyStartSlot(t *testing.T) {
 	params.UseMinimalConfig()
 	defer params.UseMainnetConfig()
 
-	cfg := &Config{BeaconDB: db}
+	cfg := &Config{BeaconDB: db, ForkChoiceStore: protoarray.New(0, 0, [32]byte{})}
 	service, err := NewService(ctx, cfg)
 	require.NoError(t, err)
 	service.genesisTime = time.Now()
@@ -456,7 +457,7 @@ func TestAncestor_HandleSkipSlot(t *testing.T) {
 	ctx := context.Background()
 	db, _ := testDB.SetupDB(t)
 
-	cfg := &Config{BeaconDB: db}
+	cfg := &Config{BeaconDB: db, ForkChoiceStore: protoarray.New(0, 0, [32]byte{})}
 	service, err := NewService(ctx, cfg)
 	require.NoError(t, err)
 
@@ -548,7 +549,7 @@ func TestFinalizedImpliesNewJustified(t *testing.T) {
 	for _, test := range tests {
 		beaconState := testutil.NewBeaconState()
 		require.NoError(t, beaconState.SetCurrentJustifiedCheckpoint(test.args.stateCheckPoint))
-		service, err := NewService(ctx, &Config{BeaconDB: db, StateGen: stategen.New(db, sc)})
+		service, err := NewService(ctx, &Config{BeaconDB: db, StateGen: stategen.New(db, sc), ForkChoiceStore: protoarray.New(0, 0, [32]byte{})})
 		require.NoError(t, err)
 		service.justifiedCheckpt = test.args.cachedCheckPoint
 		require.NoError(t, service.beaconDB.SaveStateSummary(ctx, &pb.StateSummary{Root: bytesutil.PadTo(test.want.Root, 32)}))
@@ -641,7 +642,7 @@ func TestVerifyBlkDescendant(t *testing.T) {
 		},
 	}
 	for _, test := range tests {
-		service, err := NewService(ctx, &Config{BeaconDB: db, StateGen: stategen.New(db, sc)})
+		service, err := NewService(ctx, &Config{BeaconDB: db, StateGen: stategen.New(db, sc), ForkChoiceStore: protoarray.New(0, 0, [32]byte{})})
 		require.NoError(t, err)
 		service.finalizedCheckpt = &ethpb.Checkpoint{
 			Root: test.args.finalizedRoot[:],

beacon-chain/forkchoice/interfaces.go

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

beacon-chain/forkchoice/protoarray/BUILD.bazel

@@ -34,6 +34,7 @@ go_test(
     ],
     embed = [":go_default_library"],
     deps = [
+        "//shared/bytesutil:go_default_library",
         "//shared/hashutil:go_default_library",
         "//shared/params:go_default_library",
         "//shared/testutil/assert:go_default_library",

beacon-chain/forkchoice/protoarray/nodes_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"testing"
 
+	"github.com/prysmaticlabs/prysm/shared/bytesutil"
 	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
 	"github.com/prysmaticlabs/prysm/shared/testutil/require"
 )
@@ -364,3 +365,52 @@ func TestStore_ViableForHead(t *testing.T) {
 		assert.Equal(t, tc.want, s.viableForHead(tc.n))
 	}
 }
+
+func TestStore_HasParent(t *testing.T) {
+	tests := []struct {
+		m    map[[32]byte]uint64
+		n    []*Node
+		r    [32]byte
+		want bool
+	}{
+		{r: [32]byte{'a'}, want: false},
+		{m: map[[32]byte]uint64{[32]byte{'a'}: 0}, r: [32]byte{'a'}, want: false},
+		{m: map[[32]byte]uint64{[32]byte{'a'}: 0}, r: [32]byte{'a'},
+			n: []*Node{{Parent: NonExistentNode}}, want: false},
+		{m: map[[32]byte]uint64{[32]byte{'a'}: 0},
+			n: []*Node{{Parent: 0}}, r: [32]byte{'a'},
+			want: true},
+	}
+	for _, tc := range tests {
+		f := &ForkChoice{store: &Store{
+			NodeIndices: tc.m,
+			Nodes:       tc.n,
+		}}
+		assert.Equal(t, tc.want, f.HasParent(tc.r))
+	}
+}
+
+func TestStore_AncestorRoot(t *testing.T) {
+	ctx := context.Background()
+	f := &ForkChoice{store: &Store{}}
+	f.store.NodeIndices = map[[32]byte]uint64{}
+	_, err := f.AncestorRoot(ctx, [32]byte{'a'}, 0)
+	assert.ErrorContains(t, "node does not exist", err)
+	f.store.NodeIndices[[32]byte{'a'}] = 0
+	_, err = f.AncestorRoot(ctx, [32]byte{'a'}, 0)
+	assert.ErrorContains(t, "node index out of range", err)
+	f.store.NodeIndices[[32]byte{'b'}] = 1
+	f.store.NodeIndices[[32]byte{'c'}] = 2
+	f.store.Nodes = []*Node{
+		{Slot: 1, Root: [32]byte{'a'}, Parent: NonExistentNode},
+		{Slot: 2, Root: [32]byte{'b'}, Parent: 0},
+		{Slot: 3, Root: [32]byte{'c'}, Parent: 1},
+	}
+
+	r, err := f.AncestorRoot(ctx, [32]byte{'c'}, 1)
+	require.NoError(t, err)
+	assert.Equal(t, bytesutil.ToBytes32(r), [32]byte{'a'})
+	r, err = f.AncestorRoot(ctx, [32]byte{'c'}, 2)
+	require.NoError(t, err)
+	assert.Equal(t, bytesutil.ToBytes32(r), [32]byte{'b'})
+}

beacon-chain/forkchoice/protoarray/store.go

@@ -138,3 +138,41 @@ func (f *ForkChoice) HasNode(root [32]byte) bool {
 	_, ok := f.store.NodeIndices[root]
 	return ok
 }
+
+// HasParent returns true if the node parent exists in fork choice store,
+// false else wise.
+func (f *ForkChoice) HasParent(root [32]byte) bool {
+	f.store.nodeIndicesLock.RLock()
+	defer f.store.nodeIndicesLock.RUnlock()
+
+	i, ok := f.store.NodeIndices[root]
+	if !ok || i >= uint64(len(f.store.Nodes)) {
+		return false
+	}
+
+	return f.store.Nodes[i].Parent != NonExistentNode
+}
+
+// AncestorRoot returns the ancestor root of input block root at a given slot.
+func (f *ForkChoice) AncestorRoot(ctx context.Context, root [32]byte, slot uint64) ([]byte, error) {
+	i, ok := f.store.NodeIndices[root]
+	if !ok {
+		return nil, errors.New("node does not exist")
+	}
+	if i >= uint64(len(f.store.Nodes)) {
+		return nil, errors.New("node index out of range")
+	}
+
+	for f.store.Nodes[i].Slot > slot {
+		if ctx.Err() != nil {
+			return nil, ctx.Err()
+		}
+
+		i = f.store.Nodes[i].Parent
+	}
+	if i >= uint64(len(f.store.Nodes)) {
+		return nil, errors.New("node index out of range")
+	}
+
+	return f.store.Nodes[i].Root[:], nil
+}