From f64d12482229eec023ddc71dd3d576de0ddd7e45 Mon Sep 17 00:00:00 2001
From: Yueh-Hsuan Chiang <93241502+yhchiang-sol@users.noreply.github.com>
Date: Wed, 15 Nov 2023 15:54:53 -0800
Subject: [PATCH] [TieredStorage] HotStorageReader::get_owner_address (#34053)

#### Problem
The HotStorageReader does not have an API to obtain owners_address.

#### Summary of Changes
This PR adds HotStorageReader::get_owner_address().

#### Test Plan
A new unit-test is added to this PR.
---
 accounts-db/src/tiered_storage/hot.rs | 45 +++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/accounts-db/src/tiered_storage/hot.rs b/accounts-db/src/tiered_storage/hot.rs
index 4812acde14098c..349d4297792783 100644
--- a/accounts-db/src/tiered_storage/hot.rs
+++ b/accounts-db/src/tiered_storage/hot.rs
@@ -12,6 +12,7 @@ use {
             index::{AccountOffset, IndexBlockFormat, IndexOffset},
             meta::{AccountMetaFlags, AccountMetaOptionalFields, TieredAccountMeta},
             mmap_utils::get_type,
+            owners::{OwnerOffset, OwnersBlock},
             TieredStorageFormat, TieredStorageResult,
         },
     },
@@ -244,6 +245,12 @@ impl HotStorageReader {
             .index_block_format
             .get_account_address(&self.mmap, &self.footer, index)
     }
+
+    /// 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)
+    }
 }
 
 #[cfg(test)]
@@ -531,4 +538,42 @@ pub mod tests {
             assert_eq!(account_address, index_writer_entry.address);
         }
     }
+
+    #[test]
+    fn test_hot_storage_get_owner_address() {
+        // 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_get_owner_address");
+        const NUM_OWNERS: usize = 10;
+
+        let addresses: Vec<_> = std::iter::repeat_with(Pubkey::new_unique)
+            .take(NUM_OWNERS)
+            .collect();
+
+        let footer = TieredStorageFooter {
+            account_meta_format: AccountMetaFormat::Hot,
+            // Set owners_block_offset to 0 as we didn't write any account
+            // meta/data nor index block in this test
+            owners_block_offset: 0,
+            ..TieredStorageFooter::default()
+        };
+
+        {
+            let file = TieredStorageFile::new_writable(&path).unwrap();
+
+            OwnersBlock::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.
+            footer.write_footer_block(&file).unwrap();
+        }
+
+        let hot_storage = HotStorageReader::new_from_path(&path).unwrap();
+        for (i, address) in addresses.iter().enumerate() {
+            assert_eq!(
+                hot_storage.get_owner_address(OwnerOffset(i)).unwrap(),
+                address,
+            );
+        }
+    }
 }