[TieredStorage] Use BufWriter in TieredWritableFile (pyth-network#261)

#### Problem
TieredWritableFile currently uses File instead of BufWriter.
This will introduce more syscall when doing file writes.

#### Summary of Changes
This PR makes TieredWritableFile uses BufWriter to allow the
write-call to be more optimized to reduce the number of syscalls.

#### Test Plan
Existing tiered-storage test.
Will run experiments to verify its performance improvement.

#### Dependency
anza-xyz/agave#260

Author: yhchiang-sol

accounts-db/src/tiered_storage.rs

@@ -125,7 +125,7 @@ impl TieredStorage {
 
         if format == &HOT_FORMAT {
             let result = {
-                let writer = HotStorageWriter::new(&self.path)?;
+                let mut writer = HotStorageWriter::new(&self.path)?;
                 writer.write_accounts(accounts, skip)
             };
 

accounts-db/src/tiered_storage/file.rs

@@ -2,7 +2,7 @@ use {
     bytemuck::{AnyBitPattern, NoUninit},
     std::{
         fs::{File, OpenOptions},
-        io::{Read, Result as IoResult, Seek, SeekFrom, Write},
+        io::{BufWriter, Read, Result as IoResult, Seek, SeekFrom, Write},
         mem,
         path::Path,
     },
@@ -77,22 +77,22 @@ impl TieredReadableFile {
 }
 
 #[derive(Debug)]
-pub struct TieredWritableFile(pub File);
+pub struct TieredWritableFile(pub BufWriter<File>);
 
 impl TieredWritableFile {
     pub fn new(file_path: impl AsRef<Path>) -> IoResult<Self> {
-        Ok(Self(
+        Ok(Self(BufWriter::new(
             OpenOptions::new()
                 .create_new(true)
                 .write(true)
                 .open(file_path)?,
-        ))
+        )))
     }
 
     /// Writes `value` to the file.
     ///
     /// `value` must be plain ol' data.
-    pub fn write_pod<T: NoUninit>(&self, value: &T) -> IoResult<usize> {
+    pub fn write_pod<T: NoUninit>(&mut self, value: &T) -> IoResult<usize> {
         // SAFETY: Since T is NoUninit, it does not contain any uninitialized bytes.
         unsafe { self.write_type(value) }
     }
@@ -107,22 +107,22 @@ impl TieredWritableFile {
     /// Caller must ensure casting T to bytes is safe.
     /// Refer to the Safety sections in std::slice::from_raw_parts()
     /// and bytemuck's Pod and NoUninit for more information.
-    pub unsafe fn write_type<T>(&self, value: &T) -> IoResult<usize> {
+    pub unsafe fn write_type<T>(&mut self, value: &T) -> IoResult<usize> {
         let ptr = value as *const _ as *const u8;
         let bytes = unsafe { std::slice::from_raw_parts(ptr, mem::size_of::<T>()) };
         self.write_bytes(bytes)
     }
 
-    pub fn seek(&self, offset: u64) -> IoResult<u64> {
-        (&self.0).seek(SeekFrom::Start(offset))
+    pub fn seek(&mut self, offset: u64) -> IoResult<u64> {
+        self.0.seek(SeekFrom::Start(offset))
     }
 
-    pub fn seek_from_end(&self, offset: i64) -> IoResult<u64> {
-        (&self.0).seek(SeekFrom::End(offset))
+    pub fn seek_from_end(&mut self, offset: i64) -> IoResult<u64> {
+        self.0.seek(SeekFrom::End(offset))
     }
 
-    pub fn write_bytes(&self, bytes: &[u8]) -> IoResult<usize> {
-        (&self.0).write_all(bytes)?;
+    pub fn write_bytes(&mut self, bytes: &[u8]) -> IoResult<usize> {
+        self.0.write_all(bytes)?;
 
         Ok(bytes.len())
     }

accounts-db/src/tiered_storage/footer.rs

@@ -190,7 +190,7 @@ impl TieredStorageFooter {
         Self::new_from_footer_block(&file)
     }
 
-    pub fn write_footer_block(&self, file: &TieredWritableFile) -> TieredStorageResult<()> {
+    pub fn write_footer_block(&self, file: &mut TieredWritableFile) -> TieredStorageResult<()> {
         // SAFETY: The footer does not contain any uninitialized bytes.
         unsafe { file.write_type(self)? };
         file.write_pod(&TieredStorageMagicNumber::default())?;
@@ -356,8 +356,8 @@ mod tests {
 
         // Persist the expected footer.
         {
-            let file = TieredWritableFile::new(&path.path).unwrap();
-            expected_footer.write_footer_block(&file).unwrap();
+            let mut file = TieredWritableFile::new(&path.path).unwrap();
+            expected_footer.write_footer_block(&mut file).unwrap();
         }
 
         // Reopen the same storage, and expect the persisted footer is

accounts-db/src/tiered_storage/hot.rs

@@ -542,7 +542,7 @@ impl HotStorageReader {
 }
 
 fn write_optional_fields(
-    file: &TieredWritableFile,
+    file: &mut TieredWritableFile,
     opt_fields: &AccountMetaOptionalFields,
 ) -> TieredStorageResult<usize> {
     let mut size = 0;
@@ -572,7 +572,7 @@ impl HotStorageWriter {
     /// Persists an account with the specified information and returns
     /// the stored size of the account.
     fn write_account(
-        &self,
+        &mut self,
         lamports: u64,
         owner_offset: OwnerOffset,
         account_data: &[u8],
@@ -599,7 +599,7 @@ impl HotStorageWriter {
         stored_size += self
             .storage
             .write_bytes(&PADDING_BUFFER[0..(padding_len as usize)])?;
-        stored_size += write_optional_fields(&self.storage, &optional_fields)?;
+        stored_size += write_optional_fields(&mut self.storage, &optional_fields)?;
 
         Ok(stored_size)
     }
@@ -614,7 +614,7 @@ impl HotStorageWriter {
         U: StorableAccounts<'a, T>,
         V: Borrow<AccountHash>,
     >(
-        &self,
+        &mut self,
         accounts: &StorableAccountsWithHashesAndWriteVersions<'a, 'b, T, U, V>,
         skip: usize,
     ) -> TieredStorageResult<Vec<StoredAccountInfo>> {
@@ -677,7 +677,7 @@ impl HotStorageWriter {
         footer.index_block_offset = cursor as u64;
         cursor += footer
             .index_block_format
-            .write_index_block(&self.storage, &index)?;
+            .write_index_block(&mut self.storage, &index)?;
         if cursor % HOT_BLOCK_ALIGNMENT != 0 {
             // In case it is not yet aligned, it is due to the fact that
             // the index block has an odd number of entries.  In such case,
@@ -692,9 +692,9 @@ impl HotStorageWriter {
         footer.owner_count = owners_table.len() as u32;
         footer
             .owners_block_format
-            .write_owners_block(&self.storage, &owners_table)?;
+            .write_owners_block(&mut self.storage, &owners_table)?;
 
-        footer.write_footer_block(&self.storage)?;
+        footer.write_footer_block(&mut self.storage)?;
 
         Ok(stored_infos)
     }
@@ -892,8 +892,8 @@ pub mod tests {
         };
 
         {
-            let file = TieredWritableFile::new(&path).unwrap();
-            expected_footer.write_footer_block(&file).unwrap();
+            let mut file = TieredWritableFile::new(&path).unwrap();
+            expected_footer.write_footer_block(&mut file).unwrap();
         }
 
         // Reopen the same storage, and expect the persisted footer is
@@ -928,7 +928,7 @@ pub mod tests {
             ..TieredStorageFooter::default()
         };
         {
-            let file = TieredWritableFile::new(&path).unwrap();
+            let mut file = TieredWritableFile::new(&path).unwrap();
             let mut current_offset = 0;
 
             account_offsets = hot_account_metas
@@ -942,7 +942,7 @@ pub mod tests {
             // while the test only focuses on account metas, writing a footer
             // here is necessary to make it a valid tiered-storage file.
             footer.index_block_offset = current_offset as u64;
-            footer.write_footer_block(&file).unwrap();
+            footer.write_footer_block(&mut file).unwrap();
         }
 
         let hot_storage = HotStorageReader::new_from_path(&path).unwrap();
@@ -971,8 +971,8 @@ pub mod tests {
         };
 
         {
-            let file = TieredWritableFile::new(&path).unwrap();
-            footer.write_footer_block(&file).unwrap();
+            let mut file = TieredWritableFile::new(&path).unwrap();
+            footer.write_footer_block(&mut file).unwrap();
         }
 
         let hot_storage = HotStorageReader::new_from_path(&path).unwrap();
@@ -1016,14 +1016,14 @@ pub mod tests {
             ..TieredStorageFooter::default()
         };
         {
-            let file = TieredWritableFile::new(&path).unwrap();
+            let mut file = TieredWritableFile::new(&path).unwrap();
 
             let cursor = footer
                 .index_block_format
-                .write_index_block(&file, &index_writer_entries)
+                .write_index_block(&mut file, &index_writer_entries)
                 .unwrap();
             footer.owners_block_offset = cursor as u64;
-            footer.write_footer_block(&file).unwrap();
+            footer.write_footer_block(&mut file).unwrap();
         }
 
         let hot_storage = HotStorageReader::new_from_path(&path).unwrap();
@@ -1059,20 +1059,20 @@ pub mod tests {
         };
 
         {
-            let file = TieredWritableFile::new(&path).unwrap();
+            let mut file = TieredWritableFile::new(&path).unwrap();
 
             let mut owners_table = OwnersTable::default();
             addresses.iter().for_each(|owner_address| {
                 owners_table.insert(owner_address);
             });
             footer
                 .owners_block_format
-                .write_owners_block(&file, &owners_table)
+                .write_owners_block(&mut file, &owners_table)
                 .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();
+            footer.write_footer_block(&mut file).unwrap();
         }
 
         let hot_storage = HotStorageReader::new_from_path(&path).unwrap();
@@ -1118,7 +1118,7 @@ pub mod tests {
         let account_offsets: Vec<_>;
 
         {
-            let file = TieredWritableFile::new(&path).unwrap();
+            let mut file = TieredWritableFile::new(&path).unwrap();
             let mut current_offset = 0;
 
             account_offsets = hot_account_metas
@@ -1141,12 +1141,12 @@ pub mod tests {
             });
             footer
                 .owners_block_format
-                .write_owners_block(&file, &owners_table)
+                .write_owners_block(&mut file, &owners_table)
                 .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();
+            footer.write_footer_block(&mut file).unwrap();
         }
 
         let hot_storage = HotStorageReader::new_from_path(&path).unwrap();
@@ -1237,7 +1237,7 @@ pub mod tests {
         };
 
         {
-            let file = TieredWritableFile::new(&path).unwrap();
+            let mut file = TieredWritableFile::new(&path).unwrap();
             let mut current_offset = 0;
 
             // write accounts blocks
@@ -1264,7 +1264,7 @@ pub mod tests {
             footer.index_block_offset = current_offset as u64;
             current_offset += footer
                 .index_block_format
-                .write_index_block(&file, &index_writer_entries)
+                .write_index_block(&mut file, &index_writer_entries)
                 .unwrap();
 
             // write owners block
@@ -1275,10 +1275,10 @@ pub mod tests {
             });
             footer
                 .owners_block_format
-                .write_owners_block(&file, &owners_table)
+                .write_owners_block(&mut file, &owners_table)
                 .unwrap();
 
-            footer.write_footer_block(&file).unwrap();
+            footer.write_footer_block(&mut file).unwrap();
         }
 
         let hot_storage = HotStorageReader::new_from_path(&path).unwrap();
@@ -1358,7 +1358,7 @@ pub mod tests {
         let temp_dir = TempDir::new().unwrap();
         let path = temp_dir.path().join("test_write_account_and_index_blocks");
         let stored_infos = {
-            let writer = HotStorageWriter::new(&path).unwrap();
+            let mut writer = HotStorageWriter::new(&path).unwrap();
             writer.write_accounts(&storable_accounts, 0).unwrap()
         };
 

accounts-db/src/tiered_storage/index.rs

@@ -59,7 +59,7 @@ impl IndexBlockFormat {
     /// the total number of bytes written.
     pub fn write_index_block(
         &self,
-        file: &TieredWritableFile,
+        file: &mut TieredWritableFile,
         index_entries: &[AccountIndexWriterEntry<impl AccountOffset>],
     ) -> TieredStorageResult<usize> {
         match self {
@@ -181,9 +181,11 @@ mod tests {
             .collect();
 
         {
-            let file = TieredWritableFile::new(&path).unwrap();
+            let mut file = TieredWritableFile::new(&path).unwrap();
             let indexer = IndexBlockFormat::AddressesThenOffsets;
-            let cursor = indexer.write_index_block(&file, &index_entries).unwrap();
+            let cursor = indexer
+                .write_index_block(&mut file, &index_entries)
+                .unwrap();
             footer.owners_block_offset = cursor as u64;
         }
 
@@ -223,8 +225,8 @@ mod tests {
         {
             // we only write a footer here as the test should hit an assert
             // failure before it actually reads the file.
-            let file = TieredWritableFile::new(&path).unwrap();
-            footer.write_footer_block(&file).unwrap();
+            let mut file = TieredWritableFile::new(&path).unwrap();
+            footer.write_footer_block(&mut file).unwrap();
         }
 
         let file = OpenOptions::new()
@@ -259,8 +261,8 @@ mod tests {
         {
             // we only write a footer here as the test should hit an assert
             // failure before it actually reads the file.
-            let file = TieredWritableFile::new(&path).unwrap();
-            footer.write_footer_block(&file).unwrap();
+            let mut file = TieredWritableFile::new(&path).unwrap();
+            footer.write_footer_block(&mut file).unwrap();
         }
 
         let file = OpenOptions::new()
@@ -294,8 +296,8 @@ mod tests {
         {
             // we only write a footer here as the test should hit an assert
             // failure before we actually read the file.
-            let file = TieredWritableFile::new(&path).unwrap();
-            footer.write_footer_block(&file).unwrap();
+            let mut file = TieredWritableFile::new(&path).unwrap();
+            footer.write_footer_block(&mut file).unwrap();
         }
 
         let file = OpenOptions::new()
@@ -334,8 +336,8 @@ mod tests {
         {
             // we only write a footer here as the test should hit an assert
             // failure before we actually read the file.
-            let file = TieredWritableFile::new(&path).unwrap();
-            footer.write_footer_block(&file).unwrap();
+            let mut file = TieredWritableFile::new(&path).unwrap();
+            footer.write_footer_block(&mut file).unwrap();
         }
 
         let file = OpenOptions::new()

accounts-db/src/tiered_storage/owners.rs

@@ -47,7 +47,7 @@ impl OwnersBlockFormat {
     /// Persists the provided owners' addresses into the specified file.
     pub fn write_owners_block(
         &self,
-        file: &TieredWritableFile,
+        file: &mut TieredWritableFile,
         owners_table: &OwnersTable,
     ) -> TieredStorageResult<usize> {
         match self {
@@ -139,20 +139,20 @@ mod tests {
         };
 
         {
-            let file = TieredWritableFile::new(&path).unwrap();
+            let mut file = TieredWritableFile::new(&path).unwrap();
 
             let mut owners_table = OwnersTable::default();
             addresses.iter().for_each(|owner_address| {
                 owners_table.insert(owner_address);
             });
             footer
                 .owners_block_format
-                .write_owners_block(&file, &owners_table)
+                .write_owners_block(&mut file, &owners_table)
                 .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();
+            footer.write_footer_block(&mut file).unwrap();
         }
 
         let file = OpenOptions::new().read(true).open(path).unwrap();