Skip to content

Commit

Permalink
Sketch op pool changes
Browse files Browse the repository at this point in the history
  • Loading branch information
michaelsproul authored and realbigsean committed Jun 25, 2024
1 parent 33aedf5 commit 6b1d05d
Show file tree
Hide file tree
Showing 3 changed files with 61 additions and 99 deletions.
143 changes: 52 additions & 91 deletions beacon_node/operation_pool/src/attestation_storage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ pub struct SplitAttestation<E: EthSpec> {
pub indexed: CompactIndexedAttestation<E>,
}

// TODO(electra): rename this type
#[derive(Debug, Clone)]
pub struct CompactAttestationRef<'a, E: EthSpec> {
pub checkpoint: &'a CheckpointKey,
Expand Down Expand Up @@ -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,
}
Expand All @@ -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()
Expand All @@ -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
Expand All @@ -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> {
Expand All @@ -239,116 +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);
}
}

/// Aggregate Electra attestations for the same attestation data signed by different
/// committees.
///
/// Non-Electra attestations are left as-is.
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_indexed_attestations in attestation_map.attestations.values_mut() {
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<CompactIndexedAttestation<E>> = vec![];
let mut aggregated_attestations = vec![];

// Aggregate the best attestations for each committee and leave the rest.
let mut best_attestations_by_committee: BTreeMap<
u64,
CompactIndexedAttestationElectra<E>,
> = BTreeMap::new();
let mut best_attestations_by_committee = BTreeMap::new();

for committee_attestation in unaggregated_attestations {
let mut electra_attestation = match committee_attestation {
CompactIndexedAttestation::Electra(att)
if att.committee_bits.num_set_bits() == 1 =>
{
att
}
CompactIndexedAttestation::Electra(att) => {
// Aggregate already covers multiple committees, leave it as-is.
aggregated_attestations.push(CompactIndexedAttestation::Electra(att));
continue;
}
CompactIndexedAttestation::Base(att) => {
// Leave as-is.
aggregated_attestations.push(CompactIndexedAttestation::Base(att));
continue;
}
};
if let Some(committee_index) = electra_attestation.committee_index() {
if let Some(existing_attestation) =
best_attestations_by_committee.get_mut(&committee_index)
{
// Search for the best (most aggregation bits) attestation for this committee
// index.
if electra_attestation.aggregation_bits.num_set_bits()
> existing_attestation.aggregation_bits.num_set_bits()
{
// New attestation is better than the previously known one for this
// committee. Replace it.
std::mem::swap(existing_attestation, &mut electra_attestation);
}
// Put the inferior attestation into the list of aggregated attestations
// without performing any cross-committee aggregation.
aggregated_attestations
.push(CompactIndexedAttestation::Electra(electra_attestation));
} else {
// First attestation seen for this committee. Place it in the map
// provisionally.
best_attestations_by_committee.insert(committee_index, electra_attestation);
}
// 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,
);
}
}

if let Some(on_chain_aggregate) =
Self::compute_on_chain_aggregate(best_attestations_by_committee)
{
aggregated_attestations
.push(CompactIndexedAttestation::Electra(on_chain_aggregate));
}
// TODO(electra): aggregate all the best attestations by committee
// (use btreemap sort order to get order by committee index)

*compact_indexed_attestations = aggregated_attestations;
}
}

pub fn compute_on_chain_aggregate(
mut attestations_by_committee: BTreeMap<u64, CompactIndexedAttestationElectra<E>>,
) -> Option<CompactIndexedAttestationElectra<E>> {
let (_, mut on_chain_aggregate) = attestations_by_committee.pop_first()?;
for (_, attestation) in attestations_by_committee {
on_chain_aggregate.aggregate_with_disjoint_committees(&attestation);
}
Some(on_chain_aggregate)
}

/// Iterate all attestations matching the given `checkpoint_key`.
pub fn get_attestations<'a>(
&'a self,
Expand Down
14 changes: 6 additions & 8 deletions beacon_node/operation_pool/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,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,
};

Expand Down Expand Up @@ -286,10 +286,8 @@ impl<E: EthSpec> OperationPool<E> {
// TODO(electra): Work out how to do this more elegantly. This is a bit of a hack.
let mut all_attestations = self.attestations.write();

if fork_name.electra_enabled() {
all_attestations.aggregate_across_committees(prev_epoch_key);
all_attestations.aggregate_across_committees(curr_epoch_key);
}
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);

Expand All @@ -316,10 +314,10 @@ impl<E: EthSpec> OperationPool<E> {
)
.inspect(|_| num_curr_valid += 1);

let curr_epoch_limit = if fork_name.electra_enabled() {
E::MaxAttestationsElectra::to_usize()
} else {
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(
Expand Down
3 changes: 3 additions & 0 deletions consensus/types/src/attestation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -263,6 +263,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)
Expand Down

0 comments on commit 6b1d05d

Please sign in to comment.