Prune during on_tick (#11387)

* Prune during on_tick

* add unit test

* fix tests

Co-authored-by: terencechain <terence@prysmaticlabs.com>
Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>

beacon-chain/blockchain/process_block_test.go

@@ -1171,6 +1171,7 @@ func TestOnBlock_CanFinalize_WithOnTick(t *testing.T) {
 
 	gs, keys := util.DeterministicGenesisState(t, 32)
 	require.NoError(t, service.saveGenesisData(ctx, gs))
+	require.NoError(t, fcs.UpdateFinalizedCheckpoint(&forkchoicetypes.Checkpoint{Root: service.originBlockRoot}))
 
 	testState := gs.Copy()
 	for i := types.Slot(1); i <= 4*params.BeaconConfig().SlotsPerEpoch; i++ {

beacon-chain/forkchoice/doubly-linked-tree/forkchoice.go

@@ -36,7 +36,6 @@ func New() *ForkChoice {
 		nodeByRoot:                    make(map[[fieldparams.RootLength]byte]*Node),
 		nodeByPayload:                 make(map[[fieldparams.RootLength]byte]*Node),
 		slashedIndices:                make(map[types.ValidatorIndex]bool),
-		pruneThreshold:                defaultPruneThreshold,
 		receivedBlocksLastEpoch:       [fieldparams.SlotsPerEpoch]types.Slot{},
 	}
 

beacon-chain/forkchoice/doubly-linked-tree/node.go

@@ -11,15 +11,6 @@ import (
 	v1 "github.com/prysmaticlabs/prysm/v3/proto/eth/v1"
 )
 
-// depth returns the length of the path to the root of Fork Choice
-func (n *Node) depth() uint64 {
-	ret := uint64(0)
-	for node := n.parent; node != nil; node = node.parent {
-		ret += 1
-	}
-	return ret
-}
-
 // applyWeightChanges recomputes the weight of the node passed as an argument and all of its descendants,
 // using the current balance stored in each node. This function requires a lock
 // in Store.nodesLock

beacon-chain/forkchoice/doubly-linked-tree/node_test.go

@@ -140,25 +140,6 @@ func TestNode_UpdateBestDescendant_LowerWeightChild(t *testing.T) {
 	assert.Equal(t, s.treeRootNode.children[0], s.treeRootNode.bestDescendant)
 }
 
-func TestNode_TestDepth(t *testing.T) {
-	f := setup(1, 1)
-	ctx := context.Background()
-	// Input child is best descendant
-	state, blkRoot, err := prepareForkchoiceState(ctx, 1, indexToHash(1), params.BeaconConfig().ZeroHash, params.BeaconConfig().ZeroHash, 1, 1)
-	require.NoError(t, err)
-	require.NoError(t, f.InsertNode(ctx, state, blkRoot))
-	state, blkRoot, err = prepareForkchoiceState(ctx, 2, indexToHash(2), indexToHash(1), params.BeaconConfig().ZeroHash, 1, 1)
-	require.NoError(t, err)
-	require.NoError(t, f.InsertNode(ctx, state, blkRoot))
-	state, blkRoot, err = prepareForkchoiceState(ctx, 3, indexToHash(3), params.BeaconConfig().ZeroHash, params.BeaconConfig().ZeroHash, 1, 1)
-	require.NoError(t, err)
-	require.NoError(t, f.InsertNode(ctx, state, blkRoot))
-
-	s := f.store
-	require.Equal(t, s.nodeByRoot[indexToHash(2)].depth(), uint64(2))
-	require.Equal(t, s.nodeByRoot[indexToHash(3)].depth(), uint64(1))
-}
-
 func TestNode_ViableForHead(t *testing.T) {
 	tests := []struct {
 		n              *Node

beacon-chain/forkchoice/doubly-linked-tree/on_tick.go

@@ -70,5 +70,5 @@ func (f *ForkChoice) NewSlot(ctx context.Context, slot types.Slot) error {
 	if !features.Get().DisablePullTips {
 		f.updateUnrealizedCheckpoints()
 	}
-	return nil
+	return f.store.prune(ctx)
 }

beacon-chain/forkchoice/doubly-linked-tree/store.go

@@ -12,10 +12,6 @@ import (
 	"go.opencensus.io/trace"
 )
 
-// This defines the minimal number of block nodes that can be in the tree
-// before getting pruned upon new finalization.
-const defaultPruneThreshold = 256
-
 // applyProposerBoostScore applies the current proposer boost scores to the
 // relevant nodes. This function requires a lock in Store.nodesLock.
 func (s *Store) applyProposerBoostScore(newBalances []uint64) error {
@@ -59,11 +55,6 @@ func (s *Store) proposerBoost() [fieldparams.RootLength]byte {
 	return s.proposerBoostRoot
 }
 
-// PruneThreshold of fork choice store.
-func (s *Store) PruneThreshold() uint64 {
-	return s.pruneThreshold
-}
-
 // head starts from justified root and then follows the best descendant links
 // to find the best block for head. This function assumes a lock on s.nodesLock
 func (s *Store) head(ctx context.Context) ([32]byte, error) {
@@ -216,8 +207,7 @@ func (s *Store) pruneFinalizedNodeByRootMap(ctx context.Context, node, finalized
 	return nil
 }
 
-// prune prunes the fork choice store with the new finalized root. The store is only pruned if the input
-// root is different than the current store finalized root, and the number of the store has met prune threshold.
+// prune prunes the fork choice store. It removes all nodes that compete with the finalized root.
 // This function does not prune for invalid optimistically synced nodes, it deals only with pruning upon finalization
 func (s *Store) prune(ctx context.Context) error {
 	ctx, span := trace.StartSpan(ctx, "doublyLinkedForkchoice.Prune")
@@ -233,10 +223,8 @@ func (s *Store) prune(ctx context.Context) error {
 	if !ok || finalizedNode == nil {
 		return errUnknownFinalizedRoot
 	}
-
-	// The number of the nodes has not met the prune threshold.
-	// Pruning at small numbers incurs more cost than benefit.
-	if finalizedNode.depth() < s.pruneThreshold {
+	// return early if we haven't changed the finalized checkpoint
+	if finalizedNode.parent == nil {
 		return nil
 	}
 

beacon-chain/forkchoice/doubly-linked-tree/store_test.go

@@ -12,15 +12,6 @@ import (
 	"github.com/prysmaticlabs/prysm/v3/testing/require"
 )
 
-func TestStore_PruneThreshold(t *testing.T) {
-	s := &Store{
-		pruneThreshold: defaultPruneThreshold,
-	}
-	if got := s.PruneThreshold(); got != defaultPruneThreshold {
-		t.Errorf("PruneThreshold() = %v, want %v", got, defaultPruneThreshold)
-	}
-}
-
 func TestStore_JustifiedEpoch(t *testing.T) {
 	j := types.Epoch(100)
 	f := setup(j, j)
@@ -154,7 +145,7 @@ func TestStore_Insert(t *testing.T) {
 	assert.Equal(t, indexToHash(100), child.root, "Incorrect root")
 }
 
-func TestStore_Prune_LessThanThreshold(t *testing.T) {
+func TestStore_Prune_MoreThanThreshold(t *testing.T) {
 	// Define 100 nodes in store.
 	numOfNodes := uint64(100)
 	f := setup(0, 0)
@@ -169,16 +160,14 @@ func TestStore_Prune_LessThanThreshold(t *testing.T) {
 	}
 
 	s := f.store
-	s.pruneThreshold = 100
 
-	// Finalized root has depth 99 so everything before it should be pruned,
-	// but PruneThreshold is at 100 so nothing will be pruned.
+	// Finalized root is at index 99 so everything before 99 should be pruned.
 	s.finalizedCheckpoint.Root = indexToHash(99)
 	require.NoError(t, s.prune(context.Background()))
-	assert.Equal(t, 100, len(s.nodeByRoot), "Incorrect nodes count")
+	assert.Equal(t, 1, len(s.nodeByRoot), "Incorrect nodes count")
 }
 
-func TestStore_Prune_MoreThanThreshold(t *testing.T) {
+func TestStore_Prune_MoreThanOnce(t *testing.T) {
 	// Define 100 nodes in store.
 	numOfNodes := uint64(100)
 	f := setup(0, 0)
@@ -193,15 +182,19 @@ func TestStore_Prune_MoreThanThreshold(t *testing.T) {
 	}
 
 	s := f.store
-	s.pruneThreshold = 0
 
-	// Finalized root is at index 99 so everything before 99 should be pruned.
-	s.finalizedCheckpoint.Root = indexToHash(99)
+	// Finalized root is at index 11 so everything before 11 should be pruned.
+	s.finalizedCheckpoint.Root = indexToHash(10)
 	require.NoError(t, s.prune(context.Background()))
-	assert.Equal(t, 1, len(s.nodeByRoot), "Incorrect nodes count")
+	assert.Equal(t, 90, len(s.nodeByRoot), "Incorrect nodes count")
+
+	// One more time.
+	s.finalizedCheckpoint.Root = indexToHash(20)
+	require.NoError(t, s.prune(context.Background()))
+	assert.Equal(t, 80, len(s.nodeByRoot), "Incorrect nodes count")
 }
 
-func TestStore_Prune_MoreThanOnce(t *testing.T) {
+func TestStore_Prune_ReturnEarly(t *testing.T) {
 	// Define 100 nodes in store.
 	numOfNodes := uint64(100)
 	f := setup(0, 0)
@@ -214,19 +207,10 @@ func TestStore_Prune_MoreThanOnce(t *testing.T) {
 		require.NoError(t, err)
 		require.NoError(t, f.InsertNode(ctx, state, blkRoot))
 	}
-
-	s := f.store
-	s.pruneThreshold = 0
-
-	// Finalized root is at index 11 so everything before 11 should be pruned.
-	s.finalizedCheckpoint.Root = indexToHash(10)
-	require.NoError(t, s.prune(context.Background()))
-	assert.Equal(t, 90, len(s.nodeByRoot), "Incorrect nodes count")
-
-	// One more time.
-	s.finalizedCheckpoint.Root = indexToHash(20)
-	require.NoError(t, s.prune(context.Background()))
-	assert.Equal(t, 80, len(s.nodeByRoot), "Incorrect nodes count")
+	require.NoError(t, f.store.prune(ctx))
+	nodeCount := f.NodeCount()
+	require.NoError(t, f.store.prune(ctx))
+	require.Equal(t, nodeCount, f.NodeCount())
 }
 
 // This unit tests starts with a simple branch like this
@@ -245,7 +229,6 @@ func TestStore_Prune_NoDanglingBranch(t *testing.T) {
 	state, blkRoot, err = prepareForkchoiceState(ctx, 2, indexToHash(2), params.BeaconConfig().ZeroHash, params.BeaconConfig().ZeroHash, 0, 0)
 	require.NoError(t, err)
 	require.NoError(t, f.InsertNode(ctx, state, blkRoot))
-	f.store.pruneThreshold = 0
 
 	s := f.store
 	s.finalizedCheckpoint.Root = indexToHash(1)
@@ -330,7 +313,6 @@ func TestStore_PruneMapsNodes(t *testing.T) {
 	require.NoError(t, f.InsertNode(ctx, state, blkRoot))
 
 	s := f.store
-	s.pruneThreshold = 0
 	s.finalizedCheckpoint.Root = indexToHash(1)
 	require.NoError(t, s.prune(context.Background()))
 	require.Equal(t, len(s.nodeByRoot), 1)

beacon-chain/forkchoice/doubly-linked-tree/types.go

@@ -24,7 +24,6 @@ type Store struct {
 	unrealizedFinalizedCheckpoint *forkchoicetypes.Checkpoint            // best unrealized finalized checkpoint in store.
 	prevJustifiedCheckpoint       *forkchoicetypes.Checkpoint            // previous justified checkpoint in store.
 	finalizedCheckpoint           *forkchoicetypes.Checkpoint            // latest finalized epoch in store.
-	pruneThreshold                uint64                                 // do not prune tree unless threshold is reached.
 	proposerBoostRoot             [fieldparams.RootLength]byte           // latest block root that was boosted after being received in a timely manner.
 	previousProposerBoostRoot     [fieldparams.RootLength]byte           // previous block root that was boosted after being received in a timely manner.
 	previousProposerBoostScore    uint64                                 // previous proposer boosted root score.

beacon-chain/forkchoice/doubly-linked-tree/vote_test.go

@@ -261,14 +261,6 @@ func TestVotes_CanFindHead(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, indexToHash(10), r, "Incorrect head for with justified epoch at 3")
 
-	// Verify pruning below the prune threshold does not affect head.
-	f.store.pruneThreshold = 1000
-	prevRoot := f.store.finalizedCheckpoint.Root
-	f.store.finalizedCheckpoint.Root = indexToHash(5)
-	require.NoError(t, f.store.prune(context.Background()))
-	assert.Equal(t, 11, len(f.store.nodeByRoot), "Incorrect nodes length after prune")
-
-	f.store.finalizedCheckpoint.Root = prevRoot
 	r, err = f.Head(context.Background(), balances)
 	require.NoError(t, err)
 	assert.Equal(t, indexToHash(10), r, "Incorrect head for with justified epoch at 3")
@@ -289,13 +281,11 @@ func TestVotes_CanFindHead(t *testing.T) {
 	//          8
 	//         / \
 	//        9  10
-	f.store.pruneThreshold = 1
 	f.store.finalizedCheckpoint.Root = indexToHash(5)
 	require.NoError(t, f.store.prune(context.Background()))
 	assert.Equal(t, 5, len(f.store.nodeByRoot), "Incorrect nodes length after prune")
 	// we pruned artificially the justified root.
 	f.store.justifiedCheckpoint.Root = indexToHash(5)
-	f.store.finalizedCheckpoint.Root = prevRoot
 
 	r, err = f.Head(context.Background(), balances)
 	require.NoError(t, err)