[TieredStorage] Code refactoring for OwnersBlock (solana-labs#34854)

#### Problem The OwnersBlockFormat is currently defined inside footer.rs instead of inside owners.rs. In addition, the implementation of OwnersBlock doesn't honor OwnersBlockFormat. #### Summary of Changes This PR moves OwnersBlockFormat from footer.rs to owners.rs and repurpose OwnersBlock as OwnersBlockFormat (just like the IndexBlockFormat inside index.rs) #### Test Plan Existing unit-tests.
yhchiang-sol · Jan 19, 2024 · 81b3210 · 81b3210
1 parent 39687a0
commit 81b3210
Show file tree

Hide file tree

Showing 4 changed files with 76 additions and 54 deletions.
diff --git a/accounts-db/src/tiered_storage.rs b/accounts-db/src/tiered_storage.rs
@@ -19,8 +19,9 @@ use {
         storable_accounts::StorableAccounts,
     },
     error::TieredStorageError,
-    footer::{AccountBlockFormat, AccountMetaFormat, OwnersBlockFormat},
+    footer::{AccountBlockFormat, AccountMetaFormat},
     index::IndexBlockFormat,
+    owners::OwnersBlockFormat,
     readable::TieredStorageReader,
     solana_sdk::account::ReadableAccount,
     std::{

diff --git a/accounts-db/src/tiered_storage/footer.rs b/accounts-db/src/tiered_storage/footer.rs
@@ -4,6 +4,7 @@ use {
         file::TieredStorageFile,
         index::IndexBlockFormat,
         mmap_utils::{get_pod, get_type},
+        owners::OwnersBlockFormat,
         TieredStorageResult,
     },
     bytemuck::{Pod, Zeroable},
@@ -78,23 +79,6 @@ pub enum AccountBlockFormat {
     Lz4 = 1,
 }
 
-#[repr(u16)]
-#[derive(
-    Clone,
-    Copy,
-    Debug,
-    Default,
-    Eq,
-    Hash,
-    PartialEq,
-    num_enum::IntoPrimitive,
-    num_enum::TryFromPrimitive,
-)]
-pub enum OwnersBlockFormat {
-    #[default]
-    LocalIndex = 0,
-}
-
 #[derive(Debug, PartialEq, Eq, Clone, Copy)]
 #[repr(C)]
 pub struct TieredStorageFooter {
@@ -353,7 +337,7 @@ mod tests {
         let path = get_append_vec_path("test_file_footer");
         let expected_footer = TieredStorageFooter {
             account_meta_format: AccountMetaFormat::Hot,
-            owners_block_format: OwnersBlockFormat::LocalIndex,
+            owners_block_format: OwnersBlockFormat::AddressesOnly,
             index_block_format: IndexBlockFormat::AddressesThenOffsets,
             account_block_format: AccountBlockFormat::AlignedRaw,
             account_entry_count: 300,

diff --git a/accounts-db/src/tiered_storage/hot.rs b/accounts-db/src/tiered_storage/hot.rs
@@ -8,13 +8,11 @@ use {
         tiered_storage::{
             byte_block,
             file::TieredStorageFile,
-            footer::{
-                AccountBlockFormat, AccountMetaFormat, OwnersBlockFormat, TieredStorageFooter,
-            },
+            footer::{AccountBlockFormat, AccountMetaFormat, TieredStorageFooter},
             index::{AccountOffset, IndexBlockFormat, IndexOffset},
             meta::{AccountMetaFlags, AccountMetaOptionalFields, TieredAccountMeta},
             mmap_utils::{get_pod, get_slice},
-            owners::{OwnerOffset, OwnersBlock},
+            owners::{OwnerOffset, OwnersBlockFormat},
             readable::TieredReadableAccount,
             TieredStorageError, TieredStorageFormat, TieredStorageResult,
         },
@@ -29,7 +27,7 @@ use {
 pub const HOT_FORMAT: TieredStorageFormat = TieredStorageFormat {
     meta_entry_size: std::mem::size_of::<HotAccountMeta>(),
     account_meta_format: AccountMetaFormat::Hot,
-    owners_block_format: OwnersBlockFormat::LocalIndex,
+    owners_block_format: OwnersBlockFormat::AddressesOnly,
     index_block_format: IndexBlockFormat::AddressesThenOffsets,
     account_block_format: AccountBlockFormat::AlignedRaw,
 };
@@ -331,7 +329,9 @@ impl HotStorageReader {
     /// Returns the address of the account owner given the specified
     /// owner_offset.
     fn get_owner_address(&self, owner_offset: OwnerOffset) -> TieredStorageResult<&Pubkey> {
-        OwnersBlock::get_owner_address(&self.mmap, &self.footer, owner_offset)
+        self.footer
+            .owners_block_format
+            .get_owner_address(&self.mmap, &self.footer, owner_offset)
     }
 
     /// Returns Ok(index_of_matching_owner) if the account owner at
@@ -466,13 +466,11 @@ pub mod tests {
         crate::tiered_storage::{
             byte_block::ByteBlockWriter,
             file::TieredStorageFile,
-            footer::{
-                AccountBlockFormat, AccountMetaFormat, OwnersBlockFormat, TieredStorageFooter,
-                FOOTER_SIZE,
-            },
+            footer::{AccountBlockFormat, AccountMetaFormat, TieredStorageFooter, FOOTER_SIZE},
             hot::{HotAccountMeta, HotStorageReader},
             index::{AccountIndexWriterEntry, IndexBlockFormat, IndexOffset},
             meta::{AccountMetaFlags, AccountMetaOptionalFields, TieredAccountMeta},
+            owners::OwnersBlockFormat,
         },
         assert_matches::assert_matches,
         memoffset::offset_of,
@@ -638,7 +636,7 @@ pub mod tests {
         let path = temp_dir.path().join("test_hot_storage_footer");
         let expected_footer = TieredStorageFooter {
             account_meta_format: AccountMetaFormat::Hot,
-            owners_block_format: OwnersBlockFormat::LocalIndex,
+            owners_block_format: OwnersBlockFormat::AddressesOnly,
             index_block_format: IndexBlockFormat::AddressesThenOffsets,
             account_block_format: AccountBlockFormat::AlignedRaw,
             account_entry_count: 300,
@@ -825,7 +823,10 @@ pub mod tests {
         {
             let file = TieredStorageFile::new_writable(&path).unwrap();
 
-            OwnersBlock::write_owners_block(&file, &addresses).unwrap();
+            footer
+                .owners_block_format
+                .write_owners_block(&file, &addresses)
+                .unwrap();
 
             // while the test only focuses on account metas, writing a footer
             // here is necessary to make it a valid tiered-storage file.
@@ -892,7 +893,10 @@ pub mod tests {
             // 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();
+            footer
+                .owners_block_format
+                .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.
@@ -1025,7 +1029,10 @@ pub mod tests {
 
             // write owners block
             footer.owners_block_offset = current_offset as u64;
-            OwnersBlock::write_owners_block(&file, &owners).unwrap();
+            footer
+                .owners_block_format
+                .write_owners_block(&file, &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
@@ -7,12 +7,6 @@ use {
     solana_sdk::pubkey::Pubkey,
 };
 
-/// Owner block holds a set of unique addresses of account owners,
-/// and an account meta has a owner_offset field for accessing
-/// it's owner address.
-#[derive(Debug)]
-pub struct OwnersBlock;
-
 /// The offset to an owner entry in the owners block.
 /// This is used to obtain the address of the account owner.
 ///
@@ -21,35 +15,65 @@ pub struct OwnersBlock;
 #[derive(Clone, Copy, Debug, Eq, PartialEq, PartialOrd)]
 pub struct OwnerOffset(pub u32);
 
-/// 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 OwnersBlock {
+/// Owner block holds a set of unique addresses of account owners,
+/// and an account meta has a owner_offset field for accessing
+/// it's owner address.
+#[repr(u16)]
+#[derive(
+    Clone,
+    Copy,
+    Debug,
+    Default,
+    Eq,
+    Hash,
+    PartialEq,
+    num_enum::IntoPrimitive,
+    num_enum::TryFromPrimitive,
+)]
+pub enum OwnersBlockFormat {
+    /// This format persists OwnerBlock 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.
+    #[default]
+    AddressesOnly = 0,
+}
+
+impl OwnersBlockFormat {
     /// Persists the provided owners' addresses into the specified file.
     pub fn write_owners_block(
+        &self,
         file: &TieredStorageFile,
         addresses: &[Pubkey],
     ) -> TieredStorageResult<usize> {
-        let mut bytes_written = 0;
-        for address in addresses {
-            bytes_written += file.write_pod(address)?;
-        }
+        match self {
+            Self::AddressesOnly => {
+                let mut bytes_written = 0;
+                for address in addresses {
+                    bytes_written += file.write_pod(address)?;
+                }
 
-        Ok(bytes_written)
+                Ok(bytes_written)
+            }
+        }
     }
 
     /// Returns the owner address associated with the specified owner_offset
     /// and footer inside the input mmap.
     pub fn get_owner_address<'a>(
+        &self,
         mmap: &'a Mmap,
         footer: &TieredStorageFooter,
         owner_offset: OwnerOffset,
     ) -> TieredStorageResult<&'a Pubkey> {
-        let offset = footer.owners_block_offset as usize
-            + (std::mem::size_of::<Pubkey>() * owner_offset.0 as usize);
-        let (pubkey, _) = get_pod::<Pubkey>(mmap, offset)?;
+        match self {
+            Self::AddressesOnly => {
+                let offset = footer.owners_block_offset as usize
+                    + (std::mem::size_of::<Pubkey>() * owner_offset.0 as usize);
+                let (pubkey, _) = get_pod::<Pubkey>(mmap, offset)?;
 
-        Ok(pubkey)
+                Ok(pubkey)
+            }
+        }
     }
 }
 
@@ -81,7 +105,10 @@ mod tests {
         {
             let file = TieredStorageFile::new_writable(&path).unwrap();
 
-            OwnersBlock::write_owners_block(&file, &addresses).unwrap();
+            footer
+                .owners_block_format
+                .write_owners_block(&file, &addresses)
+                .unwrap();
 
             // while the test only focuses on account metas, writing a footer
             // here is necessary to make it a valid tiered-storage file.
@@ -93,7 +120,10 @@ mod tests {
 
         for (i, address) in addresses.iter().enumerate() {
             assert_eq!(
-                OwnersBlock::get_owner_address(&mmap, &footer, OwnerOffset(i as u32)).unwrap(),
+                footer
+                    .owners_block_format
+                    .get_owner_address(&mmap, &footer, OwnerOffset(i as u32))
+                    .unwrap(),
                 address
             );
         }