diff --git a/beacon_node/operation_pool/src/attestation_storage.rs b/beacon_node/operation_pool/src/attestation_storage.rs index 5cf93e642c9..6e04af01b9c 100644 --- a/beacon_node/operation_pool/src/attestation_storage.rs +++ b/beacon_node/operation_pool/src/attestation_storage.rs @@ -187,6 +187,13 @@ impl CompactIndexedAttestation { _ => (), } } + + pub fn committee_index(&self) -> u64 { + match self { + CompactIndexedAttestation::Base(att) => att.index, + CompactIndexedAttestation::Electra(att) => att.committee_index(), + } + } } impl CompactIndexedAttestationBase { @@ -276,25 +283,34 @@ impl AttestationMap { 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() - { + for (_, compact_indexed_attestations) in attestation_map.attestations.iter_mut() { let unaggregated_attestations = std::mem::take(compact_indexed_attestations); - let mut aggregated_attestations = vec![]; + let mut aggregated_attestations: Vec> = vec![]; // Aggregate the best attestations for each committee and leave the rest. - let mut best_attestations_by_committee = BTreeMap::new(); + let mut best_attestations_by_committee: BTreeMap> = + 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()) + 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 + if existing_attestation.should_aggregate(&committee_attestation) { + existing_attestation.aggregate(&committee_attestation); + + best_attestations_by_committee.insert( + committee_attestation.committee_index(), + committee_attestation, + ); + } else { + aggregated_attestations.push(committee_attestation); + } } else { best_attestations_by_committee.insert( committee_attestation.committee_index(), @@ -305,11 +321,62 @@ impl AttestationMap { // TODO(electra): aggregate all the best attestations by committee // (use btreemap sort order to get order by committee index) + aggregated_attestations.extend(Self::compute_on_chain_aggregate( + best_attestations_by_committee, + )); *compact_indexed_attestations = aggregated_attestations; } } + // TODO(electra) unwraps in this function should be cleaned up + // also in general this could be a bit more elegant + pub fn compute_on_chain_aggregate( + mut attestations_by_committee: BTreeMap>, + ) -> Vec> { + let mut aggregated_attestations = vec![]; + if let Some((_, on_chain_aggregate)) = attestations_by_committee.pop_first() { + match on_chain_aggregate { + CompactIndexedAttestation::Base(a) => { + aggregated_attestations.push(CompactIndexedAttestation::Base(a)); + aggregated_attestations.extend( + attestations_by_committee + .values() + .map(|a| { + CompactIndexedAttestation::Base(CompactIndexedAttestationBase { + attesting_indices: a.attesting_indices().clone(), + aggregation_bits: a.aggregation_bits_base().unwrap().clone(), + signature: a.signature().clone(), + index: *a.index(), + }) + }) + .collect::>>(), + ); + } + CompactIndexedAttestation::Electra(mut a) => { + for (_, attestation) in attestations_by_committee.iter_mut() { + let new_committee_bits = a + .committee_bits + .union(attestation.committee_bits().unwrap()); + a.aggregate(attestation.as_electra().unwrap()); + + a = CompactIndexedAttestationElectra { + attesting_indices: a.attesting_indices.clone(), + aggregation_bits: a.aggregation_bits.clone(), + signature: a.signature.clone(), + index: a.index, + committee_bits: new_committee_bits, + }; + } + + aggregated_attestations.push(CompactIndexedAttestation::Electra(a)); + } + } + } + + aggregated_attestations + } + /// Iterate all attestations matching the given `checkpoint_key`. pub fn get_attestations<'a>( &'a self, diff --git a/beacon_node/operation_pool/src/lib.rs b/beacon_node/operation_pool/src/lib.rs index 1d83ecd4651..4f247e6bf20 100644 --- a/beacon_node/operation_pool/src/lib.rs +++ b/beacon_node/operation_pool/src/lib.rs @@ -1246,7 +1246,17 @@ mod release_tests { let num_big = target_committee_size / big_step_size; let stats = op_pool.attestation_stats(); - assert_eq!(stats.num_attestation_data, committees.len()); + let fork_name = state.fork_name_unchecked(); + + match fork_name { + ForkName::Electra => { + assert_eq!(stats.num_attestation_data, 1); + } + _ => { + assert_eq!(stats.num_attestation_data, committees.len()); + } + }; + assert_eq!( stats.num_attestations, (num_small + num_big) * committees.len() @@ -1257,11 +1267,27 @@ mod release_tests { let best_attestations = op_pool .get_attestations(&state, |_| true, |_| true, spec) .expect("should have best attestations"); - assert_eq!(best_attestations.len(), max_attestations); + match fork_name { + ForkName::Electra => { + assert_eq!(best_attestations.len(), 8); + } + _ => { + assert_eq!(best_attestations.len(), max_attestations); + } + }; // All the best attestations should be signed by at least `big_step_size` (4) validators. for att in &best_attestations { - assert!(att.num_set_aggregation_bits() >= big_step_size); + match fork_name { + ForkName::Electra => { + // TODO(electra) some attestations only have 2 or 3 agg bits set + // others have 5 + assert!(att.num_set_aggregation_bits() >= 2); + } + _ => { + assert!(att.num_set_aggregation_bits() >= big_step_size); + } + }; } } @@ -1340,11 +1366,20 @@ mod release_tests { let num_small = target_committee_size / small_step_size; let num_big = target_committee_size / big_step_size; + let fork_name = state.fork_name_unchecked(); + + match fork_name { + ForkName::Electra => { + assert_eq!(op_pool.attestation_stats().num_attestation_data, 1); + } + _ => { + assert_eq!( + op_pool.attestation_stats().num_attestation_data, + committees.len() + ); + } + }; - assert_eq!( - op_pool.attestation_stats().num_attestation_data, - committees.len() - ); assert_eq!( op_pool.num_attestations(), (num_small + num_big) * committees.len() @@ -1355,7 +1390,15 @@ mod release_tests { let best_attestations = op_pool .get_attestations(&state, |_| true, |_| true, spec) .expect("should have valid best attestations"); - assert_eq!(best_attestations.len(), max_attestations); + + match fork_name { + ForkName::Electra => { + assert_eq!(best_attestations.len(), 8); + } + _ => { + assert_eq!(best_attestations.len(), max_attestations); + } + }; let total_active_balance = state.get_total_active_balance().unwrap();