replay: refactor set-root to enable alpenglow to take over

Author: AshwinSekar

Cargo.lock

@@ -134,6 +134,7 @@ members = [
     "version",
     "vortexor",
     "vote",
+    "votor",
     "watchtower",
     "wen-restart",
     "xdp",
@@ -554,6 +555,7 @@ solana-version = { path = "version", version = "=3.0.0" }
 solana-vote = { path = "vote", version = "=3.0.0" }
 solana-vote-interface = "2.2.6"
 solana-vote-program = { path = "programs/vote", version = "=3.0.0", default-features = false }
+solana-votor = { path = "votor", version = "=3.0.0" }
 solana-wen-restart = { path = "wen-restart", version = "=3.0.0" }
 solana-zk-elgamal-proof-program = { path = "programs/zk-elgamal-proof", version = "=3.0.0" }
 solana-zk-keygen = { path = "zk-keygen", version = "=3.0.0" }

Cargo.toml

@@ -164,6 +164,7 @@ solana-validator-exit = { workspace = true }
 solana-version = { workspace = true }
 solana-vote = { workspace = true }
 solana-vote-program = { workspace = true }
+solana-votor = { workspace = true }
 solana-wen-restart = { workspace = true }
 static_assertions = { workspace = true }
 strum = { workspace = true, features = ["derive"] }

core/Cargo.toml

@@ -75,6 +75,7 @@ use {
     solana_time_utils::timestamp,
     solana_transaction::Transaction,
     solana_vote::vote_transaction::VoteTransaction,
+    solana_votor::root_utils,
     std::{
         collections::{HashMap, HashSet},
         num::{NonZeroUsize, Saturating},
@@ -2400,16 +2401,23 @@ impl ReplayStage {
         let new_root = tower.record_bank_vote(bank);
 
         if let Some(new_root) = new_root {
+            let highest_super_majority_root = Some(
+                block_commitment_cache
+                    .read()
+                    .unwrap()
+                    .highest_super_majority_root(),
+            );
             Self::check_and_handle_new_root(
-                bank,
+                &identity_keypair.pubkey(),
+                bank.parent_slot(),
                 new_root,
                 bank_forks,
                 progress,
                 blockstore,
                 leader_schedule_cache,
                 snapshot_controller,
                 rpc_subscriptions,
-                block_commitment_cache,
+                highest_super_majority_root,
                 bank_notification_sender,
                 has_new_vote_been_rooted,
                 tracked_vote_transactions,
@@ -3967,119 +3975,62 @@ impl ReplayStage {
     }
 
     #[allow(clippy::too_many_arguments)]
+    /// A wrapper around `root_utils::check_and_handle_new_root` which:
+    /// - calls into `root_utils::set_bank_forks_root`
+    /// - Executes `set_progress_and_tower_bft_root` to cleanup tower bft structs and the progress map
     fn check_and_handle_new_root(
-        vote_bank: &Bank,
+        my_pubkey: &Pubkey,
+        parent_slot: Slot,
         new_root: Slot,
         bank_forks: &RwLock<BankForks>,
         progress: &mut ProgressMap,
         blockstore: &Blockstore,
         leader_schedule_cache: &Arc<LeaderScheduleCache>,
         snapshot_controller: Option<&SnapshotController>,
         rpc_subscriptions: Option<&RpcSubscriptions>,
-        block_commitment_cache: &Arc<RwLock<BlockCommitmentCache>>,
+        highest_super_majority_root: Option<Slot>,
         bank_notification_sender: &Option<BankNotificationSenderConfig>,
         has_new_vote_been_rooted: &mut bool,
         tracked_vote_transactions: &mut Vec<TrackedVoteTransaction>,
         drop_bank_sender: &Sender<Vec<BankWithScheduler>>,
         tbft_structs: &mut TowerBFTStructures,
     ) -> Result<(), SetRootError> {
-        // get the root bank before squash
-        let root_bank = bank_forks
-            .read()
-            .unwrap()
-            .get(new_root)
-            .expect("Root bank doesn't exist");
-        let mut rooted_banks = root_bank.parents();
-        let oldest_parent = rooted_banks.last().map(|last| last.parent_slot());
-        rooted_banks.push(root_bank.clone());
-        let rooted_slots: Vec<_> = rooted_banks.iter().map(|bank| bank.slot()).collect();
-        // The following differs from rooted_slots by including the parent slot of the oldest parent bank.
-        let rooted_slots_with_parents = bank_notification_sender
-            .as_ref()
-            .is_some_and(|sender| sender.should_send_parents)
-            .then(|| {
-                let mut new_chain = rooted_slots.clone();
-                new_chain.push(oldest_parent.unwrap_or_else(|| vote_bank.parent_slot()));
-                new_chain
-            });
-
-        // Call leader schedule_cache.set_root() before blockstore.set_root() because
-        // bank_forks.root is consumed by repair_service to update gossip, so we don't want to
-        // get shreds for repair on gossip before we update leader schedule, otherwise they may
-        // get dropped.
-        leader_schedule_cache.set_root(rooted_banks.last().unwrap());
-        blockstore
-            .set_roots(rooted_slots.iter())
-            .expect("Ledger set roots failed");
-        let highest_super_majority_root = Some(
-            block_commitment_cache
-                .read()
-                .unwrap()
-                .highest_super_majority_root(),
-        );
-        Self::handle_new_root(
+        root_utils::check_and_handle_new_root(
+            parent_slot,
             new_root,
-            bank_forks,
-            progress,
             snapshot_controller,
             highest_super_majority_root,
-            has_new_vote_been_rooted,
-            tracked_vote_transactions,
+            bank_notification_sender,
             drop_bank_sender,
-            tbft_structs,
-        )?;
-        blockstore.slots_stats.mark_rooted(new_root);
-        if let Some(rpc_subscriptions) = rpc_subscriptions {
-            rpc_subscriptions.notify_roots(rooted_slots);
-        }
-        if let Some(sender) = bank_notification_sender {
-            let dependency_work = sender
-                .dependency_tracker
-                .as_ref()
-                .map(|s| s.get_current_declared_work());
-            sender
-                .sender
-                .send((BankNotification::NewRootBank(root_bank), dependency_work))
-                .unwrap_or_else(|err| warn!("bank_notification_sender failed: {err:?}"));
-
-            if let Some(new_chain) = rooted_slots_with_parents {
-                sender
-                    .sender
-                    .send((BankNotification::NewRootedChain(new_chain), dependency_work))
-                    .unwrap_or_else(|err| warn!("bank_notification_sender failed: {err:?}"));
-            }
-        }
-        info!("new root {new_root}");
-        Ok(())
+            blockstore,
+            leader_schedule_cache,
+            bank_forks,
+            rpc_subscriptions,
+            my_pubkey,
+            move |bank_forks| {
+                Self::set_progress_and_tower_bft_root(
+                    new_root,
+                    bank_forks,
+                    progress,
+                    has_new_vote_been_rooted,
+                    tracked_vote_transactions,
+                    tbft_structs,
+                )
+            },
+        )
     }
 
-    #[allow(clippy::too_many_arguments)]
-    pub fn handle_new_root(
+    // To avoid code duplication and keep compatibility with alpenglow, we add this
+    // extra callback in the rooting path. This happens immediately after setting the bank forks root
+    fn set_progress_and_tower_bft_root(
         new_root: Slot,
-        bank_forks: &RwLock<BankForks>,
+        bank_forks: &BankForks,
         progress: &mut ProgressMap,
-        snapshot_controller: Option<&SnapshotController>,
-        highest_super_majority_root: Option<Slot>,
         has_new_vote_been_rooted: &mut bool,
         tracked_vote_transactions: &mut Vec<TrackedVoteTransaction>,
-        drop_bank_sender: &Sender<Vec<BankWithScheduler>>,
         tbft_structs: &mut TowerBFTStructures,
-    ) -> Result<(), SetRootError> {
-        bank_forks.read().unwrap().prune_program_cache(new_root);
-        let removed_banks = bank_forks.write().unwrap().set_root(
-            new_root,
-            snapshot_controller,
-            highest_super_majority_root,
-        )?;
-
-        drop_bank_sender
-            .send(removed_banks)
-            .unwrap_or_else(|err| warn!("bank drop failed: {err:?}"));
-
-        // Dropping the bank_forks write lock and reacquiring as a read lock is
-        // safe because updates to bank_forks are only made by a single thread.
-        let r_bank_forks = bank_forks.read().unwrap();
-        let new_root_bank = &r_bank_forks[new_root];
+    ) {
+        let new_root_bank = &bank_forks[new_root];
         if !*has_new_vote_been_rooted {
             for TrackedVoteTransaction {
                 message_hash,
@@ -4098,15 +4049,17 @@ impl ReplayStage {
                 std::mem::take(tracked_vote_transactions);
             }
         }
-        progress.handle_new_root(&r_bank_forks);
+
+        progress.handle_new_root(bank_forks);
         let TowerBFTStructures {
             heaviest_subtree_fork_choice,
             duplicate_slots_tracker,
             duplicate_confirmed_slots,
-            epoch_slots_frozen_slots,
             unfrozen_gossip_verified_vote_hashes,
+            epoch_slots_frozen_slots,
+            ..
         } = tbft_structs;
-        heaviest_subtree_fork_choice.set_tree_root((new_root, r_bank_forks.root_bank().hash()));
+        heaviest_subtree_fork_choice.set_tree_root((new_root, bank_forks.root_bank().hash()));
         *duplicate_slots_tracker = duplicate_slots_tracker.split_off(&new_root);
         // duplicate_slots_tracker now only contains entries >= `new_root`
 
@@ -4115,10 +4068,41 @@ impl ReplayStage {
 
         unfrozen_gossip_verified_vote_hashes.set_root(new_root);
         *epoch_slots_frozen_slots = epoch_slots_frozen_slots.split_off(&new_root);
-        Ok(())
         // epoch_slots_frozen_slots now only contains entries >= `new_root`
     }
 
+    #[allow(clippy::too_many_arguments)]
+    pub fn handle_new_root(
+        new_root: Slot,
+        bank_forks: &RwLock<BankForks>,
+        progress: &mut ProgressMap,
+        snapshot_controller: Option<&SnapshotController>,
+        highest_super_majority_root: Option<Slot>,
+        has_new_vote_been_rooted: &mut bool,
+        tracked_vote_transactions: &mut Vec<TrackedVoteTransaction>,
+        drop_bank_sender: &Sender<Vec<BankWithScheduler>>,
+        tbft_structs: &mut TowerBFTStructures,
+    ) -> Result<(), SetRootError> {
+        root_utils::set_bank_forks_root(
+            new_root,
+            bank_forks,
+            snapshot_controller,
+            highest_super_majority_root,
+            drop_bank_sender,
+            move |bank_forks| {
+                Self::set_progress_and_tower_bft_root(
+                    new_root,
+                    bank_forks,
+                    progress,
+                    has_new_vote_been_rooted,
+                    tracked_vote_transactions,
+                    tbft_structs,
+                )
+            },
+        )?;
+        Ok(())
+    }
+
     fn generate_new_bank_forks(
         blockstore: &Blockstore,
         bank_forks: &RwLock<BankForks>,