[TieredStorage] Make TieredStorage::write_accounts() thread-safe

accounts-db/src/tiered_storage.rs

@@ -29,7 +29,10 @@ use {
         borrow::Borrow,
         fs::{self, OpenOptions},
         path::{Path, PathBuf},
-        sync::OnceLock,
+        sync::{
+            atomic::{AtomicBool, Ordering},
+            OnceLock,
+        },
     },
 };
 
@@ -46,9 +49,14 @@ pub struct TieredStorageFormat {
     pub account_block_format: AccountBlockFormat,
 }
 
+/// The implementation of AccountsFile for tiered-storage.
 #[derive(Debug)]
 pub struct TieredStorage {
+    /// The internal reader instance for its accounts file.
     reader: OnceLock<TieredStorageReader>,
+    /// A status flag indicating whether its file has been already written.
+    already_written: AtomicBool,
+    /// The path to the file that stores accounts.
     path: PathBuf,
 }
 
@@ -72,6 +80,7 @@ impl TieredStorage {
     pub fn new_writable(path: impl Into<PathBuf>) -> Self {
         Self {
             reader: OnceLock::<TieredStorageReader>::new(),
+            already_written: false.into(),
             path: path.into(),
         }
     }
@@ -82,6 +91,7 @@ impl TieredStorage {
         let path = path.into();
         Ok(Self {
             reader: TieredStorageReader::new_from_path(&path).map(OnceLock::from)?,
+            already_written: true.into(),
             path,
         })
     }
@@ -94,9 +104,7 @@ impl TieredStorage {
     /// Writes the specified accounts into this TieredStorage.
     ///
     /// Note that this function can only be called once per a TieredStorage
-    /// instance.  TieredStorageError::AttemptToUpdateReadOnly will be returned
-    /// if this function is invoked more than once on the same TieredStorage
-    /// instance.
+    /// instance.  Otherwise, it will trigger panic.
     pub fn write_accounts<
         'a,
         'b,
@@ -109,10 +117,10 @@ impl TieredStorage {
         skip: usize,
         format: &TieredStorageFormat,
     ) -> TieredStorageResult<Vec<StoredAccountInfo>> {
-        if self.is_read_only() {
-            return Err(TieredStorageError::AttemptToUpdateReadOnly(
-                self.path.to_path_buf(),
-            ));
+        let was_written = self.already_written.swap(true, Ordering::AcqRel);
+
+        if was_written {
+            panic!("cannot write same tiered storage file more than once");
         }
 
         if format == &HOT_FORMAT {
@@ -122,16 +130,17 @@ impl TieredStorage {
             };
 
             // panic here if self.reader.get() is not None as self.reader can only be
-            // None since we have passed `is_read_only()` check previously, indicating
-            // self.reader is not yet set.
+            // None since a false-value `was_written` indicates the accounts file has
+            // not been written previously, implying is_read_only() was also false.
+            debug_assert!(!self.is_read_only());
             self.reader
                 .set(TieredStorageReader::new_from_path(&self.path)?)
                 .unwrap();
 
-            return result;
+            result
+        } else {
+            Err(TieredStorageError::UnknownFormat(self.path.to_path_buf()))
         }
-
-        Err(TieredStorageError::UnknownFormat(self.path.to_path_buf()))
     }
 
     /// Returns the underlying reader of the TieredStorage.  None will be
@@ -258,6 +267,7 @@ mod tests {
     }
 
     #[test]
+    #[should_panic(expected = "cannot write same tiered storage file more than once")]
     fn test_write_accounts_twice() {
         // Generate a new temp path that is guaranteed to NOT already have a file.
         let temp_dir = tempdir().unwrap();