From 7d96f78821a27e0d816a85265fb700eaf3f87b2b Mon Sep 17 00:00:00 2001 From: "Jeff Washington (jwash)" <75863576+jeffwashington@users.noreply.github.com> Date: Tue, 11 May 2021 17:06:22 -0500 Subject: [PATCH] include/exclude keys on account secondary index (#17110) * AccountSecondaryIndexes.include/exclude * use normal scan if key is not indexed * add a test to ask for a scan for an excluded secondary index * fix up cli args --- runtime/src/accounts.rs | 18 ++-- runtime/src/accounts_db.rs | 72 +++++++++++--- runtime/src/accounts_index.rs | 174 ++++++++++++++++++++++++++++++++-- runtime/src/bank.rs | 2 +- validator/src/main.rs | 88 ++++++++++++++--- 5 files changed, 314 insertions(+), 40 deletions(-) diff --git a/runtime/src/accounts.rs b/runtime/src/accounts.rs index 2e9868a4db..e33268013f 100644 --- a/runtime/src/accounts.rs +++ b/runtime/src/accounts.rs @@ -702,13 +702,17 @@ impl Accounts { index_key: &IndexKey, filter: F, ) -> Vec<(Pubkey, AccountSharedData)> { - self.accounts_db.index_scan_accounts( - ancestors, - *index_key, - |collector: &mut Vec<(Pubkey, AccountSharedData)>, some_account_tuple| { - Self::load_while_filtering(collector, some_account_tuple, |account| filter(account)) - }, - ) + self.accounts_db + .index_scan_accounts( + ancestors, + *index_key, + |collector: &mut Vec<(Pubkey, AccountSharedData)>, some_account_tuple| { + Self::load_while_filtering(collector, some_account_tuple, |account| { + filter(account) + }) + }, + ) + .0 } pub fn load_all(&self, ancestors: &Ancestors) -> Vec<(Pubkey, AccountSharedData, Slot)> { diff --git a/runtime/src/accounts_db.rs b/runtime/src/accounts_db.rs index 01b1a5b693..41a1343694 100644 --- a/runtime/src/accounts_db.rs +++ b/runtime/src/accounts_db.rs @@ -2270,11 +2270,22 @@ impl AccountsDb { ancestors: &Ancestors, index_key: IndexKey, scan_func: F, - ) -> A + ) -> (A, bool) where F: Fn(&mut A, Option<(&Pubkey, AccountSharedData, Slot)>), A: Default, { + let key = match &index_key { + IndexKey::ProgramId(key) => key, + IndexKey::SplTokenMint(key) => key, + IndexKey::SplTokenOwner(key) => key, + }; + if !self.account_indexes.include_key(key) { + // the requested key was not indexed in the secondary index, so do a normal scan + let used_index = false; + return (self.scan_accounts(ancestors, scan_func), used_index); + } + let mut collector = A::default(); self.accounts_index.index_scan_accounts( ancestors, @@ -2287,7 +2298,8 @@ impl AccountsDb { scan_func(&mut collector, account_slot) }, ); - collector + let used_index = true; + (collector, used_index) } /// Scan a specific slot through all the account storage in parallel @@ -5533,8 +5545,11 @@ impl AccountsDb { pub mod tests { use super::*; use crate::{ - accounts_hash::MERKLE_FANOUT, accounts_index::tests::*, accounts_index::RefCount, - append_vec::AccountMeta, inline_spl_token_v2_0, + accounts_hash::MERKLE_FANOUT, + accounts_index::RefCount, + accounts_index::{tests::*, AccountSecondaryIndexesIncludeExclude}, + append_vec::AccountMeta, + inline_spl_token_v2_0, }; use assert_matches::assert_matches; use rand::{thread_rng, Rng}; @@ -6853,7 +6868,7 @@ pub mod tests { fn test_clean_old_with_both_normal_and_zero_lamport_accounts() { solana_logger::setup(); - let accounts = AccountsDb::new_with_config( + let mut accounts = AccountsDb::new_with_config( Vec::new(), &ClusterType::Development, spl_token_mint_index_enabled(), @@ -6897,17 +6912,52 @@ pub mod tests { // Secondary index should still find both pubkeys let mut found_accounts = HashSet::new(); - accounts.accounts_index.index_scan_accounts( - &Ancestors::default(), - IndexKey::SplTokenMint(mint_key), - |key, _| { + let index_key = IndexKey::SplTokenMint(mint_key); + accounts + .accounts_index + .index_scan_accounts(&Ancestors::default(), index_key, |key, _| { found_accounts.insert(*key); - }, - ); + }); assert_eq!(found_accounts.len(), 2); assert!(found_accounts.contains(&pubkey1)); assert!(found_accounts.contains(&pubkey2)); + { + accounts.account_indexes.keys = Some(AccountSecondaryIndexesIncludeExclude { + exclude: true, + keys: [mint_key].iter().cloned().collect::>(), + }); + // Secondary index can't be used - do normal scan: should still find both pubkeys + let found_accounts = accounts.index_scan_accounts( + &Ancestors::default(), + index_key, + |collection: &mut HashSet, account| { + collection.insert(*account.unwrap().0); + }, + ); + assert!(!found_accounts.1); + assert_eq!(found_accounts.0.len(), 2); + assert!(found_accounts.0.contains(&pubkey1)); + assert!(found_accounts.0.contains(&pubkey2)); + + accounts.account_indexes.keys = None; + + // Secondary index can now be used since it isn't marked as excluded + let found_accounts = accounts.index_scan_accounts( + &Ancestors::default(), + index_key, + |collection: &mut HashSet, account| { + collection.insert(*account.unwrap().0); + }, + ); + assert!(found_accounts.1); + assert_eq!(found_accounts.0.len(), 2); + assert!(found_accounts.0.contains(&pubkey1)); + assert!(found_accounts.0.contains(&pubkey2)); + + accounts.account_indexes.keys = None; + } + accounts.clean_accounts(None); //both zero lamport and normal accounts are cleaned up diff --git a/runtime/src/accounts_index.rs b/runtime/src/accounts_index.rs index b5c030a394..633e0aac77 100644 --- a/runtime/src/accounts_index.rs +++ b/runtime/src/accounts_index.rs @@ -74,7 +74,32 @@ pub enum AccountIndex { SplTokenOwner, } -pub type AccountSecondaryIndexes = HashSet; +#[derive(Debug, PartialEq, Eq, Clone)] +pub struct AccountSecondaryIndexesIncludeExclude { + pub exclude: bool, + pub keys: HashSet, +} + +#[derive(Debug, Default, Clone)] +pub struct AccountSecondaryIndexes { + pub keys: Option, + pub indexes: HashSet, +} + +impl AccountSecondaryIndexes { + pub fn is_empty(&self) -> bool { + self.indexes.is_empty() + } + pub fn contains(&self, index: &AccountIndex) -> bool { + self.indexes.contains(index) + } + pub fn include_key(&self, key: &Pubkey) -> bool { + match &self.keys { + Some(options) => options.exclude ^ options.keys.contains(key), + None => true, // include all keys + } + } +} #[derive(Debug)] pub struct AccountMapEntryInner { @@ -1063,7 +1088,9 @@ impl AccountsIndex { return; } - if account_indexes.contains(&AccountIndex::ProgramId) { + if account_indexes.contains(&AccountIndex::ProgramId) + && account_indexes.include_key(account_owner) + { self.program_id_index.insert(account_owner, pubkey, slot); } // Note because of the below check below on the account data length, when an @@ -1087,7 +1114,9 @@ impl AccountsIndex { &account_data[SPL_TOKEN_ACCOUNT_OWNER_OFFSET ..SPL_TOKEN_ACCOUNT_OWNER_OFFSET + PUBKEY_BYTES], ); - self.spl_token_owner_index.insert(&owner_key, pubkey, slot); + if account_indexes.include_key(&owner_key) { + self.spl_token_owner_index.insert(&owner_key, pubkey, slot); + } } if account_indexes.contains(&AccountIndex::SplTokenMint) { @@ -1095,7 +1124,9 @@ impl AccountsIndex { &account_data[SPL_TOKEN_ACCOUNT_MINT_OFFSET ..SPL_TOKEN_ACCOUNT_MINT_OFFSET + PUBKEY_BYTES], ); - self.spl_token_mint_index.insert(&mint_key, pubkey, slot); + if account_indexes.include_key(&mint_key) { + self.spl_token_mint_index.insert(&mint_key, pubkey, slot); + } } } } @@ -1447,13 +1478,19 @@ pub mod tests { pub fn spl_token_mint_index_enabled() -> AccountSecondaryIndexes { let mut account_indexes = HashSet::new(); account_indexes.insert(AccountIndex::SplTokenMint); - account_indexes + AccountSecondaryIndexes { + indexes: account_indexes, + keys: None, + } } pub fn spl_token_owner_index_enabled() -> AccountSecondaryIndexes { let mut account_indexes = HashSet::new(); account_indexes.insert(AccountIndex::SplTokenOwner); - account_indexes + AccountSecondaryIndexes { + indexes: account_indexes, + keys: None, + } } impl<'a, T: 'static, U> AccountIndexGetResult<'a, T, U> { @@ -2070,6 +2107,51 @@ pub mod tests { assert_eq!(num, 0); } + #[test] + fn test_secondary_index_include_exclude() { + let pk1 = Pubkey::new_unique(); + let pk2 = Pubkey::new_unique(); + let mut index = AccountSecondaryIndexes::default(); + + assert!(!index.contains(&AccountIndex::ProgramId)); + index.indexes.insert(AccountIndex::ProgramId); + assert!(index.contains(&AccountIndex::ProgramId)); + assert!(index.include_key(&pk1)); + assert!(index.include_key(&pk2)); + + let exclude = false; + index.keys = Some(AccountSecondaryIndexesIncludeExclude { + keys: [pk1].iter().cloned().collect::>(), + exclude, + }); + assert!(index.include_key(&pk1)); + assert!(!index.include_key(&pk2)); + + let exclude = true; + index.keys = Some(AccountSecondaryIndexesIncludeExclude { + keys: [pk1].iter().cloned().collect::>(), + exclude, + }); + assert!(!index.include_key(&pk1)); + assert!(index.include_key(&pk2)); + + let exclude = true; + index.keys = Some(AccountSecondaryIndexesIncludeExclude { + keys: [pk1, pk2].iter().cloned().collect::>(), + exclude, + }); + assert!(!index.include_key(&pk1)); + assert!(!index.include_key(&pk2)); + + let exclude = false; + index.keys = Some(AccountSecondaryIndexesIncludeExclude { + keys: [pk1, pk2].iter().cloned().collect::>(), + exclude, + }); + assert!(index.include_key(&pk1)); + assert!(index.include_key(&pk2)); + } + #[test] fn test_insert_no_ancestors() { let key = Keypair::new(); @@ -2902,6 +2984,7 @@ pub mod tests { key_end: usize, account_index: &AccountSecondaryIndexes, ) { + let mut account_index = account_index.clone(); let account_key = Pubkey::new_unique(); let index_key = Pubkey::new_unique(); let slot = 1; @@ -2914,7 +2997,7 @@ pub mod tests { &account_key, &Pubkey::default(), &account_data, - account_index, + &account_index, true, &mut vec![], ); @@ -2927,13 +3010,49 @@ pub mod tests { &account_key, &inline_spl_token_v2_0::id(), &account_data[1..], - account_index, + &account_index, true, &mut vec![], ); assert!(secondary_index.index.is_empty()); assert!(secondary_index.reverse_index.is_empty()); + // excluded key + account_index.keys = Some(AccountSecondaryIndexesIncludeExclude { + keys: [index_key].iter().cloned().collect::>(), + exclude: true, + }); + index.upsert( + 0, + &account_key, + &inline_spl_token_v2_0::id(), + &account_data[1..], + &account_index, + true, + &mut vec![], + ); + assert!(secondary_index.index.is_empty()); + assert!(secondary_index.reverse_index.is_empty()); + + // not-included key + account_index.keys = Some(AccountSecondaryIndexesIncludeExclude { + keys: [account_key].iter().cloned().collect::>(), + exclude: false, + }); + index.upsert( + 0, + &account_key, + &inline_spl_token_v2_0::id(), + &account_data[1..], + &account_index, + true, + &mut vec![], + ); + assert!(secondary_index.index.is_empty()); + assert!(secondary_index.reverse_index.is_empty()); + + account_index.keys = None; + // Just right. Inserting the same index multiple times should be ok for _ in 0..2 { index.update_secondary_indexes( @@ -2941,7 +3060,7 @@ pub mod tests { slot, &inline_spl_token_v2_0::id(), &account_data, - account_index, + &account_index, ); check_secondary_index_unique(secondary_index, slot, &index_key, &account_key); } @@ -2950,6 +3069,43 @@ pub mod tests { assert!(!secondary_index.index.is_empty()); assert!(!secondary_index.reverse_index.is_empty()); + account_index.keys = Some(AccountSecondaryIndexesIncludeExclude { + keys: [index_key].iter().cloned().collect::>(), + exclude: false, + }); + secondary_index.index.clear(); + secondary_index.reverse_index.clear(); + index.update_secondary_indexes( + &account_key, + slot, + &inline_spl_token_v2_0::id(), + &account_data, + &account_index, + ); + assert!(!secondary_index.index.is_empty()); + assert!(!secondary_index.reverse_index.is_empty()); + check_secondary_index_unique(secondary_index, slot, &index_key, &account_key); + + // not-excluded + account_index.keys = Some(AccountSecondaryIndexesIncludeExclude { + keys: [].iter().cloned().collect::>(), + exclude: true, + }); + secondary_index.index.clear(); + secondary_index.reverse_index.clear(); + index.update_secondary_indexes( + &account_key, + slot, + &inline_spl_token_v2_0::id(), + &account_data, + &account_index, + ); + assert!(!secondary_index.index.is_empty()); + assert!(!secondary_index.reverse_index.is_empty()); + check_secondary_index_unique(secondary_index, slot, &index_key, &account_key); + + account_index.keys = None; + index .get_account_write_entry(&account_key) .unwrap() diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs index 6eafe96063..f08dc137b5 100644 --- a/runtime/src/bank.rs +++ b/runtime/src/bank.rs @@ -9046,7 +9046,7 @@ pub(crate) mod tests { fn test_get_filtered_indexed_accounts() { let (genesis_config, _mint_keypair) = create_genesis_config(500); let mut account_indexes = AccountSecondaryIndexes::default(); - account_indexes.insert(AccountIndex::ProgramId); + account_indexes.indexes.insert(AccountIndex::ProgramId); let bank = Arc::new(Bank::new_with_config( &genesis_config, account_indexes, diff --git a/validator/src/main.rs b/validator/src/main.rs index a4c0b9ada9..4eee25f6ea 100644 --- a/validator/src/main.rs +++ b/validator/src/main.rs @@ -39,7 +39,9 @@ use { solana_ledger::blockstore_db::BlockstoreRecoveryMode, solana_perf::recycler::enable_recycler_warming, solana_runtime::{ - accounts_index::AccountIndex, + accounts_index::{ + AccountIndex, AccountSecondaryIndexes, AccountSecondaryIndexesIncludeExclude, + }, bank_forks::{ArchiveFormat, SnapshotConfig, SnapshotVersion}, hardened_unpack::{unpack_genesis_archive, MAX_GENESIS_ARCHIVE_UNPACKED_SIZE}, snapshot_utils::get_highest_snapshot_archive_path, @@ -79,6 +81,9 @@ enum Operation { Run, } +const EXCLUDE_KEY: &str = "account-index-exclude-key"; +const INCLUDE_KEY: &str = "account-index-include-key"; + fn monitor_validator(ledger_path: &Path) { let dashboard = Dashboard::new(ledger_path, None, None).unwrap_or_else(|err| { println!( @@ -1707,6 +1712,25 @@ pub fn main() { .value_name("INDEX") .help("Enable an accounts index, indexed by the selected account field"), ) + .arg( + Arg::with_name("account_index_exclude_key") + .long(EXCLUDE_KEY) + .takes_value(true) + .validator(is_pubkey) + .multiple(true) + .value_name("KEY") + .help("When account indexes are enabled, exclude this key from the index."), + ) + .arg( + Arg::with_name("account_index_include_key") + .long(INCLUDE_KEY) + .takes_value(true) + .validator(is_pubkey) + .conflicts_with("account_index_exclude_key") + .multiple(true) + .value_name("KEY") + .help("When account indexes are enabled, only include specific keys in the index. This overrides --account-index-exclude-key."), + ) .arg( Arg::with_name("no_accounts_db_caching") .long("no-accounts-db-caching") @@ -2033,16 +2057,7 @@ pub fn main() { let contact_debug_interval = value_t_or_exit!(matches, "contact_debug_interval", u64); - let account_indexes: HashSet = matches - .values_of("account_indexes") - .unwrap_or_default() - .map(|value| match value { - "program-id" => AccountIndex::ProgramId, - "spl-token-mint" => AccountIndex::SplTokenMint, - "spl-token-owner" => AccountIndex::SplTokenOwner, - _ => unreachable!(), - }) - .collect(); + let account_indexes = process_account_indexes(&matches); let restricted_repair_only_mode = matches.is_present("restricted_repair_only_mode"); let mut validator_config = ValidatorConfig { @@ -2084,7 +2099,7 @@ pub fn main() { rpc_bigtable_timeout: value_t!(matches, "rpc_bigtable_timeout", u64) .ok() .map(Duration::from_secs), - account_indexes: account_indexes.clone(), + account_indexes: account_indexes.indexes.clone(), }, rpc_addrs: value_t!(matches, "rpc_port", u16).ok().map(|rpc_port| { ( @@ -2499,3 +2514,52 @@ pub fn main() { validator.join(); info!("Validator exiting.."); } + +fn process_account_indexes(matches: &ArgMatches) -> AccountSecondaryIndexes { + let account_indexes: HashSet = matches + .values_of("account_indexes") + .unwrap_or_default() + .map(|value| match value { + "program-id" => AccountIndex::ProgramId, + "spl-token-mint" => AccountIndex::SplTokenMint, + "spl-token-owner" => AccountIndex::SplTokenOwner, + _ => unreachable!(), + }) + .collect(); + + let account_indexes_include_keys: HashSet = + values_t!(matches, "account_index_include_key", Pubkey) + .unwrap_or_default() + .iter() + .cloned() + .collect(); + + let account_indexes_exclude_keys: HashSet = + values_t!(matches, "account_index_exclude_key", Pubkey) + .unwrap_or_default() + .iter() + .cloned() + .collect(); + + let exclude_keys = !account_indexes_exclude_keys.is_empty(); + let include_keys = !account_indexes_include_keys.is_empty(); + + let keys = if !account_indexes.is_empty() && (exclude_keys || include_keys) { + let account_indexes_keys = AccountSecondaryIndexesIncludeExclude { + exclude: exclude_keys, + keys: if exclude_keys { + account_indexes_exclude_keys + } else { + account_indexes_include_keys + }, + }; + Some(account_indexes_keys) + } else { + None + }; + + AccountSecondaryIndexes { + keys, + indexes: account_indexes, + } +}