Skip to content

Commit

Permalink
[TieredStorage] Refactor TieredStorage::new_readonly() code path (#195)
Browse files Browse the repository at this point in the history
#### Problem
The TieredStorage::new_readonly() function currently has the following
problems:

* It opens the file without checking the magic number before checking and loading the footer.
* It opens the file twice: first to load the footer, then open again by the reader.

#### Summary of Changes
This PR refactors TieredStorage::new_readonly() so that it first performs all
checks inside the constructor of TieredReadableFile.  The TieredReadableFile
instance is then passed to the proper reader (currently HotStorageReader)
when all checks are passed.

#### Test Plan
* Added a new test to check MagicNumberMismatch.
* Existing tiered-storage tests
  • Loading branch information
yhchiang-sol authored Mar 20, 2024
1 parent f41fb84 commit 0e932c7
Show file tree
Hide file tree
Showing 5 changed files with 105 additions and 49 deletions.
3 changes: 2 additions & 1 deletion accounts-db/src/tiered_storage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -170,7 +170,8 @@ mod tests {
use {
super::*,
crate::account_storage::meta::StoredMetaWriteVersion,
footer::{TieredStorageFooter, TieredStorageMagicNumber},
file::TieredStorageMagicNumber,
footer::TieredStorageFooter,
hot::HOT_FORMAT,
index::IndexOffset,
solana_sdk::{
Expand Down
86 changes: 75 additions & 11 deletions accounts-db/src/tiered_storage/file.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
use {
bytemuck::{AnyBitPattern, NoUninit},
super::{error::TieredStorageError, TieredStorageResult},
bytemuck::{AnyBitPattern, NoUninit, Pod, Zeroable},
std::{
fs::{File, OpenOptions},
io::{BufWriter, Read, Result as IoResult, Seek, SeekFrom, Write},
Expand All @@ -8,23 +9,37 @@ use {
},
};

/// The ending 8 bytes of a valid tiered account storage file.
pub const FILE_MAGIC_NUMBER: u64 = u64::from_le_bytes(*b"AnzaTech");

#[derive(Debug, PartialEq, Eq, Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct TieredStorageMagicNumber(pub u64);

// Ensure there are no implicit padding bytes
const _: () = assert!(std::mem::size_of::<TieredStorageMagicNumber>() == 8);

impl Default for TieredStorageMagicNumber {
fn default() -> Self {
Self(FILE_MAGIC_NUMBER)
}
}

#[derive(Debug)]
pub struct TieredReadableFile(pub File);

impl TieredReadableFile {
pub fn new(file_path: impl AsRef<Path>) -> Self {
Self(
pub fn new(file_path: impl AsRef<Path>) -> TieredStorageResult<Self> {
let file = Self(
OpenOptions::new()
.read(true)
.create(false)
.open(&file_path)
.unwrap_or_else(|err| {
panic!(
"[TieredStorageError] Unable to open {} as read-only: {err}",
file_path.as_ref().display(),
);
}),
)
.open(&file_path)?,
);

file.check_magic_number()?;

Ok(file)
}

pub fn new_writable(file_path: impl AsRef<Path>) -> IoResult<Self> {
Expand All @@ -36,6 +51,19 @@ impl TieredReadableFile {
))
}

fn check_magic_number(&self) -> TieredStorageResult<()> {
self.seek_from_end(-(std::mem::size_of::<TieredStorageMagicNumber>() as i64))?;
let mut magic_number = TieredStorageMagicNumber::zeroed();
self.read_pod(&mut magic_number)?;
if magic_number != TieredStorageMagicNumber::default() {
return Err(TieredStorageError::MagicNumberMismatch(
TieredStorageMagicNumber::default().0,
magic_number.0,
));
}
Ok(())
}

/// Reads a value of type `T` from the file.
///
/// Type T must be plain ol' data.
Expand Down Expand Up @@ -127,3 +155,39 @@ impl TieredWritableFile {
Ok(bytes.len())
}
}

#[cfg(test)]
mod tests {
use {
crate::tiered_storage::{
error::TieredStorageError,
file::{TieredReadableFile, TieredWritableFile, FILE_MAGIC_NUMBER},
},
std::path::Path,
tempfile::TempDir,
};

fn generate_test_file_with_number(path: impl AsRef<Path>, number: u64) {
let mut file = TieredWritableFile::new(path).unwrap();
file.write_pod(&number).unwrap();
}

#[test]
fn test_new() {
let temp_dir = TempDir::new().unwrap();
let path = temp_dir.path().join("test_new");
generate_test_file_with_number(&path, FILE_MAGIC_NUMBER);
assert!(TieredReadableFile::new(&path).is_ok());
}

#[test]
fn test_magic_number_mismatch() {
let temp_dir = TempDir::new().unwrap();
let path = temp_dir.path().join("test_magic_number_mismatch");
generate_test_file_with_number(&path, !FILE_MAGIC_NUMBER);
assert!(matches!(
TieredReadableFile::new(&path),
Err(TieredStorageError::MagicNumberMismatch(_, _))
));
}
}
24 changes: 4 additions & 20 deletions accounts-db/src/tiered_storage/footer.rs
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
use {
crate::tiered_storage::{
error::TieredStorageError,
file::{TieredReadableFile, TieredWritableFile},
file::{TieredReadableFile, TieredStorageMagicNumber, TieredWritableFile},
index::IndexBlockFormat,
mmap_utils::{get_pod, get_type},
owners::OwnersBlockFormat,
TieredStorageResult,
},
bytemuck::{Pod, Zeroable},
bytemuck::Zeroable,
memmap2::Mmap,
num_enum::TryFromPrimitiveError,
solana_sdk::{hash::Hash, pubkey::Pubkey},
Expand All @@ -26,22 +26,6 @@ static_assertions::const_assert_eq!(mem::size_of::<TieredStorageFooter>(), 160);
/// even when the footer's format changes.
pub const FOOTER_TAIL_SIZE: usize = 24;

/// The ending 8 bytes of a valid tiered account storage file.
pub const FOOTER_MAGIC_NUMBER: u64 = 0x502A2AB5; // SOLALABS -> SOLANA LABS

#[derive(Debug, PartialEq, Eq, Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct TieredStorageMagicNumber(pub u64);

// Ensure there are no implicit padding bytes
const _: () = assert!(std::mem::size_of::<TieredStorageMagicNumber>() == 8);

impl Default for TieredStorageMagicNumber {
fn default() -> Self {
Self(FOOTER_MAGIC_NUMBER)
}
}

#[repr(u16)]
#[derive(
Clone,
Expand Down Expand Up @@ -133,7 +117,7 @@ pub struct TieredStorageFooter {
/// The size of the footer including the magic number.
pub footer_size: u64,
// This field is persisted in the storage but not in this struct.
// The number should match FOOTER_MAGIC_NUMBER.
// The number should match FILE_MAGIC_NUMBER.
// pub magic_number: u64,
}

Expand Down Expand Up @@ -186,7 +170,7 @@ impl Default for TieredStorageFooter {

impl TieredStorageFooter {
pub fn new_from_path(path: impl AsRef<Path>) -> TieredStorageResult<Self> {
let file = TieredReadableFile::new(path);
let file = TieredReadableFile::new(path)?;
Self::new_from_footer_block(&file)
}

Expand Down
35 changes: 20 additions & 15 deletions accounts-db/src/tiered_storage/hot.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use {
accounts_hash::AccountHash,
tiered_storage::{
byte_block,
file::TieredWritableFile,
file::{TieredReadableFile, TieredWritableFile},
footer::{AccountBlockFormat, AccountMetaFormat, TieredStorageFooter},
index::{AccountIndexWriterEntry, AccountOffset, IndexBlockFormat, IndexOffset},
meta::{AccountMetaFlags, AccountMetaOptionalFields, TieredAccountMeta},
Expand All @@ -24,7 +24,7 @@ use {
account::ReadableAccount, pubkey::Pubkey, rent_collector::RENT_EXEMPT_RENT_EPOCH,
stake_history::Epoch,
},
std::{borrow::Borrow, fs::OpenOptions, option::Option, path::Path},
std::{borrow::Borrow, option::Option, path::Path},
};

pub const HOT_FORMAT: TieredStorageFormat = TieredStorageFormat {
Expand Down Expand Up @@ -346,10 +346,8 @@ pub struct HotStorageReader {
}

impl HotStorageReader {
/// Constructs a HotStorageReader from the specified path.
pub fn new_from_path(path: impl AsRef<Path>) -> TieredStorageResult<Self> {
let file = OpenOptions::new().read(true).open(path)?;
let mmap = unsafe { MmapOptions::new().map(&file)? };
pub fn new(file: TieredReadableFile) -> TieredStorageResult<Self> {
let mmap = unsafe { MmapOptions::new().map(&file.0)? };
// Here we are copying the footer, as accessing any data in a
// TieredStorage instance requires accessing its Footer.
// This can help improve cache locality and reduce the overhead
Expand Down Expand Up @@ -899,7 +897,8 @@ pub mod tests {
// Reopen the same storage, and expect the persisted footer is
// the same as what we have written.
{
let hot_storage = HotStorageReader::new_from_path(&path).unwrap();
let file = TieredReadableFile::new(&path).unwrap();
let hot_storage = HotStorageReader::new(file).unwrap();
assert_eq!(expected_footer, *hot_storage.footer());
}
}
Expand Down Expand Up @@ -945,7 +944,8 @@ pub mod tests {
footer.write_footer_block(&mut file).unwrap();
}

let hot_storage = HotStorageReader::new_from_path(&path).unwrap();
let file = TieredReadableFile::new(&path).unwrap();
let hot_storage = HotStorageReader::new(file).unwrap();

for (offset, expected_meta) in account_offsets.iter().zip(hot_account_metas.iter()) {
let meta = hot_storage.get_account_meta_from_offset(*offset).unwrap();
Expand Down Expand Up @@ -975,7 +975,8 @@ pub mod tests {
footer.write_footer_block(&mut file).unwrap();
}

let hot_storage = HotStorageReader::new_from_path(&path).unwrap();
let file = TieredReadableFile::new(&path).unwrap();
let hot_storage = HotStorageReader::new(file).unwrap();
let offset = HotAccountOffset::new(footer.index_block_offset as usize).unwrap();
// Read from index_block_offset, which offset doesn't belong to
// account blocks. Expect assert failure here
Expand Down Expand Up @@ -1026,7 +1027,8 @@ pub mod tests {
footer.write_footer_block(&mut file).unwrap();
}

let hot_storage = HotStorageReader::new_from_path(&path).unwrap();
let file = TieredReadableFile::new(&path).unwrap();
let hot_storage = HotStorageReader::new(file).unwrap();
for (i, index_writer_entry) in index_writer_entries.iter().enumerate() {
let account_offset = hot_storage
.get_account_offset(IndexOffset(i as u32))
Expand Down Expand Up @@ -1075,7 +1077,8 @@ pub mod tests {
footer.write_footer_block(&mut file).unwrap();
}

let hot_storage = HotStorageReader::new_from_path(&path).unwrap();
let file = TieredReadableFile::new(&path).unwrap();
let hot_storage = HotStorageReader::new(file).unwrap();
for (i, address) in addresses.iter().enumerate() {
assert_eq!(
hot_storage
Expand Down Expand Up @@ -1149,7 +1152,8 @@ pub mod tests {
footer.write_footer_block(&mut file).unwrap();
}

let hot_storage = HotStorageReader::new_from_path(&path).unwrap();
let file = TieredReadableFile::new(&path).unwrap();
let hot_storage = HotStorageReader::new(file).unwrap();

// First, verify whether we can find the expected owners.
let mut owner_candidates = owner_addresses.clone();
Expand Down Expand Up @@ -1281,7 +1285,8 @@ pub mod tests {
footer.write_footer_block(&mut file).unwrap();
}

let hot_storage = HotStorageReader::new_from_path(&path).unwrap();
let file = TieredReadableFile::new(&path).unwrap();
let hot_storage = HotStorageReader::new(file).unwrap();

for i in 0..NUM_ACCOUNTS {
let (stored_meta, next) = hot_storage
Expand Down Expand Up @@ -1362,10 +1367,10 @@ pub mod tests {
writer.write_accounts(&storable_accounts, 0).unwrap()
};

let hot_storage = HotStorageReader::new_from_path(&path).unwrap();
let file = TieredReadableFile::new(&path).unwrap();
let hot_storage = HotStorageReader::new(file).unwrap();

let num_accounts = account_data_sizes.len();

for i in 0..num_accounts {
let (stored_meta, next) = hot_storage
.get_account(IndexOffset(i as u32))
Expand Down
6 changes: 4 additions & 2 deletions accounts-db/src/tiered_storage/readable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ use {
account_storage::meta::StoredAccountMeta,
accounts_file::MatchAccountOwnerError,
tiered_storage::{
file::TieredReadableFile,
footer::{AccountMetaFormat, TieredStorageFooter},
hot::HotStorageReader,
index::IndexOffset,
Expand All @@ -22,9 +23,10 @@ pub enum TieredStorageReader {
impl TieredStorageReader {
/// Creates a reader for the specified tiered storage accounts file.
pub fn new_from_path(path: impl AsRef<Path>) -> TieredStorageResult<Self> {
let footer = TieredStorageFooter::new_from_path(&path)?;
let file = TieredReadableFile::new(&path)?;
let footer = TieredStorageFooter::new_from_footer_block(&file)?;
match footer.account_meta_format {
AccountMetaFormat::Hot => Ok(Self::Hot(HotStorageReader::new_from_path(path)?)),
AccountMetaFormat::Hot => Ok(Self::Hot(HotStorageReader::new(file)?)),
}
}

Expand Down

0 comments on commit 0e932c7

Please sign in to comment.