From 08cc140d4a733d183130a3e6cde9a6135860fcb3 Mon Sep 17 00:00:00 2001 From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com> Date: Thu, 17 Feb 2022 02:20:38 +0000 Subject: [PATCH] accounts_index: Add SPL Token account indexing for token-2022 accounts (#23043) (#23203) (cherry picked from commit a102453baedc9895249efbc32ed72f3c42329162) Co-authored-by: Michael Vines --- 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 80eea424fa..3c979e31eb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5699,6 +5699,7 @@ dependencies = [ "bincode", "blake3 1.2.0", "bv", + "bytemuck", "byteorder", "bzip2", "crossbeam-channel", diff --git a/programs/bpf/Cargo.lock b/programs/bpf/Cargo.lock index 1d1d7a0721..010843611b 100644 --- a/programs/bpf/Cargo.lock +++ b/programs/bpf/Cargo.lock @@ -3271,6 +3271,7 @@ dependencies = [ "bincode", "blake3", "bv", + "bytemuck", "byteorder 1.4.3", "bzip2", "crossbeam-channel", diff --git a/runtime/Cargo.toml b/runtime/Cargo.toml index bf1fb53ae2..31253bf3e2 100644 --- a/runtime/Cargo.toml +++ b/runtime/Cargo.toml @@ -14,6 +14,7 @@ arrayref = "0.3.6" bincode = "1.3.3" blake3 = "1.2.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 b8e13e4720..07b32fc532 100644 --- a/runtime/src/accounts_db.rs +++ b/runtime/src/accounts_db.rs @@ -8992,8 +8992,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 4a71f07412..bb52193dac 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::*, }, @@ -16,7 +17,7 @@ use { solana_measure::measure::Measure, solana_sdk::{ clock::{BankId, Slot}, - pubkey::{Pubkey, PUBKEY_BYTES}, + pubkey::Pubkey, }, std::{ collections::{btree_map::BTreeMap, HashSet}, @@ -1551,6 +1552,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, @@ -1580,29 +1608,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 { @@ -1988,7 +2008,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, }; @@ -3779,7 +3803,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 @@ -3942,9 +3966,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, @@ -3954,7 +3979,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 @@ -3975,7 +4000,7 @@ pub mod tests { index.upsert( 0, &account_key, - &inline_spl_token::id(), + token_id, &account_data[1..], &secondary_indexes, true, @@ -3991,7 +4016,7 @@ pub mod tests { for _ in 0..2 { index.update_secondary_indexes( &account_key, - &inline_spl_token::id(), + token_id, &account_data, &secondary_indexes, ); @@ -4008,12 +4033,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); @@ -4025,12 +4045,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); @@ -4049,31 +4064,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, @@ -4084,10 +4106,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())); @@ -4095,7 +4117,7 @@ pub mod tests { index.upsert( slot, &account_key, - &inline_spl_token::id(), + token_id, &account_data1, secondary_indexes, true, @@ -4107,7 +4129,7 @@ pub mod tests { index.upsert( slot, &account_key, - &inline_spl_token::id(), + token_id, &account_data2, secondary_indexes, true, @@ -4127,7 +4149,7 @@ pub mod tests { index.upsert( later_slot, &account_key, - &inline_spl_token::id(), + token_id, &account_data1, secondary_indexes, true, @@ -4163,26 +4185,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 d2186c88a8..a1b22e91b4 100644 --- a/runtime/src/lib.rs +++ b/runtime/src/lib.rs @@ -33,6 +33,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;