Sketch op pool changes

beacon_node/operation_pool/src/attestation_storage.rs

@@ -1,6 +1,6 @@
 use crate::AttestationStats;
 use itertools::Itertools;
-use std::collections::HashMap;
+use std::collections::{BTreeMap, HashMap};
 use types::{
     attestation::{AttestationBase, AttestationElectra},
     superstruct, AggregateSignature, Attestation, AttestationData, BeaconState, BitList, BitVector,
@@ -42,6 +42,7 @@ pub struct SplitAttestation<E: EthSpec> {
     pub indexed: CompactIndexedAttestation<E>,
 }
 
+// TODO(electra): rename this type
 #[derive(Debug, Clone)]
 pub struct AttestationRef<'a, E: EthSpec> {
     pub checkpoint: &'a CheckpointKey,
@@ -159,15 +160,15 @@ impl CheckpointKey {
 }
 
 impl<E: EthSpec> CompactIndexedAttestation<E> {
-    pub fn signers_disjoint_from(&self, other: &Self) -> bool {
+    pub fn should_aggregate(&self, other: &Self) -> bool {
         match (self, other) {
             (CompactIndexedAttestation::Base(this), CompactIndexedAttestation::Base(other)) => {
-                this.signers_disjoint_from(other)
+                this.should_aggregate(other)
             }
             (
                 CompactIndexedAttestation::Electra(this),
                 CompactIndexedAttestation::Electra(other),
-            ) => this.signers_disjoint_from(other),
+            ) => this.should_aggregate(other),
             // TODO(electra) is a mix of electra and base compact indexed attestations an edge case we need to deal with?
             _ => false,
         }
@@ -189,7 +190,7 @@ impl<E: EthSpec> CompactIndexedAttestation<E> {
 }
 
 impl<E: EthSpec> CompactIndexedAttestationBase<E> {
-    pub fn signers_disjoint_from(&self, other: &Self) -> bool {
+    pub fn should_aggregate(&self, other: &Self) -> bool {
         self.aggregation_bits
             .intersection(&other.aggregation_bits)
             .is_zero()
@@ -208,14 +209,15 @@ impl<E: EthSpec> CompactIndexedAttestationBase<E> {
 }
 
 impl<E: EthSpec> CompactIndexedAttestationElectra<E> {
-    // TODO(electra) update to match spec requirements
-    pub fn signers_disjoint_from(&self, other: &Self) -> bool {
-        self.aggregation_bits
-            .intersection(&other.aggregation_bits)
-            .is_zero()
+    pub fn should_aggregate(&self, other: &Self) -> bool {
+        // For Electra, only aggregate attestations in the same committee.
+        self.committee_bits == other.committee_bits
+            && self
+                .aggregation_bits
+                .intersection(&other.aggregation_bits)
+                .is_zero()
     }
 
-    // TODO(electra) update to match spec requirements
     pub fn aggregate(&mut self, other: &Self) {
         self.attesting_indices = self
             .attesting_indices
@@ -226,6 +228,18 @@ impl<E: EthSpec> CompactIndexedAttestationElectra<E> {
         self.aggregation_bits = self.aggregation_bits.union(&other.aggregation_bits);
         self.signature.add_assign_aggregate(&other.signature);
     }
+
+    pub fn committee_index(&self) -> u64 {
+        *self.get_committee_indices().first().unwrap_or(&0u64)
+    }
+
+    pub fn get_committee_indices(&self) -> Vec<u64> {
+        self.committee_bits
+            .iter()
+            .enumerate()
+            .filter_map(|(index, bit)| if bit { Some(index as u64) } else { None })
+            .collect()
+    }
 }
 
 impl<E: EthSpec> AttestationMap<E> {
@@ -239,34 +253,63 @@ impl<E: EthSpec> AttestationMap<E> {
         let attestation_map = self.checkpoint_map.entry(checkpoint).or_default();
         let attestations = attestation_map.attestations.entry(data).or_default();
 
-        // TODO(electra):
         // Greedily aggregate the attestation with all existing attestations.
         // NOTE: this is sub-optimal and in future we will remove this in favour of max-clique
         // aggregation.
         let mut aggregated = false;
 
-        match attestation {
-            Attestation::Base(_) => {
-                for existing_attestation in attestations.iter_mut() {
-                    if existing_attestation.signers_disjoint_from(&indexed) {
-                        existing_attestation.aggregate(&indexed);
-                        aggregated = true;
-                    } else if *existing_attestation == indexed {
-                        aggregated = true;
-                    }
-                }
+        for existing_attestation in attestations.iter_mut() {
+            if existing_attestation.should_aggregate(&indexed) {
+                existing_attestation.aggregate(&indexed);
+                aggregated = true;
+            } else if *existing_attestation == indexed {
+                aggregated = true;
             }
-            // TODO(electra) in order to be devnet ready, we can skip
-            // aggregating here for now. this will result in "poorly"
-            // constructed blocks, but that should be fine for devnet
-            Attestation::Electra(_) => (),
-        };
+        }
 
         if !aggregated {
             attestations.push(indexed);
         }
     }
 
+    pub fn aggregate_across_committees(&mut self, checkpoint_key: CheckpointKey) {
+        let Some(attestation_map) = self.checkpoint_map.get_mut(&checkpoint_key) else {
+            return;
+        };
+        for (compact_attestation_data, compact_indexed_attestations) in
+            attestation_map.attestations.iter_mut()
+        {
+            let unaggregated_attestations = std::mem::take(compact_indexed_attestations);
+            let mut aggregated_attestations = vec![];
+
+            // Aggregate the best attestations for each committee and leave the rest.
+            let mut best_attestations_by_committee = BTreeMap::new();
+
+            for committee_attestation in unaggregated_attestations {
+                // TODO(electra)
+                // compare to best attestations by committee
+                // could probably use `.entry` here
+                if let Some(existing_attestation) =
+                    best_attestations_by_committee.get_mut(committee_attestation.committee_index())
+                {
+                    // compare and swap, put the discarded one straight into
+                    // `aggregated_attestations` in case we have room to pack it without
+                    // cross-committee aggregation
+                } else {
+                    best_attestations_by_committee.insert(
+                        committee_attestation.committee_index(),
+                        committee_attestation,
+                    );
+                }
+            }
+
+            // TODO(electra): aggregate all the best attestations by committee
+            // (use btreemap sort order to get order by committee index)
+
+            *compact_indexed_attestations = aggregated_attestations;
+        }
+    }
+
     /// Iterate all attestations matching the given `checkpoint_key`.
     pub fn get_attestations<'a>(
         &'a self,

beacon_node/operation_pool/src/lib.rs

@@ -40,7 +40,7 @@ use std::ptr;
 use types::{
     sync_aggregate::Error as SyncAggregateError, typenum::Unsigned, AbstractExecPayload,
     Attestation, AttestationData, AttesterSlashing, BeaconState, BeaconStateError, ChainSpec,
-    Epoch, EthSpec, ProposerSlashing, SignedBeaconBlock, SignedBlsToExecutionChange,
+    Epoch, EthSpec, ForkName, ProposerSlashing, SignedBeaconBlock, SignedBlsToExecutionChange,
     SignedVoluntaryExit, Slot, SyncAggregate, SyncCommitteeContribution, Validator,
 };
 
@@ -256,6 +256,7 @@ impl<E: EthSpec> OperationPool<E> {
         curr_epoch_validity_filter: impl for<'a> FnMut(&AttestationRef<'a, E>) -> bool + Send,
         spec: &ChainSpec,
     ) -> Result<Vec<Attestation<E>>, OpPoolError> {
+        let fork_name = state.fork_name_unchecked();
         if !matches!(state, BeaconState::Base(_)) {
             // Epoch cache must be initialized to fetch base reward values in the max cover `score`
             // function. Currently max cover ignores items on errors. If epoch cache is not
@@ -267,7 +268,6 @@ impl<E: EthSpec> OperationPool<E> {
 
         // Attestations for the current fork, which may be from the current or previous epoch.
         let (prev_epoch_key, curr_epoch_key) = CheckpointKey::keys_for_state(state);
-        let all_attestations = self.attestations.read();
         let total_active_balance = state
             .get_total_active_balance()
             .map_err(OpPoolError::GetAttestationsTotalBalanceError)?;
@@ -284,6 +284,14 @@ impl<E: EthSpec> OperationPool<E> {
         let mut num_prev_valid = 0_i64;
         let mut num_curr_valid = 0_i64;
 
+        // TODO(electra): Work out how to do this more elegantly. This is a bit of a hack.
+        let mut all_attestations = self.attestations.write();
+
+        all_attestations.aggregate_across_committees(prev_epoch_key);
+        all_attestations.aggregate_across_committees(curr_epoch_key);
+
+        let all_attestations = parking_lot::RwLockWriteGuard::downgrade(all_attestations);
+
         let prev_epoch_att = self
             .get_valid_attestations_for_epoch(
                 &prev_epoch_key,
@@ -307,6 +315,11 @@ impl<E: EthSpec> OperationPool<E> {
             )
             .inspect(|_| num_curr_valid += 1);
 
+        let curr_epoch_limit = if fork_name < ForkName::Electra {
+            E::MaxAttestations::to_usize()
+        } else {
+            E::MaxAttestationsElectra::to_usize()
+        };
         let prev_epoch_limit = if let BeaconState::Base(base_state) = state {
             std::cmp::min(
                 E::MaxPendingAttestations::to_usize()
@@ -314,7 +327,7 @@ impl<E: EthSpec> OperationPool<E> {
                 E::MaxAttestations::to_usize(),
             )
         } else {
-            E::MaxAttestations::to_usize()
+            curr_epoch_limit
         };
 
         let (prev_cover, curr_cover) = rayon::join(
@@ -329,11 +342,7 @@ impl<E: EthSpec> OperationPool<E> {
             },
             move || {
                 let _timer = metrics::start_timer(&metrics::ATTESTATION_CURR_EPOCH_PACKING_TIME);
-                maximum_cover(
-                    curr_epoch_att,
-                    E::MaxAttestations::to_usize(),
-                    "curr_epoch_attestations",
-                )
+                maximum_cover(curr_epoch_att, curr_epoch_limit, "curr_epoch_attestations")
             },
         );
 
@@ -343,7 +352,7 @@ impl<E: EthSpec> OperationPool<E> {
         Ok(max_cover::merge_solutions(
             curr_cover,
             prev_cover,
-            E::MaxAttestations::to_usize(),
+            curr_epoch_limit,
         ))
     }
 

consensus/types/src/attestation.rs

@@ -247,6 +247,9 @@ impl<'a, E: EthSpec> AttestationRef<'a, E> {
 
 impl<E: EthSpec> AttestationElectra<E> {
     /// Are the aggregation bitfields of these attestations disjoint?
+    // TODO(electra): check whether the definition from CompactIndexedAttestation::should_aggregate
+    // is useful where this is used, i.e. only consider attestations disjoint when their committees
+    // match AND their aggregation bits do not intersect.
     pub fn signers_disjoint_from(&self, other: &Self) -> bool {
         self.aggregation_bits
             .intersection(&other.aggregation_bits)