Attestation aggregation: optimizations and benchmarks

shared/aggregation/attestations/attestations.go

@@ -16,6 +16,11 @@ const (
 
 	// MaxCoverAggregation is a strategy based on Maximum Coverage greedy algorithm.
 	MaxCoverAggregation AttestationAggregationStrategy = "max_cover"
+
+	// OptMaxCoverAggregation is a strategy based on Maximum Coverage greedy algorithm.
+	// This new variant is optimized and relies on Bitlist64 (once fully tested, `max_cover`
+	// strategy will be replaced with this one).
+	OptMaxCoverAggregation AttestationAggregationStrategy = "opt_max_cover"
 )
 
 // AttestationAggregationStrategy defines attestation aggregation strategy.
@@ -37,13 +42,23 @@ var _ = logrus.WithField("prefix", "aggregation.attestations")
 var ErrInvalidAttestationCount = errors.New("invalid number of attestations")
 
 // Aggregate aggregates attestations. The minimal number of attestations is returned.
+// Aggregation occurs in-place i.e. contents of input array will be modified. Should you need to
+// preserve input attestations, clone them before aggregating:
+//
+//   clonedAtts := make([]*ethpb.Attestation, len(atts))
+//   for i, a := range atts {
+//       clonedAtts[i] = stateTrie.CopyAttestation(a)
+//   }
+//   aggregatedAtts, err := attaggregation.Aggregate(clonedAtts)
 func Aggregate(atts []*ethpb.Attestation) ([]*ethpb.Attestation, error) {
 	strategy := AttestationAggregationStrategy(featureconfig.Get().AttestationAggregationStrategy)
 	switch strategy {
 	case "", NaiveAggregation:
 		return NaiveAttestationAggregation(atts)
 	case MaxCoverAggregation:
 		return MaxCoverAttestationAggregation(atts)
+	case OptMaxCoverAggregation:
+		return optMaxCoverAttestationAggregation(atts)
 	default:
 		return nil, errors.Wrapf(aggregation.ErrInvalidStrategy, "%q", strategy)
 	}

shared/aggregation/attestations/attestations_bench_test.go

@@ -1,9 +1,12 @@
 package attestations
 
 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/bls"
 	"github.com/prysmaticlabs/prysm/shared/bls/common"
@@ -33,75 +36,68 @@ func BenchmarkAggregateAttestations_Aggregate(b *testing.B) {
 		inputs []bitfield.Bitlist
 	}{
 		{
-			name:   "64 attestations with single bit set",
-			inputs: aggtesting.BitlistsWithSingleBitSet(64, bitlistLen),
-		},
-		{
-			name:   "64 attestations with 8 random bits set",
-			inputs: aggtesting.BitlistsWithMultipleBitSet(b, 64, bitlistLen, 8),
-		},
-		{
-			name:   "64 attestations with 16 random bits set",
-			inputs: aggtesting.BitlistsWithMultipleBitSet(b, 64, bitlistLen, 16),
-		},
-		{
-			name:   "64 attestations with 32 random bits set",
-			inputs: aggtesting.BitlistsWithMultipleBitSet(b, 64, bitlistLen, 32),
-		},
-		{
-			name:   "256 attestations with 32 random bits set",
-			inputs: aggtesting.BitlistsWithMultipleBitSet(b, 256, bitlistLen, 32),
+			name:   "256 attestations with single bit set",
+			inputs: aggtesting.BitlistsWithSingleBitSet(256, bitlistLen),
 		},
 		{
 			name:   "256 attestations with 64 random bits set",
-			inputs: aggtesting.BitlistsWithMultipleBitSet(b, 256, bitlistLen, 64),
-		},
-		{
-			name:   "128 attestations with single bit set",
-			inputs: aggtesting.BitlistsWithSingleBitSet(128, bitlistLen),
-		},
-		{
-			name:   "256 attestations with single bit set",
 			inputs: aggtesting.BitlistsWithSingleBitSet(256, bitlistLen),
 		},
 		{
 			name:   "512 attestations with single bit set",
 			inputs: aggtesting.BitlistsWithSingleBitSet(512, bitlistLen),
 		},
 		{
-			name:   "1024 attestations with single bit set",
-			inputs: aggtesting.BitlistsWithSingleBitSet(1024, bitlistLen),
+			name:   "1024 attestations with 64 random bits set",
+			inputs: aggtesting.BitlistsWithMultipleBitSet(b, 1024, bitlistLen, 64),
 		},
 	}
 
-	b.Run("max-cover", func(b *testing.B) {
-		for _, tt := range tests {
-			b.Run(tt.name, func(b *testing.B) {
-				atts := aggtesting.MakeAttestationsFromBitlists(tt.inputs)
-				b.ResetTimer()
-				for i := 0; i < b.N; i++ {
-					_, err := Aggregate(atts)
-					require.NoError(b, err)
-				}
-			})
+	runner := func(atts []*ethpb.Attestation) {
+		attsCopy := make([]*ethpb.Attestation, len(atts))
+		for i, att := range atts {
+			attsCopy[i] = stateTrie.CopyAttestation(att)
 		}
-	})
+		_, err := Aggregate(attsCopy)
+		require.NoError(b, err)
+	}
 
-	b.Run("naive", func(b *testing.B) {
-		resetCfg := featureconfig.InitWithReset(&featureconfig.Flags{
-			AttestationAggregationStrategy: string(NaiveAggregation),
+	for _, tt := range tests {
+		b.Run(fmt.Sprintf("naive_%s", tt.name), func(b *testing.B) {
+			b.StopTimer()
+			resetCfg := featureconfig.InitWithReset(&featureconfig.Flags{
+				AttestationAggregationStrategy: string(NaiveAggregation),
+			})
+			atts := aggtesting.MakeAttestationsFromBitlists(tt.inputs)
+			defer resetCfg()
+			b.StartTimer()
+			for i := 0; i < b.N; i++ {
+				runner(atts)
+			}
 		})
-		defer resetCfg()
-
-		for _, tt := range tests {
-			b.Run(tt.name, func(b *testing.B) {
-				atts := aggtesting.MakeAttestationsFromBitlists(tt.inputs)
-				b.ResetTimer()
-				for i := 0; i < b.N; i++ {
-					_, err := Aggregate(atts)
-					require.NoError(b, err)
-				}
+		b.Run(fmt.Sprintf("max-cover_%s", tt.name), func(b *testing.B) {
+			b.StopTimer()
+			resetCfg := featureconfig.InitWithReset(&featureconfig.Flags{
+				AttestationAggregationStrategy: string(MaxCoverAggregation),
 			})
-		}
-	})
+			atts := aggtesting.MakeAttestationsFromBitlists(tt.inputs)
+			defer resetCfg()
+			b.StartTimer()
+			for i := 0; i < b.N; i++ {
+				runner(atts)
+			}
+		})
+		b.Run(fmt.Sprintf("opt-max-cover_%s", tt.name), func(b *testing.B) {
+			b.StopTimer()
+			resetCfg := featureconfig.InitWithReset(&featureconfig.Flags{
+				AttestationAggregationStrategy: string(OptMaxCoverAggregation),
+			})
+			atts := aggtesting.MakeAttestationsFromBitlists(tt.inputs)
+			defer resetCfg()
+			b.StartTimer()
+			for i := 0; i < b.N; i++ {
+				runner(atts)
+			}
+		})
+	}
 }

shared/aggregation/attestations/attestations_test.go

@@ -23,7 +23,7 @@ func TestMain(m *testing.M) {
 	logrus.SetLevel(logrus.DebugLevel)
 	logrus.SetOutput(ioutil.Discard)
 	resetCfg := featureconfig.InitWithReset(&featureconfig.Flags{
-		AttestationAggregationStrategy: string(MaxCoverAggregation),
+		AttestationAggregationStrategy: string(OptMaxCoverAggregation),
 	})
 	defer resetCfg()
 	m.Run()
@@ -240,6 +240,13 @@ func TestAggregateAttestations_Aggregate(t *testing.T) {
 			defer resetCfg()
 			runner()
 		})
+		t.Run(fmt.Sprintf("%s/%s", tt.name, OptMaxCoverAggregation), func(t *testing.T) {
+			resetCfg := featureconfig.InitWithReset(&featureconfig.Flags{
+				AttestationAggregationStrategy: string(OptMaxCoverAggregation),
+			})
+			defer resetCfg()
+			runner()
+		})
 	}
 }
 

shared/aggregation/attestations/maxcover.go

@@ -62,6 +62,96 @@ func MaxCoverAttestationAggregation(atts []*ethpb.Attestation) ([]*ethpb.Attesta
 	return aggregated.merge(unaggregated.filterContained()), nil
 }
 
+// optMaxCoverAttestationAggregation relies on Maximum Coverage greedy algorithm for aggregation.
+// Aggregation occurs in many rounds, up until no more aggregation is possible (all attestations
+// are overlapping).
+// NB: this method will replace the MaxCoverAttestationAggregation() above (and will be renamed to it).
+// See https://hackmd.io/@farazdagi/in-place-attagg for design and rationale.
+func optMaxCoverAttestationAggregation(atts []*ethpb.Attestation) ([]*ethpb.Attestation, error) {
+	if len(atts) < 2 {
+		return atts, nil
+	}
+
+	if err := attList(atts).validate(); err != nil {
+		if errors.Is(err, aggregation.ErrBitsDifferentLen) {
+			return atts, nil
+		}
+		return nil, err
+	}
+
+	// In the future this conversion will be redundant, as attestation bitlist will be of a Bitlist64
+	// type, so incoming `atts` parameters can be used as candidates list directly.
+	candidates := make([]*bitfield.Bitlist64, len(atts))
+	for i := 0; i < len(atts); i++ {
+		candidates[i] = atts[i].AggregationBits.ToBitlist64()
+	}
+	coveredBitsSoFar := bitfield.NewBitlist64(candidates[0].Len())
+
+	// In order not to re-allocate anything we rely on the very same underlying array, which
+	// can only shrink (while the `aggregated` slice length can increase).
+	// The `aggregated` slice grows by combining individual attestations and appending to that slice.
+	// Both aggregated and non-aggregated slices operate on the very same underlying array.
+	aggregated := atts[:0]
+	unaggregated := atts
+
+	// Aggregation over n/2 rounds is enough to find all aggregatable items (exits earlier if there
+	// are many items that can be aggregated).
+	for i := 0; i < len(atts)/2; i++ {
+		if len(unaggregated) < 2 {
+			break
+		}
+
+		// Find maximum non-overlapping coverage for subset of still non-processed candidates.
+		roundCandidates := candidates[len(aggregated) : len(aggregated)+len(unaggregated)]
+		selectedKeys, coverage, err := aggregation.MaxCover(
+			roundCandidates, len(roundCandidates), false /* allowOverlaps */)
+		if err != nil {
+			// Return aggregated attestations, and attestations that couldn't be aggregated.
+			return append(aggregated, unaggregated...), err
+		}
+
+		// Exit earlier, if possible cover does not allow aggregation (less than two items).
+		if selectedKeys.Count() < 2 {
+			break
+		}
+
+		// Pad selected key indexes, as `roundCandidates` is a subset of `candidates`.
+		keys := padSelectedKeys(selectedKeys.BitIndices(), len(aggregated))
+
+		// Create aggregated attestation and update solution lists. Process aggregates only if they
+		// feature at least one unknown bit i.e. can increase the overall coverage.
+		if coveredBitsSoFar.XorCount(coverage) > 0 {
+			aggIdx, err := aggregateAttestations(atts, keys, coverage)
+			if err != nil {
+				return append(aggregated, unaggregated...), err
+			}
+
+			// Unless we are already at the right position, swap aggregation and the first non-aggregated item.
+			idx0 := len(aggregated)
+			if idx0 < aggIdx {
+				atts[idx0], atts[aggIdx] = atts[aggIdx], atts[idx0]
+				candidates[idx0], candidates[aggIdx] = candidates[aggIdx], candidates[idx0]
+			}
+
+			// Expand to the newly created aggregate.
+			aggregated = atts[:idx0+1]
+
+			// Shift the starting point of the slice to the right.
+			unaggregated = unaggregated[1:]
+
+			// Update covered bits map.
+			coveredBitsSoFar.NoAllocOr(coverage, coveredBitsSoFar)
+			keys = keys[1:]
+		}
+
+		// Remove processed attestations.
+		rearrangeProcessedAttestations(atts, candidates, keys)
+		unaggregated = unaggregated[:len(unaggregated)-len(keys)]
+	}
+
+	return append(aggregated, attList(unaggregated).filterContained()...), nil
+}
+
 // NewMaxCover returns initialized Maximum Coverage problem for attestations aggregation.
 func NewMaxCover(atts []*ethpb.Attestation) *aggregation.MaxCoverProblem {
 	candidates := make([]*aggregation.MaxCoverCandidate, len(atts))
@@ -91,6 +181,75 @@ func (al attList) aggregate(coverage bitfield.Bitlist) (*ethpb.Attestation, erro
 	}, nil
 }
 
+// padSelectedKeys adds additional value to every key.
+func padSelectedKeys(keys []int, pad int) []int {
+	for i, key := range keys {
+		keys[i] = key + pad
+	}
+	return keys
+}
+
+// aggregateAttestations combines signatures of selected attestations into a single aggregate attestation, and
+// pushes that aggregated attestation into the position of the first of selected attestations.
+func aggregateAttestations(atts []*ethpb.Attestation, keys []int, coverage *bitfield.Bitlist64) (targetIdx int, err error) {
+	if len(keys) < 2 || atts == nil || len(atts) < 2 {
+		return targetIdx, errors.Wrap(ErrInvalidAttestationCount, "cannot aggregate")
+	}
+	if coverage == nil || coverage.Count() == 0 {
+		return targetIdx, errors.New("invalid or empty coverage")
+	}
+
+	var data *ethpb.AttestationData
+	signs := make([]bls.Signature, 0, len(keys))
+	for i, idx := range keys {
+		sig, err := signatureFromBytes(atts[idx].Signature)
+		if err != nil {
+			return targetIdx, err
+		}
+		signs = append(signs, sig)
+		if i == 0 {
+			data = stateTrie.CopyAttestationData(atts[idx].Data)
+			targetIdx = idx
+		}
+	}
+	// Put aggregated attestation at a position of the first selected attestation.
+	atts[targetIdx] = &ethpb.Attestation{
+		// Append size byte, which will be unnecessary on switch to Bitlist64.
+		AggregationBits: coverage.ToBitlist(),
+		Data:            data,
+		Signature:       aggregateSignatures(signs).Marshal(),
+	}
+	return
+}
+
+// rearrangeProcessedAttestations pushes processed attestations to the end of the slice, returning
+// the number of items re-arranged (so that caller can cut the slice, and allow processed items to be
+// garbage collected).
+func rearrangeProcessedAttestations(atts []*ethpb.Attestation, candidates []*bitfield.Bitlist64, processedKeys []int) {
+	if atts == nil || candidates == nil || processedKeys == nil {
+		return
+	}
+	// Set all selected keys to nil.
+	for _, idx := range processedKeys {
+		atts[idx] = nil
+		candidates[idx] = nil
+	}
+	// Re-arrange nil items, move them to end of slice.
+	sort.Ints(processedKeys)
+	lastIdx := len(atts) - 1
+	for _, idx0 := range processedKeys {
+		// Make sure that nil items are swapped for non-nil items only.
+		for lastIdx > idx0 && atts[lastIdx] == nil {
+			lastIdx--
+		}
+		if idx0 == lastIdx {
+			break
+		}
+		atts[idx0], atts[lastIdx] = atts[lastIdx], atts[idx0]
+		candidates[idx0], candidates[lastIdx] = candidates[lastIdx], candidates[idx0]
+	}
+}
+
 // merge combines two attestation lists into one.
 func (al attList) merge(al1 attList) attList {
 	return append(al, al1...)