[TieredStorage] In-memory struct for writing OwnersBlock

accounts-db/src/tiered_storage/hot.rs

@@ -470,7 +470,7 @@ pub mod tests {
             hot::{HotAccountMeta, HotStorageReader},
             index::{AccountIndexWriterEntry, IndexBlockFormat, IndexOffset},
             meta::{AccountMetaFlags, AccountMetaOptionalFields, TieredAccountMeta},
-            owners::OwnersBlockFormat,
+            owners::{OwnersBlockFormat, OwnersTable},
         },
         assert_matches::assert_matches,
         memoffset::offset_of,
@@ -823,9 +823,13 @@ pub mod tests {
         {
             let file = TieredStorageFile::new_writable(&path).unwrap();
 
+            let mut owners_table = OwnersTable::new();
+            addresses.iter().for_each(|owner_address| {
+                owners_table.check_and_add(owner_address);
+            });
             footer
                 .owners_block_format
-                .write_owners_block(&file, &addresses)
+                .write_owners_block(&file, owners_table.owners())
                 .unwrap();
 
             // while the test only focuses on account metas, writing a footer
@@ -893,9 +897,13 @@ pub mod tests {
             // the owners_block_offset set to the end of the accounts blocks.
             footer.owners_block_offset = footer.index_block_offset;
 
+            let mut owners_table = OwnersTable::new();
+            owner_addresses.iter().for_each(|owner_address| {
+                owners_table.check_and_add(owner_address);
+            });
             footer
                 .owners_block_format
-                .write_owners_block(&file, &owner_addresses)
+                .write_owners_block(&file, owners_table.owners())
                 .unwrap();
 
             // while the test only focuses on account metas, writing a footer
@@ -1029,9 +1037,13 @@ pub mod tests {
 
             // write owners block
             footer.owners_block_offset = current_offset as u64;
+            let mut owners_table = OwnersTable::new();
+            owners.iter().for_each(|owner_address| {
+                owners_table.check_and_add(owner_address);
+            });
             footer
                 .owners_block_format
-                .write_owners_block(&file, &owners)
+                .write_owners_block(&file, owners_table.owners())
                 .unwrap();
 
             footer.write_footer_block(&file).unwrap();

accounts-db/src/tiered_storage/owners.rs

@@ -5,6 +5,7 @@ use {
     },
     memmap2::Mmap,
     solana_sdk::pubkey::Pubkey,
+    std::collections::HashMap,
 };
 
 /// The offset to an owner entry in the owners block.
@@ -40,16 +41,16 @@ pub enum OwnersBlockFormat {
 
 impl OwnersBlockFormat {
     /// Persists the provided owners' addresses into the specified file.
-    pub fn write_owners_block(
+    pub(crate) fn write_owners_block(
         &self,
         file: &TieredStorageFile,
-        addresses: &[Pubkey],
+        owners: &[&Pubkey],
     ) -> TieredStorageResult<usize> {
         match self {
             Self::AddressesOnly => {
                 let mut bytes_written = 0;
-                for address in addresses {
-                    bytes_written += file.write_pod(address)?;
+                for address in owners {
+                    bytes_written += file.write_pod(*address)?;
                 }
 
                 Ok(bytes_written)
@@ -77,6 +78,44 @@ impl OwnersBlockFormat {
     }
 }
 
+/// The in-memory representation of owners block for write.
+/// It manages a set of unique addresses of account owners.
+#[derive(Debug)]
+pub(crate) struct OwnersTable<'owner> {
+    owners_vec: Vec<&'owner Pubkey>,
+    owners_map: HashMap<&'owner Pubkey, OwnerOffset>,
+}
+
+/// OwnersBlock is persisted as a consecutive bytes of pubkeys without any
+/// meta-data.  For each account meta, it has a owner_offset field to
+/// access its owner's address in the OwnersBlock.
+impl<'owner> OwnersTable<'owner> {
+    pub(crate) fn new() -> Self {
+        Self {
+            owners_vec: vec![],
+            owners_map: HashMap::new(),
+        }
+    }
+
+    /// Add the specified pubkey as the owner into the OwnersWriterTable
+    /// if the specified pubkey has not existed in the OwnersWriterTable
+    /// yet.  In any case, the function returns its OwnerOffset.
+    pub(crate) fn check_and_add(&mut self, pubkey: &'owner Pubkey) -> OwnerOffset {
+        if let Some(offset) = self.owners_map.get(pubkey) {
+            return *offset;
+        }
+        let offset: u32 = self.owners_vec.len().try_into().unwrap();
+        self.owners_vec.push(pubkey);
+        self.owners_map.insert(pubkey, OwnerOffset(offset));
+
+        OwnerOffset(offset)
+    }
+
+    pub(crate) fn owners(&self) -> &[&'owner Pubkey] {
+        &self.owners_vec
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use {
@@ -105,9 +144,13 @@ mod tests {
         {
             let file = TieredStorageFile::new_writable(&path).unwrap();
 
+            let mut owners_table = OwnersTable::new();
+            addresses.iter().for_each(|owner_address| {
+                owners_table.check_and_add(owner_address);
+            });
             footer
                 .owners_block_format
-                .write_owners_block(&file, &addresses)
+                .write_owners_block(&file, owners_table.owners())
                 .unwrap();
 
             // while the test only focuses on account metas, writing a footer
@@ -128,4 +171,31 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    fn test_owners_table() {
+        let mut owners_table = OwnersTable::new();
+        const NUM_OWNERS: usize = 99;
+
+        let addresses: Vec<_> = std::iter::repeat_with(Pubkey::new_unique)
+            .take(NUM_OWNERS)
+            .collect();
+
+        // as we insert sequentially, we expect each entry has same OwnerOffset
+        // as its index inside the Vector.
+        for (i, address) in addresses.iter().enumerate() {
+            assert_eq!(owners_table.check_and_add(address), OwnerOffset(i as u32));
+        }
+
+        let cloned_addresses = addresses.clone();
+
+        // insert again and expect the same OwnerOffset
+        for (i, address) in cloned_addresses.iter().enumerate() {
+            assert_eq!(owners_table.check_and_add(address), OwnerOffset(i as u32));
+        }
+
+        // make sure the size of the resulting owner table is the same
+        // as the input
+        assert_eq!(owners_table.owners().len(), addresses.len());
+    }
 }