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

@@ -702,13 +702,17 @@ impl Accounts {
         index_key: &IndexKey,
         filter: F,
     ) -> Vec<(Pubkey, AccountSharedData)> {
-        self.accounts_db.index_scan_accounts(
+        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::load_while_filtering(collector, some_account_tuple, |account| {
+                        filter(account)
+                    })
                 },
             )
+            .0
     }
 
     pub fn load_all(&self, ancestors: &Ancestors) -> Vec<(Pubkey, AccountSharedData, Slot)> {

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,
+        accounts_hash::MERKLE_FANOUT,
-        append_vec::AccountMeta, inline_spl_token_v2_0,
+        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(
+        let index_key = IndexKey::SplTokenMint(mint_key);
-            &Ancestors::default(),
+        accounts
-            IndexKey::SplTokenMint(mint_key),
+            .accounts_index
-            |key, _| {
+            .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::<HashSet<Pubkey>>(),
+            });
+            // 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<Pubkey>, 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<Pubkey>, 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

runtime/src/accounts_index.rs

@@ -74,7 +74,32 @@ pub enum AccountIndex {
     SplTokenOwner,
 }
 
-pub type AccountSecondaryIndexes = HashSet<AccountIndex>;
+#[derive(Debug, PartialEq, Eq, Clone)]
+pub struct AccountSecondaryIndexesIncludeExclude {
+    pub exclude: bool,
+    pub keys: HashSet<Pubkey>,
+}
+
+#[derive(Debug, Default, Clone)]
+pub struct AccountSecondaryIndexes {
+    pub keys: Option<AccountSecondaryIndexesIncludeExclude>,
+    pub indexes: HashSet<AccountIndex>,
+}
+
+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<T> {
@@ -1063,7 +1088,9 @@ impl<T: 'static + Clone + IsCached + ZeroLamport> AccountsIndex<T> {
             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,18 +1114,22 @@ impl<T: 'static + Clone + IsCached + ZeroLamport> AccountsIndex<T> {
                     &account_data[SPL_TOKEN_ACCOUNT_OWNER_OFFSET
                         ..SPL_TOKEN_ACCOUNT_OWNER_OFFSET + PUBKEY_BYTES],
                 );
+                if account_indexes.include_key(&owner_key) {
                     self.spl_token_owner_index.insert(&owner_key, pubkey, slot);
                 }
+            }
 
             if account_indexes.contains(&AccountIndex::SplTokenMint) {
                 let mint_key = Pubkey::new(
                     &account_data[SPL_TOKEN_ACCOUNT_MINT_OFFSET
                         ..SPL_TOKEN_ACCOUNT_MINT_OFFSET + PUBKEY_BYTES],
                 );
+                if account_indexes.include_key(&mint_key) {
                     self.spl_token_mint_index.insert(&mint_key, pubkey, slot);
                 }
             }
         }
+    }
 
     // Same functionally to upsert, but doesn't take the read lock
     // initially on the accounts_map
@@ -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::<HashSet<_>>(),
+            exclude,
+        });
+        assert!(index.include_key(&pk1));
+        assert!(!index.include_key(&pk2));
+
+        let exclude = true;
+        index.keys = Some(AccountSecondaryIndexesIncludeExclude {
+            keys: [pk1].iter().cloned().collect::<HashSet<_>>(),
+            exclude,
+        });
+        assert!(!index.include_key(&pk1));
+        assert!(index.include_key(&pk2));
+
+        let exclude = true;
+        index.keys = Some(AccountSecondaryIndexesIncludeExclude {
+            keys: [pk1, pk2].iter().cloned().collect::<HashSet<_>>(),
+            exclude,
+        });
+        assert!(!index.include_key(&pk1));
+        assert!(!index.include_key(&pk2));
+
+        let exclude = false;
+        index.keys = Some(AccountSecondaryIndexesIncludeExclude {
+            keys: [pk1, pk2].iter().cloned().collect::<HashSet<_>>(),
+            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::<HashSet<_>>(),
+            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::<HashSet<_>>(),
+            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::<HashSet<_>>(),
+            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::<HashSet<_>>(),
+            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()

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,

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<AccountIndex> = matches
+    let account_indexes = process_account_indexes(&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 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<AccountIndex> = 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<Pubkey> =
+        values_t!(matches, "account_index_include_key", Pubkey)
+            .unwrap_or_default()
+            .iter()
+            .cloned()
+            .collect();
+
+    let account_indexes_exclude_keys: HashSet<Pubkey> =
+        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,
+    }
+}