From 55172c72ffee7fa5696e72a657837bfe0f86c987 Mon Sep 17 00:00:00 2001
From: Yueh-Hsuan Chiang <yueh-hsuan@solana.com>
Date: Fri, 19 Jan 2024 15:31:48 -0800
Subject: [PATCH] [TieredStorage] HotStorageReader::write_accounts() -- owners
 block

---
 accounts-db/src/tiered_storage/hot.rs    | 40 +++++++++++++++++-------
 accounts-db/src/tiered_storage/owners.rs | 10 ++++++
 2 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/accounts-db/src/tiered_storage/hot.rs b/accounts-db/src/tiered_storage/hot.rs
index 092b067ec77b6e..26b7a5f121b1a6 100644
--- a/accounts-db/src/tiered_storage/hot.rs
+++ b/accounts-db/src/tiered_storage/hot.rs
@@ -12,7 +12,7 @@ use {
             index::{AccountIndexWriterEntry, AccountOffset, IndexBlockFormat, IndexOffset},
             meta::{AccountMetaFlags, AccountMetaOptionalFields, TieredAccountMeta},
             mmap_utils::{get_pod, get_slice},
-            owners::{OwnerOffset, OwnersBlockFormat},
+            owners::{OwnerOffset, OwnersBlockFormat, OwnersTable},
             readable::TieredReadableAccount,
             StorableAccounts, StorableAccountsWithHashesAndWriteVersions, TieredStorageError,
             TieredStorageFormat, TieredStorageResult,
@@ -496,6 +496,7 @@ impl HotStorageWriter {
         &self,
         lamports: u64,
         rent_epoch: Epoch,
+        owner_offset: OwnerOffset,
         account_data: &[u8],
         executable: bool,
         account_hash: &AccountHash,
@@ -511,6 +512,7 @@ impl HotStorageWriter {
         let padding_len = padding_bytes(account_data.len());
         let meta = HotAccountMeta::new()
             .with_lamports(lamports)
+            .with_owner_offset(owner_offset)
             .with_account_data_size(account_data.len() as u64)
             .with_account_data_padding(padding_len)
             .with_flags(&flags);
@@ -527,8 +529,8 @@ impl HotStorageWriter {
         Ok(stored_size)
     }
 
-    /// A work-in-progress function that will eventually implements
-    /// AccountsFile::appends_account()
+    /// Persists the specified accounts from index [`skip`] into the underlying
+    /// hot accounts file associated with the HotStorageWriter.
     pub fn write_accounts<
         'a,
         'b,
@@ -542,7 +544,9 @@ impl HotStorageWriter {
     ) -> TieredStorageResult<()> {
         let mut footer = new_hot_footer();
         let mut index = vec![];
+        let mut owners_table = OwnersTable::default();
         let mut cursor = 0;
+        let zero_lamport_account_owner: Pubkey = Pubkey::default();
 
         // writing accounts blocks
         let len = accounts.accounts.len();
@@ -555,18 +559,27 @@ impl HotStorageWriter {
 
             // Obtain necessary fields from the account, or default fields
             // for a zero-lamport account in the None case.
-            let (lamports, rent_epoch, data, executable) = account
+            let (lamports, rent_epoch, data, executable, owner) = account
                 .map(|acc| {
                     (
                         acc.lamports(),
                         acc.rent_epoch(),
                         acc.data(),
                         acc.executable(),
+                        acc.owner(),
                     )
                 })
-                .unwrap_or((0, Epoch::MAX, &[], false));
-
-            cursor += self.write_account(lamports, rent_epoch, data, executable, account_hash)?;
+                .unwrap_or((0, Epoch::MAX, &[], false, &zero_lamport_account_owner));
+
+            let owner_offset = owners_table.insert(owner);
+            cursor += self.write_account(
+                lamports,
+                rent_epoch,
+                owner_offset,
+                data,
+                executable,
+                account_hash,
+            )?;
             index.push(index_entry);
         }
         footer.account_entry_count = (len - skip) as u32;
@@ -586,11 +599,13 @@ impl HotStorageWriter {
             cursor += self.storage.write_pod(&0u32)?;
         }
 
-        // TODO: owner block will be implemented in the follow-up PRs
-        // expect the offset of each block aligned.
+        // writing owners block
         assert!(cursor % HOT_BLOCK_ALIGNMENT == 0);
         footer.owners_block_offset = cursor as u64;
-        footer.owner_count = 0;
+        footer.owner_count = owners_table.len() as u32;
+        footer
+            .owners_block_format
+            .write_owners_block(&self.storage, &owners_table)?;
 
         footer.write_footer_block(&self.storage)?;
 
@@ -1238,10 +1253,12 @@ pub mod tests {
     /// and write_version.
     fn create_test_account(seed: u64) -> (StoredMeta, AccountSharedData) {
         let data_byte = seed as u8;
+        let owner_byte = u8::MAX - data_byte;
         let account = Account {
             lamports: seed + 1,
             data: std::iter::repeat(data_byte).take(seed as usize).collect(),
-            owner: Pubkey::new_unique(),
+            // this will allow some test account sharing the same owner.
+            owner: [owner_byte; 32].into(),
             executable: seed % 2 > 0,
             rent_epoch: if seed % 3 > 0 {
                 seed
@@ -1317,6 +1334,7 @@ pub mod tests {
             assert_eq!(stored_meta.data().len(), account.data().len());
             assert_eq!(stored_meta.data(), account.data());
             assert_eq!(stored_meta.executable(), account.executable());
+            assert_eq!(stored_meta.owner(), account.owner());
             assert_eq!(stored_meta.pubkey(), address);
             assert_eq!(stored_meta.hash(), hash);
 
diff --git a/accounts-db/src/tiered_storage/owners.rs b/accounts-db/src/tiered_storage/owners.rs
index 45bfafe1645430..0ac8e90d50a6ed 100644
--- a/accounts-db/src/tiered_storage/owners.rs
+++ b/accounts-db/src/tiered_storage/owners.rs
@@ -97,6 +97,16 @@ impl<'a> OwnersTable<'a> {
 
         OwnerOffset(offset as u32)
     }
+
+    /// Returns the number of unique owner addresses in the table.
+    pub fn len(&self) -> usize {
+        self.owners_set.len()
+    }
+
+    /// Returns true if the OwnersTable is empty
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
 }
 
 #[cfg(test)]