From 99232be729fae90cd09e5a14955236e6924cafda Mon Sep 17 00:00:00 2001 From: Yueh-Hsuan Chiang Date: Fri, 19 Jan 2024 01:58:26 -0800 Subject: [PATCH] [TieredStorage] In-memory struct for writing OwnersBlock --- accounts-db/src/tiered_storage/hot.rs | 20 ++++-- accounts-db/src/tiered_storage/owners.rs | 80 ++++++++++++++++++++++-- 2 files changed, 91 insertions(+), 9 deletions(-) diff --git a/accounts-db/src/tiered_storage/hot.rs b/accounts-db/src/tiered_storage/hot.rs index ace6649ba26f49..ea79e4e964e4fc 100644 --- a/accounts-db/src/tiered_storage/hot.rs +++ b/accounts-db/src/tiered_storage/hot.rs @@ -470,7 +470,7 @@ pub mod tests { hot::{HotAccountMeta, HotStorageReader}, index::{AccountIndexWriterEntry, IndexBlockFormat, IndexOffset}, meta::{AccountMetaFlags, AccountMetaOptionalFields, TieredAccountMeta}, - owners::OwnersBlockFormat, + owners::{OwnersBlockFormat, OwnersTable}, }, assert_matches::assert_matches, memoffset::offset_of, @@ -823,9 +823,13 @@ pub mod tests { { let file = TieredStorageFile::new_writable(&path).unwrap(); + let mut owners_table = OwnersTable::new(); + addresses.iter().for_each(|owner_address| { + owners_table.check_and_add(owner_address); + }); footer .owners_block_format - .write_owners_block(&file, &addresses) + .write_owners_block(&file, owners_table.owners()) .unwrap(); // while the test only focuses on account metas, writing a footer @@ -893,9 +897,13 @@ pub mod tests { // the owners_block_offset set to the end of the accounts blocks. footer.owners_block_offset = footer.index_block_offset; + let mut owners_table = OwnersTable::new(); + owner_addresses.iter().for_each(|owner_address| { + owners_table.check_and_add(owner_address); + }); footer .owners_block_format - .write_owners_block(&file, &owner_addresses) + .write_owners_block(&file, owners_table.owners()) .unwrap(); // while the test only focuses on account metas, writing a footer @@ -1029,9 +1037,13 @@ pub mod tests { // write owners block footer.owners_block_offset = current_offset as u64; + let mut owners_table = OwnersTable::new(); + owners.iter().for_each(|owner_address| { + owners_table.check_and_add(owner_address); + }); footer .owners_block_format - .write_owners_block(&file, &owners) + .write_owners_block(&file, owners_table.owners()) .unwrap(); footer.write_footer_block(&file).unwrap(); diff --git a/accounts-db/src/tiered_storage/owners.rs b/accounts-db/src/tiered_storage/owners.rs index d8a963ce143401..1bbcefdc2a449e 100644 --- a/accounts-db/src/tiered_storage/owners.rs +++ b/accounts-db/src/tiered_storage/owners.rs @@ -5,6 +5,7 @@ use { }, memmap2::Mmap, solana_sdk::pubkey::Pubkey, + std::collections::HashMap, }; /// The offset to an owner entry in the owners block. @@ -40,16 +41,16 @@ pub enum OwnersBlockFormat { impl OwnersBlockFormat { /// Persists the provided owners' addresses into the specified file. - pub fn write_owners_block( + pub(crate) fn write_owners_block( &self, file: &TieredStorageFile, - addresses: &[Pubkey], + owners: &[&Pubkey], ) -> TieredStorageResult { match self { Self::AddressesOnly => { let mut bytes_written = 0; - for address in addresses { - bytes_written += file.write_pod(address)?; + for address in owners { + bytes_written += file.write_pod(*address)?; } Ok(bytes_written) @@ -77,6 +78,44 @@ impl OwnersBlockFormat { } } +/// The in-memory representation of owners block for write. +/// It manages a set of unique addresses of account owners. +#[derive(Debug)] +pub(crate) struct OwnersTable<'owner> { + owners_vec: Vec<&'owner Pubkey>, + owners_map: HashMap<&'owner Pubkey, OwnerOffset>, +} + +/// OwnersBlock is persisted as a consecutive bytes of pubkeys without any +/// meta-data. For each account meta, it has a owner_offset field to +/// access its owner's address in the OwnersBlock. +impl<'owner> OwnersTable<'owner> { + pub(crate) fn new() -> Self { + Self { + owners_vec: vec![], + owners_map: HashMap::new(), + } + } + + /// Add the specified pubkey as the owner into the OwnersWriterTable + /// if the specified pubkey has not existed in the OwnersWriterTable + /// yet. In any case, the function returns its OwnerOffset. + pub(crate) fn check_and_add(&mut self, pubkey: &'owner Pubkey) -> OwnerOffset { + if let Some(offset) = self.owners_map.get(pubkey) { + return *offset; + } + let offset: u32 = self.owners_vec.len().try_into().unwrap(); + self.owners_vec.push(pubkey); + self.owners_map.insert(pubkey, OwnerOffset(offset)); + + OwnerOffset(offset) + } + + pub(crate) fn owners(&self) -> &[&'owner Pubkey] { + &self.owners_vec + } +} + #[cfg(test)] mod tests { use { @@ -105,9 +144,13 @@ mod tests { { let file = TieredStorageFile::new_writable(&path).unwrap(); + let mut owners_table = OwnersTable::new(); + addresses.iter().for_each(|owner_address| { + owners_table.check_and_add(owner_address); + }); footer .owners_block_format - .write_owners_block(&file, &addresses) + .write_owners_block(&file, owners_table.owners()) .unwrap(); // while the test only focuses on account metas, writing a footer @@ -128,4 +171,31 @@ mod tests { ); } } + + #[test] + fn test_owners_table() { + let mut owners_table = OwnersTable::new(); + const NUM_OWNERS: usize = 99; + + let addresses: Vec<_> = std::iter::repeat_with(Pubkey::new_unique) + .take(NUM_OWNERS) + .collect(); + + // as we insert sequentially, we expect each entry has same OwnerOffset + // as its index inside the Vector. + for (i, address) in addresses.iter().enumerate() { + assert_eq!(owners_table.check_and_add(address), OwnerOffset(i as u32)); + } + + let cloned_addresses = addresses.clone(); + + // insert again and expect the same OwnerOffset + for (i, address) in cloned_addresses.iter().enumerate() { + assert_eq!(owners_table.check_and_add(address), OwnerOffset(i as u32)); + } + + // make sure the size of the resulting owner table is the same + // as the input + assert_eq!(owners_table.owners().len(), addresses.len()); + } }