accounts_index: Add SPL Token account indexing for token-2022 accounts (#23043) (#23203)

(cherry picked from commit a102453baedc9895249efbc32ed72f3c42329162)

Co-authored-by: Michael Vines <mvines@gmail.com>
This commit is contained in:
mergify[bot] 2022-02-17 02:20:38 +00:00 committed by GitHub
parent 2120ef5808
commit 08cc140d4a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
8 changed files with 174 additions and 84 deletions

1
Cargo.lock generated
View File

@ -5699,6 +5699,7 @@ dependencies = [
"bincode",
"blake3 1.2.0",
"bv",
"bytemuck",
"byteorder",
"bzip2",
"crossbeam-channel",

View File

@ -3271,6 +3271,7 @@ dependencies = [
"bincode",
"blake3",
"bv",
"bytemuck",
"byteorder 1.4.3",
"bzip2",
"crossbeam-channel",

View File

@ -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"] }

View File

@ -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());

View File

@ -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<T: IndexValue> AccountsIndex<T> {
max_root
}
fn update_spl_token_secondary_indexes<G: GenericTokenAccount>(
&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<T: IndexValue> AccountsIndex<T> {
// 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::Account>(
&inline_spl_token::id(),
pubkey,
account_owner,
account_data,
account_indexes,
);
self.update_spl_token_secondary_indexes::<inline_spl_token_2022::Account>(
&inline_spl_token_2022::id(),
pubkey,
account_owner,
account_data,
account_indexes,
);
}
fn get_account_maps_write_lock(&self, pubkey: &Pubkey) -> AccountMapsWriteLock<T> {
@ -1988,7 +2008,11 @@ impl<T: IndexValue> AccountsIndex<T> {
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<bool>,
secondary_index: &SecondaryIndex<SecondaryIndexEntryType>,
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::<bool>::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::<bool>::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<bool>,
secondary_index: &SecondaryIndex<SecondaryIndexEntryType>,
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::<bool>::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::<bool>::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 {}

View File

@ -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 {

View File

@ -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)
}
}

View File

@ -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;