From b01d7923fc46de22b688b27f19122cd62227550e Mon Sep 17 00:00:00 2001
From: carllin <wumu727@gmail.com>
Date: Tue, 26 Mar 2024 00:34:15 -0400
Subject: [PATCH] Add local cluster utitlity functions (#355)

---
 core/src/consensus.rs              |  1 -
 local-cluster/src/cluster_tests.rs | 54 ++++++++++++++++++++++++++++--
 local-cluster/src/local_cluster.rs | 39 ++++++++++++++++++---
 3 files changed, 86 insertions(+), 8 deletions(-)

diff --git a/core/src/consensus.rs b/core/src/consensus.rs
index ab316d7c7da612..d4f2345aa14ab8 100644
--- a/core/src/consensus.rs
+++ b/core/src/consensus.rs
@@ -631,7 +631,6 @@ impl Tower {
         }
     }
 
-    #[cfg(test)]
     pub fn record_vote(&mut self, slot: Slot, hash: Hash) -> Option<Slot> {
         self.record_bank_vote_and_update_lockouts(slot, hash)
     }
diff --git a/local-cluster/src/cluster_tests.rs b/local-cluster/src/cluster_tests.rs
index 90337bb272460f..dffe2a8713ab08 100644
--- a/local-cluster/src/cluster_tests.rs
+++ b/local-cluster/src/cluster_tests.rs
@@ -10,7 +10,10 @@ use {
         connection_cache::{ConnectionCache, Protocol},
         thin_client::ThinClient,
     },
-    solana_core::consensus::VOTE_THRESHOLD_DEPTH,
+    solana_core::consensus::{
+        tower_storage::{FileTowerStorage, SavedTower, SavedTowerVersions, TowerStorage},
+        VOTE_THRESHOLD_DEPTH,
+    },
     solana_entry::entry::{Entry, EntrySlice},
     solana_gossip::{
         cluster_info::{self, ClusterInfo},
@@ -43,7 +46,7 @@ use {
         borrow::Borrow,
         collections::{HashMap, HashSet, VecDeque},
         net::{IpAddr, Ipv4Addr, SocketAddr, TcpListener},
-        path::Path,
+        path::{Path, PathBuf},
         sync::{
             atomic::{AtomicBool, Ordering},
             Arc, RwLock,
@@ -334,6 +337,53 @@ pub fn kill_entry_and_spend_and_verify_rest(
     }
 }
 
+pub fn apply_votes_to_tower(node_keypair: &Keypair, votes: Vec<(Slot, Hash)>, tower_path: PathBuf) {
+    let tower_storage = FileTowerStorage::new(tower_path);
+    let mut tower = tower_storage.load(&node_keypair.pubkey()).unwrap();
+    for (slot, hash) in votes {
+        tower.record_vote(slot, hash);
+    }
+    let saved_tower = SavedTowerVersions::from(SavedTower::new(&tower, node_keypair).unwrap());
+    tower_storage.store(&saved_tower).unwrap();
+}
+
+pub fn check_min_slot_is_rooted(
+    min_slot: Slot,
+    contact_infos: &[ContactInfo],
+    connection_cache: &Arc<ConnectionCache>,
+    test_name: &str,
+) {
+    let mut last_print = Instant::now();
+    let loop_start = Instant::now();
+    let loop_timeout = Duration::from_secs(180);
+    for ingress_node in contact_infos.iter() {
+        let (rpc, tpu) = LegacyContactInfo::try_from(ingress_node)
+            .map(|node| get_client_facing_addr(connection_cache.protocol(), node))
+            .unwrap();
+        let client = ThinClient::new(rpc, tpu, connection_cache.clone());
+        loop {
+            let root_slot = client
+                .get_slot_with_commitment(CommitmentConfig::finalized())
+                .unwrap_or(0);
+            if root_slot >= min_slot || last_print.elapsed().as_secs() > 3 {
+                info!(
+                    "{} waiting for node {} to see root >= {}.. observed latest root: {}",
+                    test_name,
+                    ingress_node.pubkey(),
+                    min_slot,
+                    root_slot
+                );
+                last_print = Instant::now();
+                if root_slot >= min_slot {
+                    break;
+                }
+            }
+            sleep(Duration::from_millis(clock::DEFAULT_MS_PER_SLOT / 2));
+            assert!(loop_start.elapsed() < loop_timeout);
+        }
+    }
+}
+
 pub fn check_for_new_roots(
     num_new_roots: usize,
     contact_infos: &[ContactInfo],
diff --git a/local-cluster/src/local_cluster.rs b/local-cluster/src/local_cluster.rs
index d06c001bcc7ed1..3d8df638fbbb81 100644
--- a/local-cluster/src/local_cluster.rs
+++ b/local-cluster/src/local_cluster.rs
@@ -30,7 +30,7 @@ use {
     solana_sdk::{
         account::{Account, AccountSharedData},
         client::SyncClient,
-        clock::{DEFAULT_DEV_SLOTS_PER_EPOCH, DEFAULT_TICKS_PER_SLOT},
+        clock::{Slot, DEFAULT_DEV_SLOTS_PER_EPOCH, DEFAULT_TICKS_PER_SLOT},
         commitment_config::CommitmentConfig,
         epoch_schedule::EpochSchedule,
         feature_set,
@@ -555,12 +555,11 @@ impl LocalCluster {
         Self::transfer_with_client(&client, source_keypair, dest_pubkey, lamports)
     }
 
-    pub fn check_for_new_roots(
+    fn discover_nodes(
         &self,
-        num_new_roots: usize,
-        test_name: &str,
         socket_addr_space: SocketAddrSpace,
-    ) {
+        test_name: &str,
+    ) -> Vec<ContactInfo> {
         let alive_node_contact_infos: Vec<_> = self
             .validators
             .values()
@@ -575,6 +574,36 @@ impl LocalCluster {
         )
         .unwrap();
         info!("{} discovered {} nodes", test_name, cluster_nodes.len());
+        alive_node_contact_infos
+    }
+
+    pub fn check_min_slot_is_rooted(
+        &self,
+        min_root: Slot,
+        test_name: &str,
+        socket_addr_space: SocketAddrSpace,
+    ) {
+        let alive_node_contact_infos = self.discover_nodes(socket_addr_space, test_name);
+        info!(
+            "{} looking minimum root {} on all nodes",
+            min_root, test_name
+        );
+        cluster_tests::check_min_slot_is_rooted(
+            min_root,
+            &alive_node_contact_infos,
+            &self.connection_cache,
+            test_name,
+        );
+        info!("{} done waiting for roots", test_name);
+    }
+
+    pub fn check_for_new_roots(
+        &self,
+        num_new_roots: usize,
+        test_name: &str,
+        socket_addr_space: SocketAddrSpace,
+    ) {
+        let alive_node_contact_infos = self.discover_nodes(socket_addr_space, test_name);
         info!("{} looking for new roots on all nodes", test_name);
         cluster_tests::check_for_new_roots(
             num_new_roots,