From a102453baedc9895249efbc32ed72f3c42329162 Mon Sep 17 00:00:00 2001 From: Michael Vines Date: Wed, 16 Feb 2022 16:23:25 -0800 Subject: [PATCH] accounts_index: Add SPL Token account indexing for token-2022 accounts (#23043) --- Cargo.lock | 1 + programs/bpf/Cargo.lock | 1 + runtime/Cargo.toml | 1 + runtime/src/accounts_db.rs | 3 +- runtime/src/accounts_index.rs | 178 ++++++++++++++++----------- runtime/src/inline_spl_token.rs | 55 +++++++-- runtime/src/inline_spl_token_2022.rs | 18 +++ runtime/src/lib.rs | 1 + 8 files changed, 174 insertions(+), 84 deletions(-) create mode 100644 runtime/src/inline_spl_token_2022.rs diff --git a/Cargo.lock b/Cargo.lock index 89eca905f2..c58951bf77 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5897,6 +5897,7 @@ dependencies = [ "bincode", "blake3 1.3.0", "bv", + "bytemuck", "byteorder", "bzip2", "crossbeam-channel", diff --git a/programs/bpf/Cargo.lock b/programs/bpf/Cargo.lock index 9a9a3359f4..e09b3a4961 100644 --- a/programs/bpf/Cargo.lock +++ b/programs/bpf/Cargo.lock @@ -5244,6 +5244,7 @@ dependencies = [ "bincode", "blake3 1.3.0", "bv", + "bytemuck", "byteorder 1.4.3", "bzip2", "crossbeam-channel", diff --git a/runtime/Cargo.toml b/runtime/Cargo.toml index 0320c375a4..491806fa5a 100644 --- a/runtime/Cargo.toml +++ b/runtime/Cargo.toml @@ -14,6 +14,7 @@ arrayref = "0.3.6" bincode = "1.3.3" blake3 = "1.3.0" bv = { version = "0.11.1", features = ["serde"] } +bytemuck = "1.7.2" byteorder = "1.4.3" bzip2 = "0.4.3" dashmap = { version = "4.0.2", features = ["rayon", "raw-api"] } diff --git a/runtime/src/accounts_db.rs b/runtime/src/accounts_db.rs index da82cc4c86..f361d9a8cc 100644 --- a/runtime/src/accounts_db.rs +++ b/runtime/src/accounts_db.rs @@ -9015,8 +9015,7 @@ pub mod tests { // Set up account to be added to secondary index let mint_key = Pubkey::new_unique(); - let mut account_data_with_mint = - vec![0; inline_spl_token::state::Account::get_packed_len()]; + let mut account_data_with_mint = vec![0; inline_spl_token::Account::get_packed_len()]; account_data_with_mint[..PUBKEY_BYTES].clone_from_slice(&(mint_key.to_bytes())); let mut normal_account = AccountSharedData::new(1, 0, AccountSharedData::default().owner()); diff --git a/runtime/src/accounts_index.rs b/runtime/src/accounts_index.rs index 3b38ef9180..8dfc03a58a 100644 --- a/runtime/src/accounts_index.rs +++ b/runtime/src/accounts_index.rs @@ -5,7 +5,8 @@ use { bucket_map_holder::{Age, BucketMapHolder}, contains::Contains, in_mem_accounts_index::{InMemAccountsIndex, InsertNewEntryResults}, - inline_spl_token::{self, SPL_TOKEN_ACCOUNT_MINT_OFFSET, SPL_TOKEN_ACCOUNT_OWNER_OFFSET}, + inline_spl_token::{self, GenericTokenAccount}, + inline_spl_token_2022, pubkey_bins::PubkeyBinCalculator24, secondary_index::*, }, @@ -20,7 +21,7 @@ use { solana_measure::measure::Measure, solana_sdk::{ clock::{BankId, Slot}, - pubkey::{Pubkey, PUBKEY_BYTES}, + pubkey::Pubkey, }, std::{ collections::{btree_map::BTreeMap, HashSet}, @@ -1573,6 +1574,33 @@ impl AccountsIndex { max_root } + fn update_spl_token_secondary_indexes( + &self, + token_id: &Pubkey, + pubkey: &Pubkey, + account_owner: &Pubkey, + account_data: &[u8], + account_indexes: &AccountSecondaryIndexes, + ) { + if *account_owner == *token_id { + if account_indexes.contains(&AccountIndex::SplTokenOwner) { + if let Some(owner_key) = G::unpack_account_owner(account_data) { + if account_indexes.include_key(owner_key) { + self.spl_token_owner_index.insert(owner_key, pubkey); + } + } + } + + if account_indexes.contains(&AccountIndex::SplTokenMint) { + if let Some(mint_key) = G::unpack_account_mint(account_data) { + if account_indexes.include_key(mint_key) { + self.spl_token_mint_index.insert(mint_key, pubkey); + } + } + } + } + } + pub(crate) fn update_secondary_indexes( &self, pubkey: &Pubkey, @@ -1602,29 +1630,21 @@ impl AccountsIndex { // 2) When the fetch from storage occurs, it will return AccountSharedData::Default // (as persisted tombstone for snapshots). This will then ultimately be // filtered out by post-scan filters, like in `get_filtered_spl_token_accounts_by_owner()`. - if *account_owner == inline_spl_token::id() - && account_data.len() == inline_spl_token::state::Account::get_packed_len() - { - if account_indexes.contains(&AccountIndex::SplTokenOwner) { - let owner_key = Pubkey::new( - &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); - } - } - 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); - } - } - } + self.update_spl_token_secondary_indexes::( + &inline_spl_token::id(), + pubkey, + account_owner, + account_data, + account_indexes, + ); + self.update_spl_token_secondary_indexes::( + &inline_spl_token_2022::id(), + pubkey, + account_owner, + account_data, + account_indexes, + ); } fn get_account_maps_write_lock(&self, pubkey: &Pubkey) -> AccountMapsWriteLock { @@ -2010,7 +2030,11 @@ impl AccountsIndex { pub mod tests { use { super::*, - solana_sdk::signature::{Keypair, Signer}, + crate::inline_spl_token::*, + solana_sdk::{ + pubkey::PUBKEY_BYTES, + signature::{Keypair, Signer}, + }, std::ops::RangeInclusive, }; @@ -3840,7 +3864,7 @@ pub mod tests { let index_key = Pubkey::new_unique(); let account_key = Pubkey::new_unique(); - let mut account_data = vec![0; inline_spl_token::state::Account::get_packed_len()]; + let mut account_data = vec![0; inline_spl_token::Account::get_packed_len()]; account_data[key_start..key_end].clone_from_slice(&(index_key.to_bytes())); // Insert slots into secondary index @@ -4003,9 +4027,10 @@ pub mod tests { ); } - fn run_test_secondary_indexes< + fn run_test_spl_token_secondary_indexes< SecondaryIndexEntryType: SecondaryIndexEntry + Default + Sync + Send, >( + token_id: &Pubkey, index: &AccountsIndex, secondary_index: &SecondaryIndex, key_start: usize, @@ -4015,7 +4040,7 @@ pub mod tests { let mut secondary_indexes = secondary_indexes.clone(); let account_key = Pubkey::new_unique(); let index_key = Pubkey::new_unique(); - let mut account_data = vec![0; inline_spl_token::state::Account::get_packed_len()]; + let mut account_data = vec![0; inline_spl_token::Account::get_packed_len()]; account_data[key_start..key_end].clone_from_slice(&(index_key.to_bytes())); // Wrong program id @@ -4036,7 +4061,7 @@ pub mod tests { index.upsert( 0, &account_key, - &inline_spl_token::id(), + token_id, &account_data[1..], &secondary_indexes, true, @@ -4052,7 +4077,7 @@ pub mod tests { for _ in 0..2 { index.update_secondary_indexes( &account_key, - &inline_spl_token::id(), + token_id, &account_data, &secondary_indexes, ); @@ -4069,12 +4094,7 @@ pub mod tests { }); secondary_index.index.clear(); secondary_index.reverse_index.clear(); - index.update_secondary_indexes( - &account_key, - &inline_spl_token::id(), - &account_data, - &secondary_indexes, - ); + index.update_secondary_indexes(&account_key, token_id, &account_data, &secondary_indexes); assert!(!secondary_index.index.is_empty()); assert!(!secondary_index.reverse_index.is_empty()); check_secondary_index_mapping_correct(secondary_index, &[index_key], &account_key); @@ -4086,12 +4106,7 @@ pub mod tests { }); secondary_index.index.clear(); secondary_index.reverse_index.clear(); - index.update_secondary_indexes( - &account_key, - &inline_spl_token::id(), - &account_data, - &secondary_indexes, - ); + index.update_secondary_indexes(&account_key, token_id, &account_data, &secondary_indexes); assert!(!secondary_index.index.is_empty()); assert!(!secondary_index.reverse_index.is_empty()); check_secondary_index_mapping_correct(secondary_index, &[index_key], &account_key); @@ -4110,31 +4125,38 @@ pub mod tests { fn test_dashmap_secondary_index() { let (key_start, key_end, secondary_indexes) = create_dashmap_secondary_index_state(); let index = AccountsIndex::::default_for_tests(); - run_test_secondary_indexes( - &index, - &index.spl_token_mint_index, - key_start, - key_end, - &secondary_indexes, - ); + for token_id in [inline_spl_token::id(), inline_spl_token_2022::id()] { + run_test_spl_token_secondary_indexes( + &token_id, + &index, + &index.spl_token_mint_index, + key_start, + key_end, + &secondary_indexes, + ); + } } #[test] fn test_rwlock_secondary_index() { let (key_start, key_end, secondary_indexes) = create_rwlock_secondary_index_state(); let index = AccountsIndex::::default_for_tests(); - run_test_secondary_indexes( - &index, - &index.spl_token_owner_index, - key_start, - key_end, - &secondary_indexes, - ); + for token_id in [inline_spl_token::id(), inline_spl_token_2022::id()] { + run_test_spl_token_secondary_indexes( + &token_id, + &index, + &index.spl_token_owner_index, + key_start, + key_end, + &secondary_indexes, + ); + } } fn run_test_secondary_indexes_same_slot_and_forks< SecondaryIndexEntryType: SecondaryIndexEntry + Default + Sync + Send, >( + token_id: &Pubkey, index: &AccountsIndex, secondary_index: &SecondaryIndex, index_key_start: usize, @@ -4145,10 +4167,10 @@ pub mod tests { let secondary_key1 = Pubkey::new_unique(); let secondary_key2 = Pubkey::new_unique(); let slot = 1; - let mut account_data1 = vec![0; inline_spl_token::state::Account::get_packed_len()]; + let mut account_data1 = vec![0; inline_spl_token::Account::get_packed_len()]; account_data1[index_key_start..index_key_end] .clone_from_slice(&(secondary_key1.to_bytes())); - let mut account_data2 = vec![0; inline_spl_token::state::Account::get_packed_len()]; + let mut account_data2 = vec![0; inline_spl_token::Account::get_packed_len()]; account_data2[index_key_start..index_key_end] .clone_from_slice(&(secondary_key2.to_bytes())); @@ -4156,7 +4178,7 @@ pub mod tests { index.upsert( slot, &account_key, - &inline_spl_token::id(), + token_id, &account_data1, secondary_indexes, true, @@ -4168,7 +4190,7 @@ pub mod tests { index.upsert( slot, &account_key, - &inline_spl_token::id(), + token_id, &account_data2, secondary_indexes, true, @@ -4188,7 +4210,7 @@ pub mod tests { index.upsert( later_slot, &account_key, - &inline_spl_token::id(), + token_id, &account_data1, secondary_indexes, true, @@ -4224,26 +4246,32 @@ pub mod tests { fn test_dashmap_secondary_index_same_slot_and_forks() { let (key_start, key_end, account_index) = create_dashmap_secondary_index_state(); let index = AccountsIndex::::default_for_tests(); - run_test_secondary_indexes_same_slot_and_forks( - &index, - &index.spl_token_mint_index, - key_start, - key_end, - &account_index, - ); + for token_id in [inline_spl_token::id(), inline_spl_token_2022::id()] { + run_test_secondary_indexes_same_slot_and_forks( + &token_id, + &index, + &index.spl_token_mint_index, + key_start, + key_end, + &account_index, + ); + } } #[test] fn test_rwlock_secondary_index_same_slot_and_forks() { let (key_start, key_end, account_index) = create_rwlock_secondary_index_state(); let index = AccountsIndex::::default_for_tests(); - run_test_secondary_indexes_same_slot_and_forks( - &index, - &index.spl_token_owner_index, - key_start, - key_end, - &account_index, - ); + for token_id in [inline_spl_token::id(), inline_spl_token_2022::id()] { + run_test_secondary_indexes_same_slot_and_forks( + &token_id, + &index, + &index.spl_token_owner_index, + key_start, + key_end, + &account_index, + ); + } } impl IndexValue for bool {} diff --git a/runtime/src/inline_spl_token.rs b/runtime/src/inline_spl_token.rs index 4e788fa1da..80d23a6265 100644 --- a/runtime/src/inline_spl_token.rs +++ b/runtime/src/inline_spl_token.rs @@ -1,4 +1,6 @@ -// Partial SPL Token declarations inlined to avoid an external dependency on the spl-token crate +/// Partial SPL Token declarations inlined to avoid an external dependency on the spl-token crate +use solana_sdk::pubkey::{Pubkey, PUBKEY_BYTES}; + solana_sdk::declare_id!("TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA"); pub(crate) mod new_token_program { @@ -19,15 +21,54 @@ pub(crate) mod new_token_program { */ pub const SPL_TOKEN_ACCOUNT_MINT_OFFSET: usize = 0; pub const SPL_TOKEN_ACCOUNT_OWNER_OFFSET: usize = 32; +const SPL_TOKEN_ACCOUNT_LENGTH: usize = 165; -pub mod state { - const LEN: usize = 165; - pub struct Account; - impl Account { - pub fn get_packed_len() -> usize { - LEN +pub(crate) trait GenericTokenAccount { + fn valid_account_data(account_data: &[u8]) -> bool; + + // Call after account length has already been verified + fn unpack_account_owner_unchecked(account_data: &[u8]) -> &Pubkey { + Self::unpack_pubkey_unchecked(account_data, SPL_TOKEN_ACCOUNT_OWNER_OFFSET) + } + + // Call after account length has already been verified + fn unpack_account_mint_unchecked(account_data: &[u8]) -> &Pubkey { + Self::unpack_pubkey_unchecked(account_data, SPL_TOKEN_ACCOUNT_MINT_OFFSET) + } + + // Call after account length has already been verified + fn unpack_pubkey_unchecked(account_data: &[u8], offset: usize) -> &Pubkey { + bytemuck::from_bytes(&account_data[offset..offset + PUBKEY_BYTES]) + } + + fn unpack_account_owner(account_data: &[u8]) -> Option<&Pubkey> { + if Self::valid_account_data(account_data) { + Some(Self::unpack_account_owner_unchecked(account_data)) + } else { + None } } + + fn unpack_account_mint(account_data: &[u8]) -> Option<&Pubkey> { + if Self::valid_account_data(account_data) { + Some(Self::unpack_account_mint_unchecked(account_data)) + } else { + None + } + } +} + +pub struct Account; +impl Account { + pub fn get_packed_len() -> usize { + SPL_TOKEN_ACCOUNT_LENGTH + } +} + +impl GenericTokenAccount for Account { + fn valid_account_data(account_data: &[u8]) -> bool { + account_data.len() == SPL_TOKEN_ACCOUNT_LENGTH + } } pub mod native_mint { diff --git a/runtime/src/inline_spl_token_2022.rs b/runtime/src/inline_spl_token_2022.rs new file mode 100644 index 0000000000..3aed652b96 --- /dev/null +++ b/runtime/src/inline_spl_token_2022.rs @@ -0,0 +1,18 @@ +/// Partial SPL Token declarations inlined to avoid an external dependency on the spl-token-2022 crate +use crate::inline_spl_token::{self, GenericTokenAccount}; + +solana_sdk::declare_id!("TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb"); + +// `spl_token_program_2022::extension::AccountType::Account` ordinal value +const ACCOUNTTYPE_ACCOUNT: u8 = 2; + +pub struct Account; +impl GenericTokenAccount for Account { + fn valid_account_data(account_data: &[u8]) -> bool { + inline_spl_token::Account::valid_account_data(account_data) + || ACCOUNTTYPE_ACCOUNT + == *account_data + .get(inline_spl_token::Account::get_packed_len()) + .unwrap_or(&0) + } +} diff --git a/runtime/src/lib.rs b/runtime/src/lib.rs index 5e294cb576..77d7f555f1 100644 --- a/runtime/src/lib.rs +++ b/runtime/src/lib.rs @@ -35,6 +35,7 @@ pub mod hardened_unpack; pub mod in_mem_accounts_index; pub mod inline_spl_associated_token_account; pub mod inline_spl_token; +pub mod inline_spl_token_2022; pub mod loader_utils; pub mod message_processor; pub mod non_circulating_supply;