[TieredStorage] Make AccountOffset a trait, introduce HotAccountOffset (#34335)

#### Problem
Hot and cold accounts storage have different implementations of
their offsets.  As a result, a single struct AccountOffset isn't suitable
to describe the offsets used by hot and cold accounts storage.

#### Summary of Changes
This PR makes AccountOffset a trait.  On top of that, introduces
HotAccountOffset that implements AccountOffset.

#### Test Plan
Updated existing unit-tests.

accounts-db/src/tiered_storage/error.rs

@@ -25,4 +25,10 @@ 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),
 }

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,13 @@ 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
-/// offset.  This increases the maximum size of a hot accounts file.
-const HOT_ACCOUNT_OFFSET_MULTIPLIER: usize = 8;
+/// The alignment for HotAccountOffset.  It is also a multiplier for converting
+/// HotAccountOffset to the internal hot account offset that increases the maximum
+/// size of a hot accounts file.
+pub(crate) const HOT_ACCOUNT_OFFSET_ALIGNMENT: usize = 8;
+
+/// The maximum supported offset for hot accounts storage.
+const MAX_HOT_ACCOUNT_OFFSET: usize = u32::MAX as usize * HOT_ACCOUNT_OFFSET_ALIGNMENT;
 
 #[bitfield(bits = 32)]
 #[repr(C)]
@@ -57,6 +61,42 @@ 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 {}
+
+impl HotAccountOffset {
+    /// Creates a new AccountOffset instance
+    pub fn new(offset: usize) -> TieredStorageResult<Self> {
+        if offset > MAX_HOT_ACCOUNT_OFFSET {
+            return Err(TieredStorageError::OffsetOutOfBounds(
+                offset,
+                MAX_HOT_ACCOUNT_OFFSET,
+            ));
+        }
+
+        // Hot accounts are aligned based on HOT_ACCOUNT_OFFSET_ALIGNMENT.
+        if offset % HOT_ACCOUNT_OFFSET_ALIGNMENT != 0 {
+            return Err(TieredStorageError::OffsetAlignmentError(
+                offset,
+                HOT_ACCOUNT_OFFSET_ALIGNMENT,
+            ));
+        }
+
+        Ok(HotAccountOffset(
+            (offset / HOT_ACCOUNT_OFFSET_ALIGNMENT) as u32,
+        ))
+    }
+
+    /// Returns the offset to the account.
+    fn offset(&self) -> usize {
+        self.0 as usize * HOT_ACCOUNT_OFFSET_ALIGNMENT
+    }
+}
+
 /// The storage and in-memory representation of the metadata entry for a
 /// hot account.
 #[derive(Debug, PartialEq, Eq)]
@@ -229,19 +269,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.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 +313,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 +515,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).unwrap()
                 })
                 .collect();
             // while the test only focuses on account metas, writing a footer
@@ -511,8 +551,10 @@ pub mod tests {
             .iter()
             .map(|address| AccountIndexWriterEntry {
                 address,
-                block_offset: rng.gen_range(0..u32::MAX),
-                intra_block_offset: rng.gen_range(0..4096),
+                offset: HotAccountOffset::new(
+                    rng.gen_range(0..u32::MAX) as usize * HOT_ACCOUNT_OFFSET_ALIGNMENT,
+                )
+                .unwrap(),
             })
             .collect();
 
@@ -542,7 +584,7 @@ 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, index_writer_entry.offset);
 
             let account_address = hot_storage
                 .get_account_address(IndexOffset(i as u32))

accounts-db/src/tiered_storage/index.rs

@@ -8,23 +8,16 @@ use {
 };
 
 /// The in-memory struct for the writing index block.
-/// The actual storage format of a tiered account index entry might be different
-/// from this.
 #[derive(Debug)]
-pub struct AccountIndexWriterEntry<'a> {
+pub struct AccountIndexWriterEntry<'a, Offset: AccountOffset> {
+    /// The account address.
     pub address: &'a Pubkey,
-    pub block_offset: u32,
-    pub intra_block_offset: u32,
+    /// The offset to the account.
+    pub offset: Offset,
 }
 
-/// 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,
-}
+/// The offset to an account.
+pub trait AccountOffset {}
 
 /// The offset to an account/address entry in the accounts index block.
 /// This can be used to obtain the AccountOffset and address by looking through
@@ -59,7 +52,7 @@ impl IndexBlockFormat {
     pub fn write_index_block(
         &self,
         file: &TieredStorageFile,
-        index_entries: &[AccountIndexWriterEntry],
+        index_entries: &[AccountIndexWriterEntry<impl AccountOffset>],
     ) -> TieredStorageResult<usize> {
         match self {
             Self::AddressAndBlockOffsetOnly => {
@@ -68,7 +61,7 @@ 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(&index_entry.offset)?;
                 }
                 Ok(bytes_written)
             }
@@ -93,31 +86,29 @@ impl IndexBlockFormat {
     }
 
     /// Returns the offset to the account given the specified index.
-    pub fn get_account_offset(
+    pub fn get_account_offset<Offset: AccountOffset + Copy>(
         &self,
         mmap: &Mmap,
         footer: &TieredStorageFooter,
         index_offset: IndexOffset,
-    ) -> TieredStorageResult<AccountOffset> {
+    ) -> TieredStorageResult<Offset> {
         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)?;
+                    + std::mem::size_of::<Offset>() * index_offset.0 as usize;
+                let (account_offset, _) = get_type::<Offset>(mmap, offset)?;
 
-                Ok(AccountOffset {
-                    block: *block_offset,
-                })
+                Ok(*account_offset)
             }
         }
     }
 
     /// Returns the size of one index entry.
-    pub fn entry_size(&self) -> usize {
+    pub fn entry_size<Offset: AccountOffset>(&self) -> usize {
         match self {
             Self::AddressAndBlockOffsetOnly => {
-                std::mem::size_of::<Pubkey>() + std::mem::size_of::<u32>()
+                std::mem::size_of::<Pubkey>() + std::mem::size_of::<Offset>()
             }
         }
     }
@@ -126,8 +117,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_ALIGNMENT},
+        },
+        memmap2::MmapOptions,
+        rand::Rng,
+        std::fs::OpenOptions,
+        tempfile::TempDir,
     };
 
     #[test]
@@ -147,8 +145,10 @@ mod tests {
             .iter()
             .map(|address| AccountIndexWriterEntry {
                 address,
-                block_offset: rng.gen_range(128..2048),
-                intra_block_offset: 0,
+                offset: HotAccountOffset::new(
+                    rng.gen_range(0..u32::MAX) as usize * HOT_ACCOUNT_OFFSET_ALIGNMENT,
+                )
+                .unwrap(),
             })
             .collect();
 
@@ -167,9 +167,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.offset, account_offset);
             let address = indexer
                 .get_account_address(&mmap, &footer, IndexOffset(i as u32))
                 .unwrap();