Prune canonical attestations when head changes (#11771)

* Only prune canonical attestations

* Feedback from potuz

* Merge conflict

---------

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

beacon-chain/blockchain/forkchoice_update_execution.go

@@ -65,6 +65,11 @@ func (s *Service) forkchoiceUpdateWithExecution(ctx context.Context, newHeadRoot
 		if err := s.saveHead(ctx, newHeadRoot, headBlock, headState); err != nil {
 			log.WithError(err).Error("could not save head")
 		}
+
+		// Only need to prune attestations from pool if the head has changed.
+		if err := s.pruneAttsFromPool(headBlock); err != nil {
+			return err
+		}
 	}
 
 	return nil

beacon-chain/blockchain/process_block.go

@@ -219,10 +219,6 @@ func (s *Service) onBlock(ctx context.Context, signed interfaces.SignedBeaconBlo
 		return err
 	}
 
-	if err := s.pruneCanonicalAttsFromPool(ctx, blockRoot, signed); err != nil {
-		return err
-	}
-
 	// Send notification of the processed block to the state feed.
 	s.cfg.StateNotifier.StateFeed().Send(&feed.Event{
 		Type: statefeed.BlockProcessed,
@@ -600,17 +596,8 @@ func (s *Service) savePostStateInfo(ctx context.Context, r [32]byte, b interface
 	return nil
 }
 
-// This removes the attestations from the mem pool. It will only remove the attestations if input root `r` is canonical,
-// meaning the block `b` is part of the canonical chain.
-func (s *Service) pruneCanonicalAttsFromPool(ctx context.Context, r [32]byte, b interfaces.SignedBeaconBlock) error {
-	canonical, err := s.IsCanonical(ctx, r)
-	if err != nil {
-		return err
-	}
-	if !canonical {
-		return nil
-	}
-
+// This removes the attestations in block `b` from the attestation mem pool.
+func (s *Service) pruneAttsFromPool(b interfaces.SignedBeaconBlock) error {
 	atts := b.Block().Body().Attestations()
 	for _, att := range atts {
 		if helpers.IsAggregated(att) {

beacon-chain/blockchain/process_block_test.go

@@ -1072,7 +1072,7 @@ func TestInsertFinalizedDeposits_MultipleFinalizedRoutines(t *testing.T) {
 	}
 }
 
-func TestRemoveBlockAttestationsInPool_Canonical(t *testing.T) {
+func TestRemoveBlockAttestationsInPool(t *testing.T) {
 	genesis, keys := util.DeterministicGenesisState(t, 64)
 	b, err := util.GenerateFullBlock(genesis, keys, util.DefaultBlockGenConfig(), 1)
 	assert.NoError(t, err)
@@ -1089,29 +1089,10 @@ func TestRemoveBlockAttestationsInPool_Canonical(t *testing.T) {
 	require.NoError(t, service.cfg.AttPool.SaveAggregatedAttestations(atts))
 	wsb, err := consensusblocks.NewSignedBeaconBlock(b)
 	require.NoError(t, err)
-	require.NoError(t, service.pruneCanonicalAttsFromPool(ctx, r, wsb))
+	require.NoError(t, service.pruneAttsFromPool(wsb))
 	require.Equal(t, 0, service.cfg.AttPool.AggregatedAttestationCount())
 }
 
-func TestRemoveBlockAttestationsInPool_NonCanonical(t *testing.T) {
-	genesis, keys := util.DeterministicGenesisState(t, 64)
-	b, err := util.GenerateFullBlock(genesis, keys, util.DefaultBlockGenConfig(), 1)
-	assert.NoError(t, err)
-	r, err := b.Block.HashTreeRoot()
-	require.NoError(t, err)
-
-	ctx := context.Background()
-	beaconDB := testDB.SetupDB(t)
-	service := setupBeaconChain(t, beaconDB)
-
-	atts := b.Block.Body.Attestations
-	require.NoError(t, service.cfg.AttPool.SaveAggregatedAttestations(atts))
-	wsb, err := consensusblocks.NewSignedBeaconBlock(b)
-	require.NoError(t, err)
-	require.NoError(t, service.pruneCanonicalAttsFromPool(ctx, r, wsb))
-	require.Equal(t, 1, service.cfg.AttPool.AggregatedAttestationCount())
-}
-
 func Test_getStateVersionAndPayload(t *testing.T) {
 	tests := []struct {
 		name    string

beacon-chain/sync/subscriber_beacon_blocks.go

@@ -4,10 +4,8 @@ import (
 	"context"
 
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/blockchain"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/core/helpers"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/core/transition/interop"
 	"github.com/prysmaticlabs/prysm/v3/consensus-types/blocks"
-	ethpb "github.com/prysmaticlabs/prysm/v3/proto/prysm/v1alpha1"
 	"google.golang.org/protobuf/proto"
 )
 
@@ -47,21 +45,3 @@ func (s *Service) beaconBlockSubscriber(ctx context.Context, msg proto.Message)
 	}
 	return err
 }
-
-// The input attestations are seen by the network, this deletes them from pool
-// so proposers don't include them in a block for the future.
-func (s *Service) deleteAttsInPool(atts []*ethpb.Attestation) error {
-	for _, att := range atts {
-		if helpers.IsAggregated(att) {
-			if err := s.cfg.attPool.DeleteAggregatedAttestation(att); err != nil {
-				return err
-			}
-		} else {
-			// Ideally there's shouldn't be any unaggregated attestation in the block.
-			if err := s.cfg.attPool.DeleteUnaggregatedAttestation(att); err != nil {
-				return err
-			}
-		}
-	}
-	return nil
-}

beacon-chain/sync/subscriber_beacon_blocks_test.go

@@ -20,27 +20,6 @@ import (
 	"google.golang.org/protobuf/proto"
 )
 
-func TestDeleteAttsInPool(t *testing.T) {
-	r := &Service{
-		cfg: &config{attPool: attestations.NewPool()},
-	}
-
-	att1 := util.HydrateAttestation(&ethpb.Attestation{AggregationBits: bitfield.Bitlist{0b1101}})
-	att2 := util.HydrateAttestation(&ethpb.Attestation{AggregationBits: bitfield.Bitlist{0b1110}})
-	att3 := util.HydrateAttestation(&ethpb.Attestation{AggregationBits: bitfield.Bitlist{0b1011}})
-	att4 := util.HydrateAttestation(&ethpb.Attestation{AggregationBits: bitfield.Bitlist{0b1001}})
-	require.NoError(t, r.cfg.attPool.SaveAggregatedAttestation(att1))
-	require.NoError(t, r.cfg.attPool.SaveAggregatedAttestation(att2))
-	require.NoError(t, r.cfg.attPool.SaveAggregatedAttestation(att3))
-	require.NoError(t, r.cfg.attPool.SaveUnaggregatedAttestation(att4))
-
-	// Seen 1, 3 and 4 in block.
-	require.NoError(t, r.deleteAttsInPool([]*ethpb.Attestation{att1, att3, att4}))
-
-	// Only 2 should remain.
-	assert.DeepEqual(t, []*ethpb.Attestation{att2}, r.cfg.attPool.AggregatedAttestations(), "Did not get wanted attestations")
-}
-
 func TestService_beaconBlockSubscriber(t *testing.T) {
 	pooledAttestations := []*ethpb.Attestation{
 		// Aggregated.