From e6286e5528cf3e1b8963d44044ec656fc8b5331d Mon Sep 17 00:00:00 2001
From: Yueh-Hsuan Chiang <yueh-hsuan@solana.com>
Date: Mon, 23 Oct 2023 23:52:43 -0700
Subject: [PATCH 1/2] [TieredStorage] HotStorageReader::get_account_offset

---
 accounts-db/src/tiered_storage/hot.rs | 65 ++++++++++++++++++++++++++-
 1 file changed, 64 insertions(+), 1 deletion(-)

diff --git a/accounts-db/src/tiered_storage/hot.rs b/accounts-db/src/tiered_storage/hot.rs
index 78271700686dd2..9201ba078eb437 100644
--- a/accounts-db/src/tiered_storage/hot.rs
+++ b/accounts-db/src/tiered_storage/hot.rs
@@ -227,6 +227,17 @@ impl HotStorageReader {
         let (meta, _) = get_type::<HotAccountMeta>(&self.mmap, offset)?;
         Ok(meta)
     }
+
+    /// Returns the offset of the account associated with the specified index.
+    fn get_account_offset(&self, index: usize) -> usize {
+        // As Hot storage currently don't support compression, each account has
+        // its own account block.  As a result, the account block offset is the
+        // account offset.
+        self.footer
+            .account_index_format
+            .get_account_block_offset(&self.mmap, &self.footer, index)
+            .unwrap() as usize
+    }
 }
 
 #[cfg(test)]
@@ -241,7 +252,7 @@ pub mod tests {
                 FOOTER_SIZE,
             },
             hot::{HotAccountMeta, HotStorageReader},
-            index::AccountIndexFormat,
+            index::{AccountIndexFormat, AccountIndexWriterEntry},
             meta::{AccountMetaFlags, AccountMetaOptionalFields, TieredAccountMeta},
         },
         memoffset::offset_of,
@@ -460,4 +471,56 @@ pub mod tests {
         }
         assert_eq!(&footer, hot_storage.footer());
     }
+
+    #[test]
+    fn test_hot_storage_index_block() {
+        const NUM_ACCOUNTS: u32 = 10;
+        let mut rng = rand::thread_rng();
+
+        // Generate the index entry of each account
+        let addresses: Vec<_> = std::iter::repeat_with(Pubkey::new_unique)
+            .take(NUM_ACCOUNTS as usize)
+            .collect();
+        let index_entries: Vec<_> = addresses
+            .iter()
+            .map(|address| AccountIndexWriterEntry {
+                address,
+                block_offset: rng.gen_range(128..2048),
+                intra_block_offset: 0,
+            })
+            .collect();
+
+        // 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_footer");
+
+        let footer = TieredStorageFooter {
+            account_meta_format: AccountMetaFormat::Hot,
+            account_index_format: AccountIndexFormat::AddressAndOffset,
+            // Set account_index_offset to 0 as we don't write any account meta &
+            // data in this test, while this field is usually non-zero in prod.
+            account_index_offset: 0,
+            account_entry_count: NUM_ACCOUNTS,
+            ..TieredStorageFooter::default()
+        };
+
+        {
+            let file = TieredStorageFile::new_writable(&path).unwrap();
+            footer
+                .account_index_format
+                .write_index_block(&file, &index_entries)
+                .unwrap();
+
+            footer.write_footer_block(&file).unwrap();
+        }
+
+        let hot_storage = HotStorageReader::new_from_path(&path).unwrap();
+        for (i, index_entry) in index_entries.iter().enumerate() {
+            assert_eq!(
+                index_entry.block_offset,
+                hot_storage.get_account_offset(i) as u64,
+            );
+        }
+        assert_eq!(&footer, hot_storage.footer());
+    }
 }

From baa8616cc32f08e3b2dd0a3ffb0fcde44228d655 Mon Sep 17 00:00:00 2001
From: Yueh-Hsuan Chiang <yueh-hsuan@solana.com>
Date: Wed, 25 Oct 2023 12:38:58 -0700
Subject: [PATCH 2/2] Improve get_account_offset comment and make it return
 TieredStorageResult.

---
 accounts-db/src/tiered_storage/hot.rs | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/accounts-db/src/tiered_storage/hot.rs b/accounts-db/src/tiered_storage/hot.rs
index 9201ba078eb437..35a56582cc682a 100644
--- a/accounts-db/src/tiered_storage/hot.rs
+++ b/accounts-db/src/tiered_storage/hot.rs
@@ -228,15 +228,15 @@ impl HotStorageReader {
         Ok(meta)
     }
 
-    /// Returns the offset of the account associated with the specified index.
-    fn get_account_offset(&self, index: usize) -> usize {
+    /// Returns the offset to the account stored in the underlying TieredStorage
+    /// file associated with the specified index.
+    fn get_account_offset(&self, index: usize) -> TieredStorageResult<u64> {
         // As Hot storage currently don't support compression, each account has
         // its own account block.  As a result, the account block offset is the
         // account offset.
         self.footer
             .account_index_format
             .get_account_block_offset(&self.mmap, &self.footer, index)
-            .unwrap() as usize
     }
 }
 
@@ -518,7 +518,7 @@ pub mod tests {
         for (i, index_entry) in index_entries.iter().enumerate() {
             assert_eq!(
                 index_entry.block_offset,
-                hot_storage.get_account_offset(i) as u64,
+                hot_storage.get_account_offset(i).unwrap(),
             );
         }
         assert_eq!(&footer, hot_storage.footer());