From 42aaacf5208188f9a28f9a4e99fac703e7b56b0a Mon Sep 17 00:00:00 2001 From: carllin Date: Fri, 22 May 2020 23:23:17 -0700 Subject: [PATCH] Factor out LockedPubkeyReferences (#10198) Co-authored-by: Carl --- core/src/cluster_info_vote_listener.rs | 28 +++++++++----------------- core/src/cluster_slots.rs | 18 +++++------------ core/src/pubkey_references.rs | 25 ++++++++++++++++++++++- 3 files changed, 39 insertions(+), 32 deletions(-) diff --git a/core/src/cluster_info_vote_listener.rs b/core/src/cluster_info_vote_listener.rs index d4497e7ce6..e674a3431e 100644 --- a/core/src/cluster_info_vote_listener.rs +++ b/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>>, leader_schedule_epoch: RwLock, current_epoch: RwLock, - keys: RwLock>>, + 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 - .write() - .unwrap() - .retain(|pubkey| Arc::strong_count(pubkey) > 1); + self.keys.purge(); *self.current_epoch.write().unwrap() = root_epoch; } } @@ -513,16 +511,8 @@ impl ClusterInfoVoteListener { continue; } } - let mut unduplicated_pubkey = - vote_tracker.keys.read().unwrap().get(vote_pubkey).cloned(); - 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()); + let unduplicated_pubkey = vote_tracker.keys.get_or_insert(vote_pubkey); + diff.entry(slot).or_default().insert(unduplicated_pubkey); } 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()) diff --git a/core/src/cluster_slots.rs b/core/src/cluster_slots.rs index e798188601..85cf83fece 100644 --- a/core/src/cluster_slots.rs +++ b/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>>>; #[derive(Default)] pub struct ClusterSlots { cluster_slots: ClusterSlotsMap, - keys: RwLock>>, + keys: LockedPubkeyReferences, since: RwLock>, validator_stakes: RwLock>, epoch: RwLock>, @@ -41,20 +41,12 @@ impl ClusterSlots { if *slot <= root { continue; } - let pubkey = Arc::new(epoch_slots.from); - let exists = self.keys.read().unwrap().get(&pubkey).is_some(); - 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); + let unduplicated_pubkey = self.keys.get_or_insert(&epoch_slots.from); + self.insert_node_id(*slot, unduplicated_pubkey); } } self.cluster_slots.write().unwrap().retain(|x, _| *x > root); - self.keys - .write() - .unwrap() - .retain(|x| Arc::strong_count(x) > 1); + self.keys.purge(); *self.since.write().unwrap() = since; } diff --git a/core/src/pubkey_references.rs b/core/src/pubkey_references.rs index f456915389..0ba28a8008 100644 --- a/core/src/pubkey_references.rs +++ b/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>); @@ -19,3 +23,22 @@ impl PubkeyReferences { self.0.retain(|x| Rc::strong_count(x) > 1); } } + +#[derive(Default)] +pub struct LockedPubkeyReferences(pub RwLock>>); + +impl LockedPubkeyReferences { + pub fn get_or_insert(&self, pubkey: &Pubkey) -> Arc { + 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); + } +}