Adding an API to calculate storage size from the stored data lengths

Author: roryharr

accounts-db/src/accounts_db.rs

@@ -7829,7 +7829,8 @@ impl AccountsDb {
                         let mut offsets = offsets.iter().cloned().collect::<Vec<_>>();
                         // sort so offsets are in order. This improves efficiency of loading the accounts.
                         offsets.sort_unstable();
-                        let dead_bytes = store.accounts.get_account_sizes(&offsets).iter().sum();
+                        let data_lens = store.accounts.get_account_data_lens(&offsets);
+                        let dead_bytes = store.accounts.calculate_storage_size(&data_lens);
                         store.remove_accounts(dead_bytes, reset_accounts, offsets.len());
                         if Self::is_shrinking_productive(&store)
                             && self.is_candidate_for_shrink(&store)
@@ -8502,7 +8503,7 @@ impl AccountsDb {
         rent_collector: &RentCollector,
         storage_info: &StorageSizeAndCountMap,
     ) -> SlotIndexGenerationInfo {
-        if storage.accounts.get_account_sizes(&[0]).is_empty() {
+        if storage.accounts.get_account_data_lens(&[0]).is_empty() {
             return SlotIndexGenerationInfo::default();
         }
         let secondary = !self.account_indexes.is_empty();

accounts-db/src/accounts_file.rs

@@ -325,17 +325,29 @@ impl AccountsFile {
         })
     }
 
-    /// for each offset in `sorted_offsets`, return the account size
-    pub(crate) fn get_account_sizes(&self, sorted_offsets: &[usize]) -> Vec<usize> {
+    /// for each offset in `sorted_offsets`, get the data size
+    pub(crate) fn get_account_data_lens(&self, sorted_offsets: &[usize]) -> Vec<usize> {
         match self {
-            Self::AppendVec(av) => av.get_account_sizes(sorted_offsets),
+            Self::AppendVec(av) => av.get_account_data_lens(sorted_offsets),
             Self::TieredStorage(ts) => ts
                 .reader()
-                .and_then(|reader| reader.get_account_sizes(sorted_offsets).ok())
+                .and_then(|reader| reader.get_account_data_lens(sorted_offsets).ok())
                 .unwrap_or_default(),
         }
     }
 
+    /// Calculate the amount of storage required for the passed in data lengths
+    /// Only need data lengths, as the number of accounts can be inferred from
+    /// length of the vector, and the account size is static
+    pub(crate) fn calculate_storage_size(&self, data_lens: &[usize]) -> usize {
+        match self {
+            Self::AppendVec(av) => av.calculate_storage_size(data_lens),
+            Self::TieredStorage(ts) => ts
+                .reader()
+                .map_or(0, |reader| reader.calculate_storage_size(data_lens)),
+        }
+    }
+
     /// iterate over all entries to put in index
     pub(crate) fn scan_index(&self, callback: impl FnMut(IndexInfo)) {
         match self {

accounts-db/src/append_vec.rs

@@ -165,8 +165,6 @@ impl IsZeroLamport for IndexInfo {
 struct AccountOffsets {
     /// offset to the end of the &[u8] data
     offset_to_end_of_data: usize,
-    /// # of bytes (aligned) to store this account, including variable sized data
-    stored_size_aligned: usize,
 }
 
 #[cfg_attr(feature = "frozen-abi", derive(AbiExample))]
@@ -889,14 +887,15 @@ impl AppendVec {
         &self,
         offset: usize,
     ) -> Option<(StoredMeta, solana_account::AccountSharedData)> {
-        let sizes = self.get_account_sizes(&[offset]);
+        let data_len = self.get_account_data_lens(&[offset]);
+        let sizes = self.calculate_storage_size(&data_len);
         let result = self.get_stored_account_meta_callback(offset, |r_callback| {
             let r2 = self.get_account_shared_data(offset);
             assert!(solana_account::accounts_equal(
                 &r_callback,
                 r2.as_ref().unwrap()
             ));
-            assert_eq!(sizes, vec![r_callback.stored_size()]);
+            assert_eq!(sizes, r_callback.stored_size());
             let meta = r_callback.meta().clone();
             Some((meta, r_callback.to_account_shared_data()))
         });
@@ -906,7 +905,7 @@ impl AppendVec {
                 .is_none());
             assert!(self.get_account_shared_data(offset).is_none());
             // it has different rules for checking len and returning None
-            assert!(sizes.is_empty());
+            assert!(sizes == 0);
         }
         result.flatten()
     }
@@ -926,12 +925,10 @@ impl AppendVec {
         let stored_size_unaligned = STORE_META_OVERHEAD
             .checked_add(stored_meta.data_len as usize)
             .expect("stored size cannot overflow");
-        let stored_size_aligned = u64_align!(stored_size_unaligned);
         let offset_to_end_of_data = start_offset + stored_size_unaligned;
 
         AccountOffsets {
             offset_to_end_of_data,
-            stored_size_aligned,
         }
     }
 
@@ -1098,8 +1095,13 @@ impl AppendVec {
         }
     }
 
-    /// for each offset in `sorted_offsets`, get the size of the account. No other information is needed for the account.
-    pub(crate) fn get_account_sizes(&self, sorted_offsets: &[usize]) -> Vec<usize> {
+    /// Estimate the amount of storage required for the passed in data lengths
+    pub(crate) fn calculate_storage_size(&self, data_len: &[usize]) -> usize {
+        data_len.iter().map(|len| aligned_stored_size(*len)).sum()
+    }
+
+    /// for each offset in `sorted_offsets`, get the the amount of data stored in the account.
+    pub(crate) fn get_account_data_lens(&self, sorted_offsets: &[usize]) -> Vec<usize> {
         // self.len() is an atomic load, so only do it once
         let self_len = self.len();
         let mut account_sizes = Vec::with_capacity(sorted_offsets.len());
@@ -1115,7 +1117,7 @@ impl AppendVec {
                         // data doesn't fit, so don't include
                         break;
                     }
-                    account_sizes.push(next.stored_size_aligned);
+                    account_sizes.push(stored_meta.data_len as usize);
                 }
             }
             AppendVecFileBacking::File(file) => {
@@ -1143,7 +1145,7 @@ impl AppendVec {
                         // data doesn't fit, so don't include
                         break;
                     }
-                    account_sizes.push(next.stored_size_aligned);
+                    account_sizes.push(stored_meta.data_len as usize);
                 }
             }
         }
@@ -1699,7 +1701,7 @@ pub mod tests {
         let size = 1000;
         let mut indexes = vec![];
         let now = Instant::now();
-        let mut sizes = vec![];
+        let mut sizes: Vec<usize> = vec![];
         for sample in 0..size {
             // sample + 1 is so sample = 0 won't be used.
             // sample = 0 produces default account with default pubkey
@@ -1708,7 +1710,10 @@ pub mod tests {
             let pos = av.append_account_test(&account).unwrap();
             assert_eq!(av.get_account_test(pos).unwrap(), account);
             indexes.push(pos);
-            assert_eq!(sizes, av.get_account_sizes(&indexes));
+            assert_eq!(
+                sizes.iter().sum::<usize>(),
+                av.calculate_storage_size(&av.get_account_data_lens(indexes.as_slice()))
+            );
         }
         trace!("append time: {} ms", now.elapsed().as_millis());
 
@@ -2071,8 +2076,9 @@ pub mod tests {
         let (append_vec, _) =
             AppendVec::new_from_file(&temp_file.path, total_stored_size, storage_access).unwrap();
 
-        let account_sizes = append_vec.get_account_sizes(account_offsets.as_slice());
-        assert_eq!(account_sizes, stored_sizes);
+        let account_sizes = append_vec
+            .calculate_storage_size(&append_vec.get_account_data_lens(account_offsets.as_slice()));
+        assert_eq!(account_sizes, total_stored_size);
     }
 
     /// A helper function for testing different scenario for scan_*.

accounts-db/src/tiered_storage/hot.rs

@@ -613,8 +613,15 @@ impl HotStorageReader {
         Ok(())
     }
 
-    /// for each offset in `sorted_offsets`, return the account size
-    pub(crate) fn get_account_sizes(
+    /// Calculate the amount of storage required for the passed in data lengths
+    /// Only need data lengths, as the number of accounts can be inferred from
+    /// length of the vector, and the account size is static
+    pub(crate) fn calculate_storage_size(&self, data_len: &[usize]) -> usize {
+        data_len.iter().map(|data_len| stored_size(*data_len)).sum()
+    }
+
+    /// for each offset in `sorted_offsets`, return the length of data stored in the account
+    pub(crate) fn get_account_data_lens(
         &self,
         sorted_offsets: &[usize],
     ) -> TieredStorageResult<Vec<usize>> {
@@ -625,7 +632,7 @@ impl HotStorageReader {
             let meta = self.get_account_meta_from_offset(account_offset)?;
             let account_block = self.get_account_block(account_offset, index_offset)?;
             let data_len = meta.account_data_size(account_block);
-            result.push(stored_size(data_len));
+            result.push(data_len);
         }
         Ok(result)
     }

accounts-db/src/tiered_storage/readable.rs

@@ -171,13 +171,22 @@ impl TieredStorageReader {
         }
     }
 
-    /// for each offset in `sorted_offsets`, return the account size
-    pub(crate) fn get_account_sizes(
+    /// Calculate the amount of storage required for the passed in data lengths
+    /// Only need data lengths, as the number of accounts can be inferred from
+    /// length of the vector, and the account size is static
+    pub(crate) fn calculate_storage_size(&self, data_len: &[usize]) -> usize {
+        match self {
+            Self::Hot(hot) => hot.calculate_storage_size(data_len),
+        }
+    }
+
+    /// for each offset in `sorted_offsets`, return the length of data stored in the account
+    pub(crate) fn get_account_data_lens(
         &self,
         sorted_offsets: &[usize],
     ) -> TieredStorageResult<Vec<usize>> {
         match self {
-            Self::Hot(hot) => hot.get_account_sizes(sorted_offsets),
+            Self::Hot(hot) => hot.get_account_data_lens(sorted_offsets),
         }
     }