Proposer attestation selection using max-cover (#8571)

* Proposer attestation selection using max-cover

* better alisgn struct field

* more tests

* cleanup

* simplify expressions

* add benchmarks

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

beacon-chain/rpc/validator/BUILD.bazel

@@ -39,6 +39,7 @@ go_library(
         "//beacon-chain/sync:go_default_library",
         "//proto/beacon/db:go_default_library",
         "//proto/beacon/p2p/v1:go_default_library",
+        "//shared/aggregation:go_default_library",
         "//shared/aggregation/attestations:go_default_library",
         "//shared/bls:go_default_library",
         "//shared/bytesutil:go_default_library",
@@ -56,13 +57,15 @@ go_library(
         "@com_github_pkg_errors//:go_default_library",
         "@com_github_prysmaticlabs_eth2_types//:go_default_library",
         "@com_github_prysmaticlabs_ethereumapis//eth/v1alpha1:go_default_library",
+        "@com_github_prysmaticlabs_go_bitfield//:go_default_library",
         "@com_github_sirupsen_logrus//:go_default_library",
         "@io_opencensus_go//trace:go_default_library",
         "@org_golang_google_grpc//codes:go_default_library",
         "@org_golang_google_grpc//status:go_default_library",
     ],
 )
 
+# gazelle:exclude proposer_utils_bench_test.go
 go_test(
     name = "go_default_test",
     srcs = [
@@ -103,6 +106,7 @@ go_test(
         "//shared/bls:go_default_library",
         "//shared/bytesutil:go_default_library",
         "//shared/event:go_default_library",
+        "//shared/featureconfig:go_default_library",
         "//shared/mock:go_default_library",
         "//shared/params:go_default_library",
         "//shared/testutil:go_default_library",

beacon-chain/rpc/validator/proposer_utils.go

@@ -4,9 +4,13 @@ import (
 	"context"
 	"sort"
 
+	types "github.com/prysmaticlabs/eth2-types"
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
+	"github.com/prysmaticlabs/go-bitfield"
 	"github.com/prysmaticlabs/prysm/beacon-chain/core/blocks"
 	stateTrie "github.com/prysmaticlabs/prysm/beacon-chain/state"
+	"github.com/prysmaticlabs/prysm/shared/aggregation"
+	"github.com/prysmaticlabs/prysm/shared/featureconfig"
 	"github.com/prysmaticlabs/prysm/shared/params"
 )
 
@@ -33,6 +37,9 @@ func (a proposerAtts) sortByProfitability() proposerAtts {
 	if len(a) < 2 {
 		return a
 	}
+	if featureconfig.Get().ProposerAttsSelectionUsingMaxCover {
+		return a.sortByProfitabilityUsingMaxCover()
+	}
 	sort.Slice(a, func(i, j int) bool {
 		if a[i].Data.Slot == a[j].Data.Slot {
 			return a[i].AggregationBits.Count() > a[j].AggregationBits.Count()
@@ -42,6 +49,65 @@ func (a proposerAtts) sortByProfitability() proposerAtts {
 	return a
 }
 
+// sortByProfitabilityUsingMaxCover orders attestations by highest slot and by highest aggregation bit count.
+// Duplicate bits are counted only once, using max-cover algorithm.
+func (a proposerAtts) sortByProfitabilityUsingMaxCover() proposerAtts {
+	// Separate attestations by slot, as slot number takes higher precedence when sorting.
+	var slots []types.Slot
+	attsBySlot := map[types.Slot]proposerAtts{}
+	for _, att := range a {
+		if _, ok := attsBySlot[att.Data.Slot]; !ok {
+			slots = append(slots, att.Data.Slot)
+		}
+		attsBySlot[att.Data.Slot] = append(attsBySlot[att.Data.Slot], att)
+	}
+
+	selectAtts := func(atts proposerAtts) proposerAtts {
+		if len(atts) < 2 {
+			return atts
+		}
+		candidates := make([]*bitfield.Bitlist64, len(atts))
+		for i := 0; i < len(atts); i++ {
+			candidates[i] = atts[i].AggregationBits.ToBitlist64()
+		}
+		// Add selected candidates on top, those that are not selected - append at bottom.
+		selectedKeys, _, err := aggregation.MaxCover(candidates, len(candidates), true /* allowOverlaps */)
+		if err == nil {
+			// Pick selected attestations first, leftover attestations will be appended at the end.
+			// Both lists will be sorted by number of bits set.
+			selectedAtts := make(proposerAtts, selectedKeys.Count())
+			leftoverAtts := make(proposerAtts, selectedKeys.Not().Count())
+			for i, key := range selectedKeys.BitIndices() {
+				selectedAtts[i] = atts[key]
+			}
+			for i, key := range selectedKeys.Not().BitIndices() {
+				leftoverAtts[i] = atts[key]
+			}
+			sort.Slice(selectedAtts, func(i, j int) bool {
+				return selectedAtts[i].AggregationBits.Count() > selectedAtts[j].AggregationBits.Count()
+			})
+			sort.Slice(leftoverAtts, func(i, j int) bool {
+				return leftoverAtts[i].AggregationBits.Count() > leftoverAtts[j].AggregationBits.Count()
+			})
+			return append(selectedAtts, leftoverAtts...)
+		}
+		return atts
+	}
+
+	// Select attestations. Slots are sorted from higher to lower at this point. Within slots attestations
+	// are sorted to maximize profitability (greedily selected, with previous attestations' bits
+	// evaluated before including any new attestation).
+	var sortedAtts proposerAtts
+	sort.Slice(slots, func(i, j int) bool {
+		return slots[i] > slots[j]
+	})
+	for _, slot := range slots {
+		sortedAtts = append(sortedAtts, selectAtts(attsBySlot[slot])...)
+	}
+
+	return sortedAtts
+}
+
 // limitToMaxAttestations limits attestations to maximum attestations per block.
 func (a proposerAtts) limitToMaxAttestations() proposerAtts {
 	if uint64(len(a)) > params.BeaconConfig().MaxAttestations {

beacon-chain/rpc/validator/proposer_utils_bench_test.go

@@ -0,0 +1,82 @@
+package validator
+
+import (
+	"fmt"
+	"testing"
+
+	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
+	"github.com/prysmaticlabs/go-bitfield"
+	stateTrie "github.com/prysmaticlabs/prysm/beacon-chain/state"
+	aggtesting "github.com/prysmaticlabs/prysm/shared/aggregation/testing"
+	"github.com/prysmaticlabs/prysm/shared/featureconfig"
+	"github.com/prysmaticlabs/prysm/shared/params"
+)
+
+func BenchmarkProposerAtts_sortByProfitability(b *testing.B) {
+	bitlistLen := params.BeaconConfig().MaxValidatorsPerCommittee
+
+	tests := []struct {
+		name   string
+		inputs []bitfield.Bitlist
+	}{
+		{
+			name:   "256 attestations with single bit set",
+			inputs: aggtesting.BitlistsWithSingleBitSet(256, bitlistLen),
+		},
+		{
+			name:   "256 attestations with 64 random bits set",
+			inputs: aggtesting.BitlistsWithSingleBitSet(256, bitlistLen),
+		},
+		{
+			name:   "512 attestations with single bit set",
+			inputs: aggtesting.BitlistsWithSingleBitSet(512, bitlistLen),
+		},
+		{
+			name:   "1024 attestations with 64 random bits set",
+			inputs: aggtesting.BitlistsWithMultipleBitSet(b, 1024, bitlistLen, 64),
+		},
+		{
+			name:   "1024 attestations with 512 random bits set",
+			inputs: aggtesting.BitlistsWithMultipleBitSet(b, 1024, bitlistLen, 512),
+		},
+		{
+			name:   "1024 attestations with 1000 random bits set",
+			inputs: aggtesting.BitlistsWithMultipleBitSet(b, 1024, bitlistLen, 1000),
+		},
+	}
+
+	runner := func(atts []*ethpb.Attestation) {
+		attsCopy := make(proposerAtts, len(atts))
+		for i, att := range atts {
+			attsCopy[i] = stateTrie.CopyAttestation(att)
+		}
+		attsCopy.sortByProfitability()
+	}
+
+	for _, tt := range tests {
+		b.Run(fmt.Sprintf("naive_%s", tt.name), func(b *testing.B) {
+			b.StopTimer()
+			resetCfg := featureconfig.InitWithReset(&featureconfig.Flags{
+				ProposerAttsSelectionUsingMaxCover: false,
+			})
+			defer resetCfg()
+			atts := aggtesting.MakeAttestationsFromBitlists(tt.inputs)
+			b.StartTimer()
+			for i := 0; i < b.N; i++ {
+				runner(atts)
+			}
+		})
+		b.Run(fmt.Sprintf("max-cover_%s", tt.name), func(b *testing.B) {
+			b.StopTimer()
+			resetCfg := featureconfig.InitWithReset(&featureconfig.Flags{
+				ProposerAttsSelectionUsingMaxCover: true,
+			})
+			defer resetCfg()
+			atts := aggtesting.MakeAttestationsFromBitlists(tt.inputs)
+			b.StartTimer()
+			for i := 0; i < b.N; i++ {
+				runner(atts)
+			}
+		})
+	}
+}

beacon-chain/rpc/validator/proposer_utils_test.go

@@ -5,8 +5,10 @@ import (
 	"sort"
 	"testing"
 
+	types "github.com/prysmaticlabs/eth2-types"
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
 	"github.com/prysmaticlabs/go-bitfield"
+	"github.com/prysmaticlabs/prysm/shared/featureconfig"
 	"github.com/prysmaticlabs/prysm/shared/testutil"
 	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
 	"github.com/prysmaticlabs/prysm/shared/testutil/require"
@@ -33,6 +35,167 @@ func TestProposer_ProposerAtts_sortByProfitability(t *testing.T) {
 	require.DeepEqual(t, want, atts)
 }
 
+func TestProposer_ProposerAtts_sortByProfitabilityUsingMaxCover(t *testing.T) {
+	resetCfg := featureconfig.InitWithReset(&featureconfig.Flags{
+		ProposerAttsSelectionUsingMaxCover: true,
+	})
+	defer resetCfg()
+
+	type testData struct {
+		slot types.Slot
+		bits bitfield.Bitlist
+	}
+	getAtts := func(data []testData) proposerAtts {
+		var atts proposerAtts
+		for _, att := range data {
+			atts = append(atts, testutil.HydrateAttestation(&ethpb.Attestation{
+				Data: &ethpb.AttestationData{Slot: att.slot}, AggregationBits: att.bits}))
+		}
+		return atts
+	}
+
+	t.Run("no atts", func(t *testing.T) {
+		atts := getAtts([]testData{})
+		want := getAtts([]testData{})
+		atts = atts.sortByProfitability()
+		require.DeepEqual(t, want, atts)
+	})
+
+	t.Run("single att", func(t *testing.T) {
+		atts := getAtts([]testData{
+			{4, bitfield.Bitlist{0b11100000, 0b1}},
+		})
+		want := getAtts([]testData{
+			{4, bitfield.Bitlist{0b11100000, 0b1}},
+		})
+		atts = atts.sortByProfitability()
+		require.DeepEqual(t, want, atts)
+	})
+
+	t.Run("single att per slot", func(t *testing.T) {
+		atts := getAtts([]testData{
+			{1, bitfield.Bitlist{0b11000000, 0b1}},
+			{4, bitfield.Bitlist{0b11100000, 0b1}},
+		})
+		want := getAtts([]testData{
+			{4, bitfield.Bitlist{0b11100000, 0b1}},
+			{1, bitfield.Bitlist{0b11000000, 0b1}},
+		})
+		atts = atts.sortByProfitability()
+		require.DeepEqual(t, want, atts)
+	})
+
+	t.Run("two atts on one of the slots", func(t *testing.T) {
+		atts := getAtts([]testData{
+			{1, bitfield.Bitlist{0b11000000, 0b1}},
+			{4, bitfield.Bitlist{0b11100000, 0b1}},
+			{4, bitfield.Bitlist{0b11110000, 0b1}},
+		})
+		want := getAtts([]testData{
+			{4, bitfield.Bitlist{0b11110000, 0b1}},
+			{4, bitfield.Bitlist{0b11100000, 0b1}},
+			{1, bitfield.Bitlist{0b11000000, 0b1}},
+		})
+		atts = atts.sortByProfitability()
+		require.DeepEqual(t, want, atts)
+	})
+
+	t.Run("compare to native sort", func(t *testing.T) {
+		// The naive sort will end up with 0b11001000 being selected second (which is not optimal
+		// as it only contains a single unknown bit).
+		// The max-cover based approach will select 0b00001100 instead, despite lower bit count
+		// (since it has two new/unknown bits).
+		t.Run("naive", func(t *testing.T) {
+			resetCfg := featureconfig.InitWithReset(&featureconfig.Flags{
+				ProposerAttsSelectionUsingMaxCover: false,
+			})
+			defer resetCfg()
+
+			atts := getAtts([]testData{
+				{1, bitfield.Bitlist{0b11000011, 0b1}},
+				{1, bitfield.Bitlist{0b11001000, 0b1}},
+				{1, bitfield.Bitlist{0b00001100, 0b1}},
+			})
+			want := getAtts([]testData{
+				{1, bitfield.Bitlist{0b11000011, 0b1}},
+				{1, bitfield.Bitlist{0b11001000, 0b1}},
+				{1, bitfield.Bitlist{0b00001100, 0b1}},
+			})
+			atts = atts.sortByProfitability()
+			require.DeepEqual(t, want, atts)
+		})
+		t.Run("max-cover", func(t *testing.T) {
+			resetCfg := featureconfig.InitWithReset(&featureconfig.Flags{
+				ProposerAttsSelectionUsingMaxCover: true,
+			})
+			defer resetCfg()
+
+			atts := getAtts([]testData{
+				{1, bitfield.Bitlist{0b11000011, 0b1}},
+				{1, bitfield.Bitlist{0b11001000, 0b1}},
+				{1, bitfield.Bitlist{0b00001100, 0b1}},
+			})
+			want := getAtts([]testData{
+				{1, bitfield.Bitlist{0b11000011, 0b1}},
+				{1, bitfield.Bitlist{0b00001100, 0b1}},
+				{1, bitfield.Bitlist{0b11001000, 0b1}},
+			})
+			atts = atts.sortByProfitability()
+			require.DeepEqual(t, want, atts)
+		})
+	})
+
+	t.Run("multiple slots", func(t *testing.T) {
+		atts := getAtts([]testData{
+			{2, bitfield.Bitlist{0b11100000, 0b1}},
+			{4, bitfield.Bitlist{0b11100000, 0b1}},
+			{1, bitfield.Bitlist{0b11000000, 0b1}},
+			{4, bitfield.Bitlist{0b11110000, 0b1}},
+			{1, bitfield.Bitlist{0b11100000, 0b1}},
+			{3, bitfield.Bitlist{0b11000000, 0b1}},
+		})
+		want := getAtts([]testData{
+			{4, bitfield.Bitlist{0b11110000, 0b1}},
+			{4, bitfield.Bitlist{0b11100000, 0b1}},
+			{3, bitfield.Bitlist{0b11000000, 0b1}},
+			{2, bitfield.Bitlist{0b11100000, 0b1}},
+			{1, bitfield.Bitlist{0b11100000, 0b1}},
+			{1, bitfield.Bitlist{0b11000000, 0b1}},
+		})
+		atts = atts.sortByProfitability()
+		require.DeepEqual(t, want, atts)
+	})
+
+	t.Run("selected and non selected atts sorted by bit count", func(t *testing.T) {
+		// Items at slot 4, must be first split into two lists by max-cover, with
+		// 0b10000011 scoring higher (as it provides more info in addition to already selected
+		// attestations) than 0b11100001 (despite naive bit count suggesting otherwise). Then,
+		// both selected and non-selected attestations must be additionally sorted by bit count.
+		atts := getAtts([]testData{
+			{4, bitfield.Bitlist{0b00000001, 0b1}},
+			{4, bitfield.Bitlist{0b11100001, 0b1}},
+			{1, bitfield.Bitlist{0b11000000, 0b1}},
+			{2, bitfield.Bitlist{0b11100000, 0b1}},
+			{4, bitfield.Bitlist{0b10000011, 0b1}},
+			{4, bitfield.Bitlist{0b11111000, 0b1}},
+			{1, bitfield.Bitlist{0b11100000, 0b1}},
+			{3, bitfield.Bitlist{0b11000000, 0b1}},
+		})
+		want := getAtts([]testData{
+			{4, bitfield.Bitlist{0b11111000, 0b1}},
+			{4, bitfield.Bitlist{0b10000011, 0b1}},
+			{4, bitfield.Bitlist{0b11100001, 0b1}},
+			{4, bitfield.Bitlist{0b00000001, 0b1}},
+			{3, bitfield.Bitlist{0b11000000, 0b1}},
+			{2, bitfield.Bitlist{0b11100000, 0b1}},
+			{1, bitfield.Bitlist{0b11100000, 0b1}},
+			{1, bitfield.Bitlist{0b11000000, 0b1}},
+		})
+		atts = atts.sortByProfitability()
+		require.DeepEqual(t, want, atts)
+	})
+}
+
 func TestProposer_ProposerAtts_dedup(t *testing.T) {
 	data1 := testutil.HydrateAttestationData(&ethpb.AttestationData{
 		Slot: 4,

shared/aggregation/testing/bitlistutils.go

@@ -81,8 +81,11 @@ func MakeAttestationsFromBitlists(bl []bitfield.Bitlist) []*ethpb.Attestation {
 	for i, b := range bl {
 		atts[i] = &ethpb.Attestation{
 			AggregationBits: b,
-			Data:            nil,
-			Signature:       bls.NewAggregateSignature().Marshal(),
+			Data: &ethpb.AttestationData{
+				Slot:           42,
+				CommitteeIndex: 1,
+			},
+			Signature: bls.NewAggregateSignature().Marshal(),
 		}
 	}
 	return atts