Skip to content

Commit

Permalink
[TieredStorage] HotStorageReader::account_matches_owners (solana-labs…
Browse files Browse the repository at this point in the history
…#34350)

#### Problem
Each AccountsFile is required to implement account_matches_owners(),
a public API that checks whether the account located at the specified
offset matches any input owners.  However, the implementation of
account_matches_owners() for HotAccountsStorage is missing.

#### Summary of Changes
This PR implements HotStorageReader::account_matches_owners().

#### Test Plan
A new unit-test is added to this PR.
  • Loading branch information
yhchiang-sol authored Dec 22, 2023
1 parent 6bbd366 commit f73fd9d
Showing 1 changed file with 136 additions and 1 deletion.
137 changes: 136 additions & 1 deletion accounts-db/src/tiered_storage/hot.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
use {
crate::{
accounts_file::MatchAccountOwnerError,
accounts_hash::AccountHash,
tiered_storage::{
byte_block,
Expand Down Expand Up @@ -316,6 +317,38 @@ impl HotStorageReader {
fn get_owner_address(&self, owner_offset: OwnerOffset) -> TieredStorageResult<&Pubkey> {
OwnersBlock::get_owner_address(&self.mmap, &self.footer, owner_offset)
}

/// Returns Ok(index_of_matching_owner) if the account owner at
/// `account_offset` is one of the pubkeys in `owners`.
///
/// Returns Err(MatchAccountOwnerError::NoMatch) if the account has 0
/// lamports or the owner is not one of the pubkeys in `owners`.
///
/// Returns Err(MatchAccountOwnerError::UnableToLoad) if there is any internal
/// error that causes the data unable to load, including `account_offset`
/// causes a data overrun.
pub fn account_matches_owners(
&self,
account_offset: HotAccountOffset,
owners: &[&Pubkey],
) -> Result<usize, MatchAccountOwnerError> {
let account_meta = self
.get_account_meta_from_offset(account_offset)
.map_err(|_| MatchAccountOwnerError::UnableToLoad)?;

if account_meta.lamports() == 0 {
Err(MatchAccountOwnerError::NoMatch)
} else {
let account_owner = self
.get_owner_address(account_meta.owner_offset())
.map_err(|_| MatchAccountOwnerError::UnableToLoad)?;

owners
.iter()
.position(|candidate| &account_owner == candidate)
.ok_or(MatchAccountOwnerError::NoMatch)
}
}
}

#[cfg(test)]
Expand All @@ -335,7 +368,7 @@ pub mod tests {
},
assert_matches::assert_matches,
memoffset::offset_of,
rand::Rng,
rand::{seq::SliceRandom, Rng},
solana_sdk::{hash::Hash, pubkey::Pubkey, stake_history::Epoch},
tempfile::TempDir,
};
Expand Down Expand Up @@ -702,4 +735,106 @@ pub mod tests {
);
}
}

#[test]
fn test_account_matches_owners() {
// Generate a new temp path that is guaranteed to NOT already have a file.
let temp_dir = TempDir::new().unwrap();
let path = temp_dir.path().join("test_hot_storage_get_owner_address");
const NUM_OWNERS: u32 = 10;

let owner_addresses: Vec<_> = std::iter::repeat_with(Pubkey::new_unique)
.take(NUM_OWNERS as usize)
.collect();

const NUM_ACCOUNTS: u32 = 30;
let mut rng = rand::thread_rng();

let hot_account_metas: Vec<_> = std::iter::repeat_with({
|| {
HotAccountMeta::new()
.with_lamports(rng.gen_range(1..u64::MAX))
.with_owner_offset(OwnerOffset(rng.gen_range(0..NUM_OWNERS)))
}
})
.take(NUM_ACCOUNTS as usize)
.collect();
let mut footer = TieredStorageFooter {
account_meta_format: AccountMetaFormat::Hot,
account_entry_count: NUM_ACCOUNTS,
owner_count: NUM_OWNERS,
..TieredStorageFooter::default()
};
let account_offsets: Vec<_>;

{
let file = TieredStorageFile::new_writable(&path).unwrap();
let mut current_offset = 0;

account_offsets = hot_account_metas
.iter()
.map(|meta| {
let prev_offset = current_offset;
current_offset += file.write_pod(meta).unwrap();
HotAccountOffset::new(prev_offset).unwrap()
})
.collect();
footer.index_block_offset = current_offset as u64;
// Typically, the owners block is stored after index block, but
// since we don't write index block in this test, so we have
// the owners_block_offset set to the end of the accounts blocks.
footer.owners_block_offset = footer.index_block_offset;

OwnersBlock::write_owners_block(&file, &owner_addresses).unwrap();

// while the test only focuses on account metas, writing a footer
// here is necessary to make it a valid tiered-storage file.
footer.write_footer_block(&file).unwrap();
}

let hot_storage = HotStorageReader::new_from_path(&path).unwrap();

// First, verify whether we can find the expected owners.
let mut owner_candidates: Vec<_> = owner_addresses.iter().collect();
owner_candidates.shuffle(&mut rng);

for (account_offset, account_meta) in account_offsets.iter().zip(hot_account_metas.iter()) {
let index = hot_storage
.account_matches_owners(*account_offset, &owner_candidates)
.unwrap();
assert_eq!(
owner_candidates[index],
&owner_addresses[account_meta.owner_offset().0 as usize]
);
}

// Second, verify the MatchAccountOwnerError::NoMatch case
const NUM_UNMATCHED_OWNERS: usize = 20;
let unmatched_owners: Vec<_> = std::iter::repeat_with(Pubkey::new_unique)
.take(NUM_UNMATCHED_OWNERS)
.collect();
let unmatched_candidates: Vec<_> = unmatched_owners.iter().collect();

for account_offset in account_offsets.iter() {
assert_eq!(
hot_storage.account_matches_owners(*account_offset, &unmatched_candidates),
Err(MatchAccountOwnerError::NoMatch)
);
}

// Thirdly, we mixed two candidates and make sure we still find the
// matched owner.
owner_candidates.extend(unmatched_candidates);
owner_candidates.shuffle(&mut rng);

for (account_offset, account_meta) in account_offsets.iter().zip(hot_account_metas.iter()) {
let index = hot_storage
.account_matches_owners(*account_offset, &owner_candidates)
.unwrap();
assert_eq!(
owner_candidates[index],
&owner_addresses[account_meta.owner_offset().0 as usize]
);
}
}
}

0 comments on commit f73fd9d

Please sign in to comment.