93 lines
3.0 KiB
Rust
93 lines
3.0 KiB
Rust
|
use solana_ledger::leader_schedule_cache::LeaderScheduleCache;
|
||
|
|
use solana_ledger::shred::Shred;
|
||
|
|
use solana_runtime::bank::Bank;
|
||
|
|
use solana_sdk::pubkey::Pubkey;
|
||
|
|
use solana_sdk::timing::timestamp;
|
||
|
|
use std::collections::HashSet;
|
||
|
|
use std::sync::Arc;
|
||
|
|
use std::sync::RwLock;
|
||
|
|
|
||
|
|
///Configure a partition in the retransmit stage
|
||
|
|
#[derive(Debug, Clone)]
|
||
|
|
pub struct Partition {
|
||
|
|
pub num_partitions: usize,
|
||
|
|
pub my_partition: usize,
|
||
|
|
pub start_ts: u64,
|
||
|
|
pub end_ts: u64,
|
||
|
|
leaders: Arc<RwLock<Vec<Pubkey>>>,
|
||
|
|
}
|
||
|
|
impl Default for Partition {
|
||
|
|
fn default() -> Self {
|
||
|
|
Self {
|
||
|
|
num_partitions: 0,
|
||
|
|
my_partition: 0,
|
||
|
|
start_ts: 0,
|
||
|
|
end_ts: 0,
|
||
|
|
leaders: Arc::new(RwLock::new(vec![])),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
#[derive(Default, Debug, Clone)]
|
||
|
|
pub struct PartitionCfg {
|
||
|
|
partitions: Vec<Partition>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl PartitionCfg {
|
||
|
|
pub fn new(partitions: Vec<Partition>) -> Self {
|
||
|
|
Self { partitions }
|
||
|
|
}
|
||
|
|
pub fn is_connected(
|
||
|
|
&self,
|
||
|
|
bank: &Option<Arc<Bank>>,
|
||
|
|
leader_schedule_cache: &Arc<LeaderScheduleCache>,
|
||
|
|
shred: &Shred,
|
||
|
|
) -> bool {
|
||
|
|
if bank.is_none() {
|
||
|
|
return true;
|
||
|
|
}
|
||
|
|
let bank = bank.as_ref().unwrap().clone();
|
||
|
|
let slot_leader_pubkey = leader_schedule_cache.slot_leader_at(shred.slot(), Some(&bank));
|
||
|
|
let slot_leader_pubkey = slot_leader_pubkey.unwrap_or_default();
|
||
|
|
let time = timestamp();
|
||
|
|
for p in &self.partitions {
|
||
|
|
let is_time = (p.start_ts <= time) && (time < p.end_ts);
|
||
|
|
if !is_time {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
trace!("PARTITION_TEST partition time! {}", p.my_partition);
|
||
|
|
if p.num_partitions == 0 {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
if p.leaders.read().unwrap().is_empty() {
|
||
|
|
let mut leader_vec = p.leaders.write().unwrap();
|
||
|
|
let mut leaders: Vec<Pubkey> = bank.vote_accounts().keys().cloned().collect();
|
||
|
|
leaders.sort();
|
||
|
|
*leader_vec = leaders;
|
||
|
|
warn!("PARTITION_TEST partition enabled {}", p.my_partition);
|
||
|
|
}
|
||
|
|
let is_connected: bool = {
|
||
|
|
let leaders = p.leaders.read().unwrap();
|
||
|
|
let start = p.my_partition * leaders.len() / p.num_partitions;
|
||
|
|
let partition_size = leaders.len() / p.num_partitions;
|
||
|
|
let end = start + partition_size;
|
||
|
|
let end = if leaders.len() - end < partition_size {
|
||
|
|
leaders.len()
|
||
|
|
} else {
|
||
|
|
end
|
||
|
|
};
|
||
|
|
let my_leaders: HashSet<_> = leaders[start..end].iter().collect();
|
||
|
|
my_leaders.contains(&slot_leader_pubkey)
|
||
|
|
};
|
||
|
|
if is_connected {
|
||
|
|
trace!("PARTITION_TEST connected {}", p.my_partition);
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
trace!("PARTITION_TEST not connected {}", p.my_partition);
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
trace!("PARTITION_TEST connected");
|
||
|
|
true
|
||
|
|
}
|
||
|
|
}
|