[TieredStorage] Make AccountOffset a trait, introduce HotAccountOffset

accounts-db/src/tiered_storage/error.rs

@@ -25,4 +25,13 @@ pub enum TieredStorageError {
 
     #[error("footer is unsanitary: {0}")]
     SanitizeFooter(#[from] SanitizeFooterError),
+
+    #[error("OffsetOutOfBounds: offset {0} is larger than the supported size {1}")]
+    OffsetOutOfBounds(usize, usize),
+
+    #[error("OffsetAlignmentError: offset {0} must be multiple of {1}")]
+    OffsetAlignmentError(usize, usize),
+
+    #[error("IntraBlockOffsetOutOfBounds: offset {0} is larger than the supported size {1}")]
+    IntraBlockOffsetOutOfBounds(usize, usize),
 }

accounts-db/src/tiered_storage/hot.rs

@@ -12,7 +12,7 @@ use {
             meta::{AccountMetaFlags, AccountMetaOptionalFields, TieredAccountMeta},
             mmap_utils::get_type,
             owners::{OwnerOffset, OwnersBlock},
-            TieredStorageFormat, TieredStorageResult,
+            TieredStorageError, TieredStorageFormat, TieredStorageResult,
         },
     },
     memmap2::{Mmap, MmapOptions},
@@ -35,9 +35,12 @@ const MAX_HOT_PADDING: u8 = 7;
 /// The maximum allowed value for the owner index of a hot account.
 const MAX_HOT_OWNER_OFFSET: OwnerOffset = OwnerOffset((1 << 29) - 1);
 
-/// The multiplier for converting AccountOffset to the internal hot account
+/// The multiplier for converting HotAccountOffset to the internal hot account
 /// offset.  This increases the maximum size of a hot accounts file.
-const HOT_ACCOUNT_OFFSET_MULTIPLIER: usize = 8;
+pub(crate) const HOT_ACCOUNT_OFFSET_MULTIPLIER: usize = 8;
+
+/// The maximum supported offset for hot accounts storage.
+const MAX_HOT_ACCOUNT_OFFSET: usize = u32::MAX as usize * HOT_ACCOUNT_OFFSET_MULTIPLIER;
 
 #[bitfield(bits = 32)]
 #[repr(C)]
@@ -57,6 +60,51 @@ struct HotMetaPackedFields {
     owner_offset: B29,
 }
 
+/// The offset to access a hot account.
+#[repr(C)]
+#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
+pub struct HotAccountOffset(u32);
+
+impl AccountOffset for HotAccountOffset {
+    /// Creates a new AccountOffset instance given the block offset and the
+    /// intra-block offset.
+    fn new(block_offset: usize, intra_block_offset: usize) -> TieredStorageResult<Self> {
+        if block_offset > MAX_HOT_ACCOUNT_OFFSET {
+            return Err(TieredStorageError::OffsetOutOfBounds(
+                block_offset,
+                MAX_HOT_ACCOUNT_OFFSET,
+            ));
+        }
+
+        // Hot accounts are aligned based on HOT_ACCOUNT_OFFSET_MULTIPLIER.
+        if block_offset % HOT_ACCOUNT_OFFSET_MULTIPLIER != 0 {
+            return Err(TieredStorageError::OffsetAlignmentError(
+                block_offset,
+                HOT_ACCOUNT_OFFSET_MULTIPLIER,
+            ));
+        }
+
+        // Each hot account has its own account block.  As a result, its
+        // intra-block offset is always 0.  Any intra-block offset that is
+        // greater than 0 is invalid.
+        if intra_block_offset > 0 {
+            return Err(TieredStorageError::IntraBlockOffsetOutOfBounds(
+                block_offset,
+                0,
+            ));
+        }
+
+        Ok(HotAccountOffset(
+            (block_offset / HOT_ACCOUNT_OFFSET_MULTIPLIER) as u32,
+        ))
+    }
+
+    /// The offset to the block that contains the account.
+    fn block_offset(&self) -> usize {
+        self.0 as usize * HOT_ACCOUNT_OFFSET_MULTIPLIER
+    }
+}
+
 /// The storage and in-memory representation of the metadata entry for a
 /// hot account.
 #[derive(Debug, PartialEq, Eq)]
@@ -229,19 +277,22 @@ impl HotStorageReader {
     /// Returns the account meta located at the specified offset.
     fn get_account_meta_from_offset(
         &self,
-        account_offset: AccountOffset,
+        account_offset: HotAccountOffset,
     ) -> TieredStorageResult<&HotAccountMeta> {
-        let internal_account_offset = account_offset.block as usize * HOT_ACCOUNT_OFFSET_MULTIPLIER;
+        let internal_account_offset = account_offset.block_offset();
 
         let (meta, _) = get_type::<HotAccountMeta>(&self.mmap, internal_account_offset)?;
         Ok(meta)
     }
 
     /// Returns the offset to the account given the specified index.
-    fn get_account_offset(&self, index_offset: IndexOffset) -> TieredStorageResult<AccountOffset> {
+    fn get_account_offset(
+        &self,
+        index_offset: IndexOffset,
+    ) -> TieredStorageResult<&HotAccountOffset> {
         self.footer
             .index_block_format
-            .get_account_offset(&self.mmap, &self.footer, index_offset)
+            .get_account_offset::<HotAccountOffset>(&self.mmap, &self.footer, index_offset)
     }
 
     /// Returns the address of the account associated with the specified index.
@@ -270,7 +321,7 @@ pub mod tests {
                 FOOTER_SIZE,
             },
             hot::{HotAccountMeta, HotStorageReader},
-            index::{AccountIndexWriterEntry, AccountOffset, IndexBlockFormat, IndexOffset},
+            index::{AccountIndexWriterEntry, IndexBlockFormat, IndexOffset},
             meta::{AccountMetaFlags, AccountMetaOptionalFields, TieredAccountMeta},
         },
         memoffset::offset_of,
@@ -472,11 +523,8 @@ pub mod tests {
                 .iter()
                 .map(|meta| {
                     let prev_offset = current_offset;
-                    current_offset += file.write_type(meta).unwrap() as u32;
-                    assert_eq!(prev_offset % HOT_ACCOUNT_OFFSET_MULTIPLIER as u32, 0);
-                    AccountOffset {
-                        block: prev_offset / HOT_ACCOUNT_OFFSET_MULTIPLIER as u32,
-                    }
+                    current_offset += file.write_type(meta).unwrap();
+                    HotAccountOffset::new(prev_offset, 0).unwrap()
                 })
                 .collect();
             // while the test only focuses on account metas, writing a footer
@@ -511,8 +559,8 @@ pub mod tests {
             .iter()
             .map(|address| AccountIndexWriterEntry {
                 address,
-                block_offset: rng.gen_range(0..u32::MAX),
-                intra_block_offset: rng.gen_range(0..4096),
+                block_offset: rng.gen_range(0..u32::MAX) as usize * HOT_ACCOUNT_OFFSET_MULTIPLIER,
+                intra_block_offset: 0,
             })
             .collect();
 
@@ -529,7 +577,7 @@ pub mod tests {
 
             footer
                 .index_block_format
-                .write_index_block(&file, &index_writer_entries)
+                .write_index_block::<HotAccountOffset>(&file, &index_writer_entries)
                 .unwrap();
 
             // while the test only focuses on account metas, writing a footer
@@ -542,7 +590,10 @@ pub mod tests {
             let account_offset = hot_storage
                 .get_account_offset(IndexOffset(i as u32))
                 .unwrap();
-            assert_eq!(account_offset.block, index_writer_entry.block_offset);
+            assert_eq!(
+                account_offset.block_offset(),
+                index_writer_entry.block_offset
+            );
 
             let account_address = hot_storage
                 .get_account_address(IndexOffset(i as u32))

accounts-db/src/tiered_storage/index.rs

@@ -13,17 +13,24 @@ use {
 #[derive(Debug)]
 pub struct AccountIndexWriterEntry<'a> {
     pub address: &'a Pubkey,
-    pub block_offset: u32,
-    pub intra_block_offset: u32,
+    /// The offset to the block that contains the account.
+    pub block_offset: usize,
+    /// The offset to the account inside the account block.
+    pub intra_block_offset: usize,
 }
 
 /// The offset to an account stored inside its accounts block.
 /// This struct is used to access the meta and data of an account by looking through
 /// its accounts block.
-#[derive(Clone, Copy, Debug, Eq, PartialEq)]
-pub struct AccountOffset {
-    /// The offset to the accounts block that contains the account meta/data.
-    pub block: u32,
+pub trait AccountOffset {
+    /// Creates a new AccountOffset instance given the block offset and the
+    /// intra-block offset.
+    fn new(block_offset: usize, intra_block_offset: usize) -> TieredStorageResult<Self>
+    where
+        Self: Sized;
+
+    /// The offset to the block that contains the account.
+    fn block_offset(&self) -> usize;
 }
 
 /// The offset to an account/address entry in the accounts index block.
@@ -56,7 +63,7 @@ pub enum IndexBlockFormat {
 impl IndexBlockFormat {
     /// Persists the specified index_entries to the specified file and returns
     /// the total number of bytes written.
-    pub fn write_index_block(
+    pub fn write_index_block<T: AccountOffset>(
         &self,
         file: &TieredStorageFile,
         index_entries: &[AccountIndexWriterEntry],
@@ -68,7 +75,10 @@ impl IndexBlockFormat {
                     bytes_written += file.write_type(index_entry.address)?;
                 }
                 for index_entry in index_entries {
-                    bytes_written += file.write_type(&index_entry.block_offset)?;
+                    bytes_written += file.write_type(&T::new(
+                        index_entry.block_offset,
+                        index_entry.intra_block_offset,
+                    )?)?;
                 }
                 Ok(bytes_written)
             }
@@ -93,22 +103,20 @@ impl IndexBlockFormat {
     }
 
     /// Returns the offset to the account given the specified index.
-    pub fn get_account_offset(
+    pub fn get_account_offset<'a, T: AccountOffset>(
         &self,
-        mmap: &Mmap,
+        mmap: &'a Mmap,
         footer: &TieredStorageFooter,
         index_offset: IndexOffset,
-    ) -> TieredStorageResult<AccountOffset> {
+    ) -> TieredStorageResult<&'a T> {
         match self {
             Self::AddressAndBlockOffsetOnly => {
-                let account_offset = footer.index_block_offset as usize
+                let offset = footer.index_block_offset as usize
                     + std::mem::size_of::<Pubkey>() * footer.account_entry_count as usize
                     + std::mem::size_of::<u32>() * index_offset.0 as usize;
-                let (block_offset, _) = get_type(mmap, account_offset)?;
+                let (account_offset, _) = get_type::<T>(mmap, offset)?;
 
-                Ok(AccountOffset {
-                    block: *block_offset,
-                })
+                Ok(account_offset)
             }
         }
     }
@@ -126,8 +134,15 @@ impl IndexBlockFormat {
 #[cfg(test)]
 mod tests {
     use {
-        super::*, crate::tiered_storage::file::TieredStorageFile, memmap2::MmapOptions, rand::Rng,
-        std::fs::OpenOptions, tempfile::TempDir,
+        super::*,
+        crate::tiered_storage::{
+            file::TieredStorageFile,
+            hot::{HotAccountOffset, HOT_ACCOUNT_OFFSET_MULTIPLIER},
+        },
+        memmap2::MmapOptions,
+        rand::Rng,
+        std::fs::OpenOptions,
+        tempfile::TempDir,
     };
 
     #[test]
@@ -147,15 +162,17 @@ mod tests {
             .iter()
             .map(|address| AccountIndexWriterEntry {
                 address,
-                block_offset: rng.gen_range(128..2048),
+                block_offset: rng.gen_range(0..u32::MAX) as usize * HOT_ACCOUNT_OFFSET_MULTIPLIER,
                 intra_block_offset: 0,
             })
             .collect();
 
         {
             let file = TieredStorageFile::new_writable(&path).unwrap();
             let indexer = IndexBlockFormat::AddressAndBlockOffsetOnly;
-            indexer.write_index_block(&file, &index_entries).unwrap();
+            indexer
+                .write_index_block::<HotAccountOffset>(&file, &index_entries)
+                .unwrap();
         }
 
         let indexer = IndexBlockFormat::AddressAndBlockOffsetOnly;
@@ -167,9 +184,9 @@ mod tests {
         let mmap = unsafe { MmapOptions::new().map(&file).unwrap() };
         for (i, index_entry) in index_entries.iter().enumerate() {
             let account_offset = indexer
-                .get_account_offset(&mmap, &footer, IndexOffset(i as u32))
+                .get_account_offset::<HotAccountOffset>(&mmap, &footer, IndexOffset(i as u32))
                 .unwrap();
-            assert_eq!(index_entry.block_offset, account_offset.block);
+            assert_eq!(index_entry.block_offset, account_offset.block_offset());
             let address = indexer
                 .get_account_address(&mmap, &footer, IndexOffset(i as u32))
                 .unwrap();