Update commitment cache in alpenglow (anza-xyz#94)

Author: carllin

core/src/commitment_service.rs

@@ -21,7 +21,21 @@ use {
     },
 };
 
-pub struct CommitmentAggregationData {
+pub enum AlpenglowCommitmentType {
+    Notarize,
+    Root,
+}
+pub enum CommitmentAggregationData {
+    AlpenglowCommitmentAggregationData(AlpenglowCommitmentAggregationData),
+    TowerCommitmentAggregationData(TowerCommitmentAggregationData),
+}
+
+pub struct AlpenglowCommitmentAggregationData {
+    pub commitment_type: AlpenglowCommitmentType,
+    pub slot: Slot,
+}
+
+pub struct TowerCommitmentAggregationData {
     bank: Arc<Bank>,
     root: Slot,
     total_stake: Stake,
@@ -30,7 +44,7 @@ pub struct CommitmentAggregationData {
     node_vote_state: (Pubkey, TowerVoteState),
 }
 
-impl CommitmentAggregationData {
+impl TowerCommitmentAggregationData {
     pub fn new(
         bank: Arc<Bank>,
         root: Slot,
@@ -107,15 +121,28 @@ impl AggregateCommitmentService {
             let aggregation_data = receiver.recv_timeout(Duration::from_secs(1))?;
             let aggregation_data = receiver.try_iter().last().unwrap_or(aggregation_data);
 
-            let ancestors = aggregation_data.bank.status_cache_ancestors();
-            if ancestors.is_empty() {
-                continue;
-            }
-
             let mut aggregate_commitment_time = Measure::start("aggregate-commitment-ms");
-            let update_commitment_slots =
-                Self::update_commitment_cache(block_commitment_cache, aggregation_data, ancestors);
+            let update_commitment_slots = {
+                match aggregation_data {
+                    CommitmentAggregationData::TowerCommitmentAggregationData(data) => {
+                        let ancestors = data.bank.status_cache_ancestors();
+                        if ancestors.is_empty() {
+                            continue;
+                        }
+
+                        Self::update_commitment_cache(block_commitment_cache, data, ancestors)
+                    }
+                    CommitmentAggregationData::AlpenglowCommitmentAggregationData(data) => {
+                        Self::alpenglow_update_commitment_cache(
+                            block_commitment_cache,
+                            data.commitment_type,
+                            data.slot,
+                        )
+                    }
+                }
+            };
             aggregate_commitment_time.stop();
+
             datapoint_info!(
                 "block-commitment-cache",
                 (
@@ -134,17 +161,36 @@ impl AggregateCommitmentService {
                     i64
                 ),
             );
-
             // Triggers rpc_subscription notifications as soon as new commitment data is available,
             // sending just the commitment cache slot information that the notifications thread
             // needs
             subscriptions.notify_subscribers(update_commitment_slots);
         }
     }
 
+    fn alpenglow_update_commitment_cache(
+        block_commitment_cache: &RwLock<BlockCommitmentCache>,
+        update_type: AlpenglowCommitmentType,
+        slot: Slot,
+    ) -> CommitmentSlots {
+        let mut w_block_commitment_cache = block_commitment_cache.write().unwrap();
+
+        match update_type {
+            AlpenglowCommitmentType::Notarize => {
+                w_block_commitment_cache.set_slot(slot);
+            }
+            AlpenglowCommitmentType::Root => {
+                w_block_commitment_cache.set_highest_confirmed_slot(slot);
+                w_block_commitment_cache.set_root(slot);
+                w_block_commitment_cache.set_highest_super_majority_root(slot);
+            }
+        }
+        w_block_commitment_cache.commitment_slots()
+    }
+
     fn update_commitment_cache(
         block_commitment_cache: &RwLock<BlockCommitmentCache>,
-        aggregation_data: CommitmentAggregationData,
+        aggregation_data: TowerCommitmentAggregationData,
         ancestors: Vec<u64>,
     ) -> CommitmentSlots {
         let (block_commitment, rooted_stake) = Self::aggregate_commitment(
@@ -615,7 +661,7 @@ mod tests {
         let ancestors = working_bank.status_cache_ancestors();
         let _ = AggregateCommitmentService::update_commitment_cache(
             &block_commitment_cache,
-            CommitmentAggregationData {
+            TowerCommitmentAggregationData {
                 bank: working_bank,
                 root: 0,
                 total_stake: 100,
@@ -654,7 +700,7 @@ mod tests {
         let ancestors = working_bank.status_cache_ancestors();
         let _ = AggregateCommitmentService::update_commitment_cache(
             &block_commitment_cache,
-            CommitmentAggregationData {
+            TowerCommitmentAggregationData {
                 bank: working_bank,
                 root: 1,
                 total_stake: 100,
@@ -703,7 +749,7 @@ mod tests {
         let ancestors = working_bank.status_cache_ancestors();
         let _ = AggregateCommitmentService::update_commitment_cache(
             &block_commitment_cache,
-            CommitmentAggregationData {
+            TowerCommitmentAggregationData {
                 bank: working_bank,
                 root: 0,
                 total_stake: 100,

core/src/replay_stage.rs

@@ -15,7 +15,10 @@ use {
             DuplicateConfirmedSlotsReceiver, GossipVerifiedVoteHashReceiver, VoteTracker,
         },
         cluster_slots_service::{cluster_slots::ClusterSlots, ClusterSlotsUpdateSender},
-        commitment_service::{AggregateCommitmentService, CommitmentAggregationData},
+        commitment_service::{
+            AggregateCommitmentService, AlpenglowCommitmentAggregationData,
+            AlpenglowCommitmentType, CommitmentAggregationData, TowerCommitmentAggregationData,
+        },
         consensus::{
             fork_choice::{select_vote_and_reset_forks, ForkChoice, SelectVoteAndResetForkResult},
             heaviest_subtree_fork_choice::HeaviestSubtreeForkChoice,
@@ -70,7 +73,7 @@ use {
         accounts_background_service::AbsRequestSender,
         bank::{bank_hash_details, Bank, NewBankOptions},
         bank_forks::{BankForks, SetRootError, MAX_ROOT_DISTANCE_FOR_VOTE_ONLY},
-        commitment::{BlockCommitmentCache, CommitmentSlots},
+        commitment::BlockCommitmentCache,
         installed_scheduler_pool::BankWithScheduler,
         prioritization_fee_cache::PrioritizationFeeCache,
         vote_sender_types::ReplayVoteSender,
@@ -1289,19 +1292,14 @@ impl ReplayStage {
                             &drop_bank_sender,
                             &mut cert_pool,
                             &mut vote_history,
+                            &lockouts_sender,
                         ) {
                             error!(
                                 "Unable to set alpenglow root {}, error {e}",
                                 new_root_bank.slot()
                             );
                             return;
                         }
-                        rpc_subscriptions.notify_subscribers(CommitmentSlots {
-                            slot: new_root_slot,
-                            root: new_root_slot,
-                            highest_confirmed_slot: new_root_slot,
-                            highest_super_majority_root: new_root_slot,
-                        });
                     }
                 }
 
@@ -2750,7 +2748,6 @@ impl ReplayStage {
         vote_account_pubkey: &Pubkey,
         identity_keypair: &Keypair,
         authorized_voter_keypairs: &[Arc<Keypair>],
-        _lockouts_sender: &Sender<CommitmentAggregationData>,
         vote_signatures: &mut Vec<Signature>,
         has_new_vote_been_rooted: &mut bool,
         replay_timing: &mut ReplayLoopTiming,
@@ -3393,7 +3390,6 @@ impl ReplayStage {
                 vote_account_pubkey,
                 identity_keypair,
                 &authorized_voter_keypairs.read().unwrap(),
-                lockouts_sender,
                 voted_signatures,
                 has_new_vote_been_rooted,
                 replay_timing,
@@ -3409,6 +3405,11 @@ impl ReplayStage {
                     None,
                 ),
             );
+            Self::alpenglow_update_commitment_cache(
+                AlpenglowCommitmentType::Notarize,
+                highest_frozen_bank.slot(),
+                lockouts_sender,
+            );
         }
 
         // Try to finalize the highest notarized block
@@ -3431,7 +3432,6 @@ impl ReplayStage {
                         vote_account_pubkey,
                         identity_keypair,
                         &authorized_voter_keypairs.read().unwrap(),
-                        lockouts_sender,
                         voted_signatures,
                         has_new_vote_been_rooted,
                         replay_timing,
@@ -3504,19 +3504,35 @@ impl ReplayStage {
         }
     }
 
+    fn alpenglow_update_commitment_cache(
+        commitment_type: AlpenglowCommitmentType,
+        slot: Slot,
+        lockouts_sender: &Sender<CommitmentAggregationData>,
+    ) {
+        if let Err(e) = lockouts_sender.send(
+            CommitmentAggregationData::AlpenglowCommitmentAggregationData(
+                AlpenglowCommitmentAggregationData {
+                    commitment_type,
+                    slot,
+                },
+            ),
+        ) {
+            trace!("lockouts_sender failed: {:?}", e);
+        }
+    }
+
     fn update_commitment_cache(
         bank: Arc<Bank>,
         root: Slot,
         total_stake: Stake,
         node_vote_state: (Pubkey, TowerVoteState),
         lockouts_sender: &Sender<CommitmentAggregationData>,
     ) {
-        if let Err(e) = lockouts_sender.send(CommitmentAggregationData::new(
-            bank,
-            root,
-            total_stake,
-            node_vote_state,
-        )) {
+        if let Err(e) =
+            lockouts_sender.send(CommitmentAggregationData::TowerCommitmentAggregationData(
+                TowerCommitmentAggregationData::new(bank, root, total_stake, node_vote_state),
+            ))
+        {
             trace!("lockouts_sender failed: {:?}", e);
         }
     }
@@ -4740,6 +4756,7 @@ impl ReplayStage {
         drop_bank_sender: &Sender<Vec<BankWithScheduler>>,
         cert_pool: &mut CertificatePool,
         vote_history: &mut VoteHistory,
+        lockouts_sender: &Sender<CommitmentAggregationData>,
     ) -> Result<(), SetRootError> {
         vote_history.set_root(new_root);
         cert_pool.purge(new_root);
@@ -4762,7 +4779,13 @@ impl ReplayStage {
             voted_signatures,
             epoch_slots_frozen_slots,
             drop_bank_sender,
-        )
+        )?;
+        Self::alpenglow_update_commitment_cache(
+            AlpenglowCommitmentType::Root,
+            new_root,
+            lockouts_sender,
+        );
+        Ok(())
     }
 
     #[allow(clippy::too_many_arguments)]

runtime/src/commitment.rs

@@ -180,12 +180,22 @@ impl BlockCommitmentCache {
         }
     }
 
+    pub fn set_slot(&mut self, slot: Slot) {
+        self.commitment_slots.slot = std::cmp::max(self.commitment_slots.slot, slot);
+    }
+
     pub fn set_highest_confirmed_slot(&mut self, slot: Slot) {
-        self.commitment_slots.highest_confirmed_slot = slot;
+        self.commitment_slots.highest_confirmed_slot =
+            std::cmp::max(self.commitment_slots.highest_confirmed_slot, slot);
     }
 
-    pub fn set_highest_super_majority_root(&mut self, root: Slot) {
-        self.commitment_slots.highest_super_majority_root = root;
+    pub fn set_root(&mut self, slot: Slot) {
+        self.commitment_slots.root = std::cmp::max(self.commitment_slots.root, slot);
+    }
+
+    pub fn set_highest_super_majority_root(&mut self, slot: Slot) {
+        self.commitment_slots.highest_super_majority_root =
+            std::cmp::max(self.commitment_slots.highest_super_majority_root, slot);
     }
 
     pub fn initialize_slots(&mut self, slot: Slot, root: Slot) {
@@ -337,4 +347,28 @@ mod tests {
 
         assert_eq!(block_commitment_cache.calculate_highest_confirmed_slot(), 0);
     }
+
+    #[test]
+    fn test_setters_getters() {
+        let mut block_commitment_cache = BlockCommitmentCache::default();
+        // Setting bigger slots should be ok
+        block_commitment_cache.set_slot(1);
+        assert_eq!(block_commitment_cache.slot(), 1);
+        block_commitment_cache.set_highest_confirmed_slot(2);
+        assert_eq!(block_commitment_cache.highest_confirmed_slot(), 2);
+        block_commitment_cache.set_root(3);
+        assert_eq!(block_commitment_cache.root(), 3);
+        block_commitment_cache.set_highest_super_majority_root(4);
+        assert_eq!(block_commitment_cache.highest_super_majority_root(), 4);
+
+        // Setting smaller slots shuold be ignored
+        block_commitment_cache.set_slot(0);
+        assert_eq!(block_commitment_cache.slot(), 1);
+        block_commitment_cache.set_highest_confirmed_slot(1);
+        assert_eq!(block_commitment_cache.highest_confirmed_slot(), 2);
+        block_commitment_cache.set_root(2);
+        assert_eq!(block_commitment_cache.root(), 3);
+        block_commitment_cache.set_highest_super_majority_root(3);
+        assert_eq!(block_commitment_cache.highest_super_majority_root(), 4);
+    }
 }