Factor out LockedPubkeyReferences (#10198)

Co-authored-by: Carl <carl@solana.com>

core/src/cluster_info_vote_listener.rs

@@ -3,6 +3,7 @@ use crate::{
     consensus::VOTE_THRESHOLD_SIZE,
     crds_value::CrdsValueLabel,
     poh_recorder::PohRecorder,
+    pubkey_references::LockedPubkeyReferences,
     result::{Error, Result},
     rpc_subscriptions::RpcSubscriptions,
     sigverify,
@@ -66,7 +67,7 @@ pub struct VoteTracker {
     epoch_authorized_voters: RwLock<HashMap<Epoch, Arc<EpochAuthorizedVoters>>>,
     leader_schedule_epoch: RwLock<Epoch>,
     current_epoch: RwLock<Epoch>,
-    keys: RwLock<HashSet<Arc<Pubkey>>>,
+    keys: LockedPubkeyReferences,
     epoch_schedule: EpochSchedule,
 }
 
@@ -132,7 +133,7 @@ impl VoteTracker {
             w_slot_vote_tracker.updates = Some(vec![pubkey.clone()]);
         }
 
-        self.keys.write().unwrap().insert(pubkey);
+        self.keys.get_or_insert(&pubkey);
     }
 
     fn update_leader_schedule_epoch(&self, root_bank: &Bank) {
@@ -182,10 +183,7 @@ impl VoteTracker {
                 .write()
                 .unwrap()
                 .retain(|epoch, _| epoch >= &root_epoch);
-            self.keys
+            self.keys.purge();
-                .write()
-                .unwrap()
-                .retain(|pubkey| Arc::strong_count(pubkey) > 1);
             *self.current_epoch.write().unwrap() = root_epoch;
         }
     }
@@ -513,16 +511,8 @@ impl ClusterInfoVoteListener {
                                 continue;
                             }
                         }
-                        let mut unduplicated_pubkey =
+                        let unduplicated_pubkey = vote_tracker.keys.get_or_insert(vote_pubkey);
-                            vote_tracker.keys.read().unwrap().get(vote_pubkey).cloned();
+                        diff.entry(slot).or_default().insert(unduplicated_pubkey);
-                        if unduplicated_pubkey.is_none() {
-                            let new_key = Arc::new(*vote_pubkey);
-                            vote_tracker.keys.write().unwrap().insert(new_key.clone());
-                            unduplicated_pubkey = Some(new_key);
-                        }
-                        diff.entry(slot)
-                            .or_default()
-                            .insert(unduplicated_pubkey.unwrap());
                     }
 
                     subscriptions.notify_vote(&vote);
@@ -731,7 +721,7 @@ mod tests {
 
         // Check `keys` and `epoch_authorized_voters` are purged when new
         // root bank moves to the next epoch
-        assert!(vote_tracker.keys.read().unwrap().contains(&new_voter));
+        assert!(vote_tracker.keys.0.read().unwrap().contains(&new_voter));
         let current_epoch = bank.epoch();
         let new_epoch_bank = Bank::new_from_parent(
             &bank,
@@ -740,7 +730,7 @@ mod tests {
                 .get_first_slot_in_epoch(current_epoch + 1),
         );
         vote_tracker.process_new_root_bank(&new_epoch_bank);
-        assert!(!vote_tracker.keys.read().unwrap().contains(&new_voter));
+        assert!(!vote_tracker.keys.0.read().unwrap().contains(&new_voter));
         assert_eq!(
             *vote_tracker.current_epoch.read().unwrap(),
             current_epoch + 1
@@ -994,6 +984,7 @@ mod tests {
         let ref_count = Arc::strong_count(
             &vote_tracker
                 .keys
+                .0
                 .read()
                 .unwrap()
                 .get(&validator0_keypairs.vote_keypair.pubkey())
@@ -1045,6 +1036,7 @@ mod tests {
         let ref_count = Arc::strong_count(
             &vote_tracker
                 .keys
+                .0
                 .read()
                 .unwrap()
                 .get(&validator0_keypairs.vote_keypair.pubkey())

core/src/cluster_slots.rs

@@ -1,6 +1,6 @@
 use crate::{
     cluster_info::ClusterInfo, contact_info::ContactInfo, epoch_slots::EpochSlots,
-    serve_repair::RepairType,
+    pubkey_references::LockedPubkeyReferences, serve_repair::RepairType,
 };
 use solana_ledger::bank_forks::BankForks;
 use solana_runtime::epoch_stakes::NodeIdToVoteAccounts;
@@ -16,7 +16,7 @@ pub type ClusterSlotsMap = RwLock<HashMap<Slot, Arc<RwLock<SlotPubkeys>>>>;
 #[derive(Default)]
 pub struct ClusterSlots {
     cluster_slots: ClusterSlotsMap,
-    keys: RwLock<HashSet<Arc<Pubkey>>>,
+    keys: LockedPubkeyReferences,
     since: RwLock<Option<u64>>,
     validator_stakes: RwLock<Arc<NodeIdToVoteAccounts>>,
     epoch: RwLock<Option<u64>>,
@@ -41,20 +41,12 @@ impl ClusterSlots {
                 if *slot <= root {
                     continue;
                 }
-                let pubkey = Arc::new(epoch_slots.from);
+                let unduplicated_pubkey = self.keys.get_or_insert(&epoch_slots.from);
-                let exists = self.keys.read().unwrap().get(&pubkey).is_some();
+                self.insert_node_id(*slot, unduplicated_pubkey);
-                if !exists {
-                    self.keys.write().unwrap().insert(pubkey.clone());
-                }
-                let from = self.keys.read().unwrap().get(&pubkey).unwrap().clone();
-                self.insert_node_id(*slot, from);
             }
         }
         self.cluster_slots.write().unwrap().retain(|x, _| *x > root);
-        self.keys
+        self.keys.purge();
-            .write()
-            .unwrap()
-            .retain(|x| Arc::strong_count(x) > 1);
         *self.since.write().unwrap() = since;
     }
 

core/src/pubkey_references.rs

@@ -1,5 +1,9 @@
 use solana_sdk::pubkey::Pubkey;
-use std::{collections::HashSet, rc::Rc};
+use std::{
+    collections::HashSet,
+    rc::Rc,
+    sync::{Arc, RwLock},
+};
 
 #[derive(Default)]
 pub struct PubkeyReferences(HashSet<Rc<Pubkey>>);
@@ -19,3 +23,22 @@ impl PubkeyReferences {
         self.0.retain(|x| Rc::strong_count(x) > 1);
     }
 }
+
+#[derive(Default)]
+pub struct LockedPubkeyReferences(pub RwLock<HashSet<Arc<Pubkey>>>);
+
+impl LockedPubkeyReferences {
+    pub fn get_or_insert(&self, pubkey: &Pubkey) -> Arc<Pubkey> {
+        let mut cached_pubkey = self.0.read().unwrap().get(pubkey).cloned();
+        if cached_pubkey.is_none() {
+            let new_pubkey = Arc::new(*pubkey);
+            self.0.write().unwrap().insert(new_pubkey.clone());
+            cached_pubkey = Some(new_pubkey);
+        }
+        cached_pubkey.unwrap()
+    }
+
+    pub fn purge(&self) {
+        self.0.write().unwrap().retain(|x| Arc::strong_count(x) > 1);
+    }
+}