diff --git a/accounts-db/src/tiered_storage.rs b/accounts-db/src/tiered_storage.rs index f0a23150e2fa70..6a5bbc9a2cc862 100644 --- a/accounts-db/src/tiered_storage.rs +++ b/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, + /// 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) -> Self { Self { reader: OnceLock::::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> { - 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();