Improve load_largest_accounts and add Bank test (#15775)
This commit is contained in:
parent
52e54e1100
commit
5991cef5f5
|
@ -538,9 +538,8 @@ impl Accounts {
|
||||||
filter_by_address: &HashSet<Pubkey>,
|
filter_by_address: &HashSet<Pubkey>,
|
||||||
filter: AccountAddressFilter,
|
filter: AccountAddressFilter,
|
||||||
) -> Vec<(Pubkey, u64)> {
|
) -> Vec<(Pubkey, u64)> {
|
||||||
let mut accounts_balances = self.accounts_db.scan_accounts(
|
self.accounts_db
|
||||||
ancestors,
|
.scan_accounts(ancestors, |collector: &mut Vec<(Pubkey, u64)>, option| {
|
||||||
|collector: &mut Vec<(Pubkey, u64)>, option| {
|
|
||||||
if let Some(data) = option
|
if let Some(data) = option
|
||||||
.filter(|(pubkey, account, _)| {
|
.filter(|(pubkey, account, _)| {
|
||||||
let should_include_pubkey = match filter {
|
let should_include_pubkey = match filter {
|
||||||
|
@ -551,14 +550,24 @@ impl Accounts {
|
||||||
})
|
})
|
||||||
.map(|(pubkey, account, _slot)| (*pubkey, account.lamports))
|
.map(|(pubkey, account, _slot)| (*pubkey, account.lamports))
|
||||||
{
|
{
|
||||||
collector.push(data)
|
let index_of_first_smaller_account =
|
||||||
|
collector.iter().position(|&(entry_pubkey, entry_balance)| {
|
||||||
|
entry_balance < data.1
|
||||||
|
|| (entry_balance == data.1 && entry_pubkey > data.0)
|
||||||
|
});
|
||||||
|
match index_of_first_smaller_account {
|
||||||
|
Some(i) => {
|
||||||
|
collector.insert(i, data);
|
||||||
|
}
|
||||||
|
None => {
|
||||||
|
if collector.len() <= num {
|
||||||
|
collector.push(data);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
collector.truncate(num);
|
||||||
}
|
}
|
||||||
},
|
})
|
||||||
);
|
|
||||||
|
|
||||||
accounts_balances.sort_by(|a, b| a.1.cmp(&b.1).reverse());
|
|
||||||
accounts_balances.truncate(num);
|
|
||||||
accounts_balances
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn calculate_capitalization(
|
pub fn calculate_capitalization(
|
||||||
|
|
|
@ -12292,4 +12292,105 @@ pub(crate) mod tests {
|
||||||
let failure_log = failure_log_info.log_messages.clone().pop().unwrap();
|
let failure_log = failure_log_info.log_messages.clone().pop().unwrap();
|
||||||
assert!(failure_log.contains(&"failed".to_string()));
|
assert!(failure_log.contains(&"failed".to_string()));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_get_largest_accounts() {
|
||||||
|
let GenesisConfigInfo { genesis_config, .. } =
|
||||||
|
create_genesis_config_with_leader(42, &solana_sdk::pubkey::new_rand(), 42);
|
||||||
|
let bank = Bank::new(&genesis_config);
|
||||||
|
|
||||||
|
let pubkeys: Vec<_> = (0..5).map(|_| Pubkey::new_unique()).collect();
|
||||||
|
let pubkeys_hashset: HashSet<_> = pubkeys.iter().cloned().collect();
|
||||||
|
|
||||||
|
let pubkeys_balances: Vec<_> = pubkeys
|
||||||
|
.iter()
|
||||||
|
.cloned()
|
||||||
|
.zip(vec![
|
||||||
|
sol_to_lamports(2.0),
|
||||||
|
sol_to_lamports(3.0),
|
||||||
|
sol_to_lamports(3.0),
|
||||||
|
sol_to_lamports(4.0),
|
||||||
|
sol_to_lamports(5.0),
|
||||||
|
])
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
// Initialize accounts; all have larger SOL balances than current Bank built-ins
|
||||||
|
let account0 = Account::new(pubkeys_balances[0].1, 0, &Pubkey::default());
|
||||||
|
bank.store_account(&pubkeys_balances[0].0, &account0);
|
||||||
|
let account1 = Account::new(pubkeys_balances[1].1, 0, &Pubkey::default());
|
||||||
|
bank.store_account(&pubkeys_balances[1].0, &account1);
|
||||||
|
let account2 = Account::new(pubkeys_balances[2].1, 0, &Pubkey::default());
|
||||||
|
bank.store_account(&pubkeys_balances[2].0, &account2);
|
||||||
|
let account3 = Account::new(pubkeys_balances[3].1, 0, &Pubkey::default());
|
||||||
|
bank.store_account(&pubkeys_balances[3].0, &account3);
|
||||||
|
let account4 = Account::new(pubkeys_balances[4].1, 0, &Pubkey::default());
|
||||||
|
bank.store_account(&pubkeys_balances[4].0, &account4);
|
||||||
|
|
||||||
|
// Create HashSet to exclude an account
|
||||||
|
let exclude4: HashSet<_> = pubkeys[4..].iter().cloned().collect();
|
||||||
|
|
||||||
|
let mut sorted_accounts = pubkeys_balances.clone();
|
||||||
|
sorted_accounts.sort_by(|a, b| a.1.cmp(&b.1).reverse());
|
||||||
|
|
||||||
|
// Return only one largest account
|
||||||
|
assert_eq!(
|
||||||
|
bank.get_largest_accounts(1, &pubkeys_hashset, AccountAddressFilter::Include),
|
||||||
|
vec![(pubkeys[4], sol_to_lamports(5.0))]
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
bank.get_largest_accounts(1, &HashSet::new(), AccountAddressFilter::Exclude),
|
||||||
|
vec![(pubkeys[4], sol_to_lamports(5.0))]
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
bank.get_largest_accounts(1, &exclude4, AccountAddressFilter::Exclude),
|
||||||
|
vec![(pubkeys[3], sol_to_lamports(4.0))]
|
||||||
|
);
|
||||||
|
|
||||||
|
// Return all added accounts
|
||||||
|
let results =
|
||||||
|
bank.get_largest_accounts(10, &pubkeys_hashset, AccountAddressFilter::Include);
|
||||||
|
assert_eq!(results.len(), sorted_accounts.len());
|
||||||
|
for pubkey_balance in sorted_accounts.iter() {
|
||||||
|
assert!(results.contains(pubkey_balance));
|
||||||
|
}
|
||||||
|
let mut sorted_results = results.clone();
|
||||||
|
sorted_results.sort_by(|a, b| a.1.cmp(&b.1).reverse());
|
||||||
|
assert_eq!(sorted_results, results);
|
||||||
|
|
||||||
|
let expected_accounts = sorted_accounts[1..].to_vec();
|
||||||
|
let results = bank.get_largest_accounts(10, &exclude4, AccountAddressFilter::Exclude);
|
||||||
|
// results include 5 Bank builtins
|
||||||
|
assert_eq!(results.len(), 10);
|
||||||
|
for pubkey_balance in expected_accounts.iter() {
|
||||||
|
assert!(results.contains(pubkey_balance));
|
||||||
|
}
|
||||||
|
let mut sorted_results = results.clone();
|
||||||
|
sorted_results.sort_by(|a, b| a.1.cmp(&b.1).reverse());
|
||||||
|
assert_eq!(sorted_results, results);
|
||||||
|
|
||||||
|
// Return 3 added accounts
|
||||||
|
let expected_accounts = sorted_accounts[0..4].to_vec();
|
||||||
|
let results = bank.get_largest_accounts(4, &pubkeys_hashset, AccountAddressFilter::Include);
|
||||||
|
assert_eq!(results.len(), expected_accounts.len());
|
||||||
|
for pubkey_balance in expected_accounts.iter() {
|
||||||
|
assert!(results.contains(pubkey_balance));
|
||||||
|
}
|
||||||
|
|
||||||
|
let expected_accounts = expected_accounts[1..4].to_vec();
|
||||||
|
let results = bank.get_largest_accounts(3, &exclude4, AccountAddressFilter::Exclude);
|
||||||
|
assert_eq!(results.len(), expected_accounts.len());
|
||||||
|
for pubkey_balance in expected_accounts.iter() {
|
||||||
|
assert!(results.contains(pubkey_balance));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Exclude more, and non-sequential, accounts
|
||||||
|
let exclude: HashSet<_> = vec![pubkeys[0], pubkeys[2], pubkeys[4]]
|
||||||
|
.iter()
|
||||||
|
.cloned()
|
||||||
|
.collect();
|
||||||
|
assert_eq!(
|
||||||
|
bank.get_largest_accounts(2, &exclude, AccountAddressFilter::Exclude),
|
||||||
|
vec![pubkeys_balances[3], pubkeys_balances[1]]
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue