Skip to content

Commit

Permalink
[TieredStorage] Make TieredStorage::write_accounts() thread-safe (sol…
Browse files Browse the repository at this point in the history
…ana-labs#35143)

#### Problem
While accounts-db might not invoke appends_account twice
for the same AccountsFile, TieredStorage::write_accounts()
itself isn't thread-safe, and it depends on the above accounts-db
assumption.

#### Summary of Changes
This PR makes TieredStorage::write_accounts() thread-safe.
So only the first thread that successfully updates the already_written
flag can proceed and write the input accounts.  All subsequent
calls to write_accounts() will be a no-op and return AttemptToUpdateReadOnly
Error.
  • Loading branch information
yhchiang-sol authored Feb 18, 2024
1 parent e406402 commit 6934589
Showing 1 changed file with 23 additions and 13 deletions.
36 changes: 23 additions & 13 deletions accounts-db/src/tiered_storage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,10 @@ use {
borrow::Borrow,
fs::{self, OpenOptions},
path::{Path, PathBuf},
sync::OnceLock,
sync::{
atomic::{AtomicBool, Ordering},
OnceLock,
},
},
};

Expand All @@ -47,9 +50,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,
}

Expand All @@ -73,6 +81,7 @@ impl TieredStorage {
pub fn new_writable(path: impl Into<PathBuf>) -> Self {
Self {
reader: OnceLock::<TieredStorageReader>::new(),
already_written: false.into(),
path: path.into(),
}
}
Expand All @@ -83,6 +92,7 @@ impl TieredStorage {
let path = path.into();
Ok(Self {
reader: TieredStorageReader::new_from_path(&path).map(OnceLock::from)?,
already_written: true.into(),
path,
})
}
Expand All @@ -95,9 +105,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,
Expand All @@ -110,10 +118,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 {
Expand All @@ -123,16 +131,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
Expand Down Expand Up @@ -255,6 +264,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();
Expand Down

0 comments on commit 6934589

Please sign in to comment.