diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index d8a62b7..9a456f7 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -2,9 +2,9 @@ name: Rust on: push: - branches: [ "main" ] + branches: [ "master" ] pull_request: - branches: [ "main" ] + branches: [ "master" ] permissions: contents: read diff --git a/Cargo.toml b/Cargo.toml index f6f775b..ca561ee 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,7 +3,7 @@ members = ["codegen", "tools/create-data-file", "tools/dump-data-file"] [package] name = "data_bucket" -version = "0.5.2" +version = "0.5.3" edition = "2021" authors = ["Handy-caT"] license = "MIT" diff --git a/src/lib.rs b/src/lib.rs index 199b56d..46812d9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -17,11 +17,14 @@ pub use page::{ get_index_page_size_from_data_length, map_data_pages_to_general, parse_data_page, parse_data_pages_batch, parse_general_header_by_index, parse_page, parse_pages_batch, persist_page, persist_pages_batch, seek_by_link, seek_to_page_start, update_at, DataPage, - GeneralHeader, GeneralPage, IndexPage, IndexPageUtility, IndexValue, Interval, PageType, - SpaceInfoPage, TableOfContentsPage, UnsizedIndexPage, UnsizedIndexPageUtility, DATA_VERSION, + GeneralHeader, GeneralPage, IndexPage, IndexPageUtility, IndexValue, Interval, + PageOverflowError, PageType, SpaceInfoPage, TableOfContentsOverflowError, TableOfContentsPage, + UnsizedIndexPage, UnsizedIndexPageUtility, DATA_VERSION, EMPTY_TABLE_OF_CONTENTS_PAGE_SIZE, GENERAL_HEADER_SIZE, INNER_PAGE_SIZE, PAGE_SIZE, }; pub use persistence::{PersistableIndex, PersistableTable}; pub use space::Id as SpaceId; pub use util::access_archived; -pub use util::{align, align8, align_vec, Persistable, SizeMeasurable, VariableSizeMeasurable}; +pub use util::{ + align, align8, align_to, align_vec, Persistable, SizeMeasurable, VariableSizeMeasurable, +}; diff --git a/src/page/data.rs b/src/page/data.rs index daad8b9..ae7ad1c 100644 --- a/src/page/data.rs +++ b/src/page/data.rs @@ -18,7 +18,11 @@ impl DataPage { )); } - if (link.offset + link.length) as usize > DATA_LENGTH { + // Sum in usize: `offset + length` in u32 can wrap past 4 GiB and + // slip under the bound with a range that is actually out of page. + let start = link.offset as usize; + let end = link.offset as usize + link.length as usize; + if end > DATA_LENGTH { return Err(eyre!( "Link range (offset: {}, length: {}) exceeds data bounds ({})", link.offset, @@ -27,16 +31,17 @@ impl DataPage { )); } - let start = link.offset as usize; - let end = (link.offset + link.length) as usize; self.data[start..end].copy_from_slice(new_data); - self.length = self.length.max(link.offset + link.length); + self.length = self.length.max(end as u32); Ok(()) } pub fn get_at(&self, link: Link) -> Result<&[u8]> { - if (link.offset + link.length) as usize > DATA_LENGTH { + // Sum in usize, see `update_at`. + let start = link.offset as usize; + let end = link.offset as usize + link.length as usize; + if end > DATA_LENGTH { return Err(eyre!( "Link range (offset: {}, length: {}) exceeds data bounds ({})", link.offset, @@ -45,8 +50,6 @@ impl DataPage { )); } - let start = link.offset as usize; - let end = (link.offset + link.length) as usize; Ok(&self.data[start..end]) } } @@ -126,6 +129,42 @@ mod tests { .contains("Link range (offset: 98, length: 3) exceeds data bounds (100)")); } + #[test] + fn test_update_at_offset_plus_length_wrapping_u32() { + let mut data = DataPage { + length: 0, + data: [0; 100], + }; + + // In u32, offset + length wraps to 5 and used to pass the bounds + // check, panicking on the slice instead of returning an error. + let link = Link { + page_id: 1.into(), + offset: u32::MAX - 2, + length: 8, + }; + + let err = data.update_at(link, &[1, 2, 3, 4, 5, 6, 7, 8]).unwrap_err(); + assert!(err.to_string().contains("exceeds data bounds")); + } + + #[test] + fn test_get_at_offset_plus_length_wrapping_u32() { + let data = DataPage { + length: 0, + data: [0; 100], + }; + + let link = Link { + page_id: 1.into(), + offset: u32::MAX - 2, + length: 8, + }; + + let err = data.get_at(link).unwrap_err(); + assert!(err.to_string().contains("exceeds data bounds")); + } + #[test] fn test_get_at_out_of_bounds() { let data = DataPage { diff --git a/src/page/index/mod.rs b/src/page/index/mod.rs index 60a9b7b..81e46d7 100644 --- a/src/page/index/mod.rs +++ b/src/page/index/mod.rs @@ -7,9 +7,10 @@ use rkyv::{Archive, Deserialize, Serialize}; use tokio::fs::File; use tokio::io::{AsyncSeekExt, AsyncWriteExt}; +use crate::page::PageOverflowError; use crate::{ - align, align8, seek_to_page_start, Link, Persistable, SizeMeasurable, VariableSizeMeasurable, - GENERAL_HEADER_SIZE, + align, align_to, seek_to_page_start, Link, Persistable, SizeMeasurable, VariableSizeMeasurable, + GENERAL_HEADER_SIZE, INNER_PAGE_SIZE, }; mod page; @@ -22,7 +23,9 @@ use crate::page::PageId; pub use page::{get_index_page_size_from_data_length, IndexPage}; pub use page_for_unsized::{UnsizedIndexPage, UnsizedIndexPageUtility}; -pub use table_of_contents_page::TableOfContentsPage; +pub use table_of_contents_page::{ + TableOfContentsOverflowError, TableOfContentsPage, EMPTY_TABLE_OF_CONTENTS_PAGE_SIZE, +}; pub trait IndexPageUtility { type Utility: Persistable + Send + Sync; @@ -38,10 +41,21 @@ pub trait IndexPageUtility { utility: Self::Utility, ) -> impl std::future::Future> + Send { async move { + let bytes = utility.as_bytes(); + let utility_length = bytes.as_ref().len(); + // An oversized utility must fail here, in its own persist, + // instead of writing past the page slot into the neighbor page. + if utility_length > INNER_PAGE_SIZE { + return Err(eyre::Report::new(PageOverflowError { + page_id, + data_length: utility_length, + capacity: INNER_PAGE_SIZE, + })); + } seek_to_page_start(file, page_id.0).await?; file.seek(SeekFrom::Current(GENERAL_HEADER_SIZE as i64)) .await?; - file.write_all(utility.as_bytes().as_ref()).await?; + file.write_all(bytes.as_ref()).await?; Ok(()) } } @@ -62,12 +76,15 @@ where T: SizeMeasurable, { fn aligned_size(&self) -> usize { - if let Some(align) = T::align() { - if align % 8 == 0 { - return align8(self.key.aligned_size() + self.link.aligned_size()); + let len = self.key.aligned_size() + self.link.aligned_size(); + if let Some(key_align) = T::align() { + if key_align % 8 == 0 { + // rkyv pads the archived value out to the key's real + // alignment (16 for u128-likes), so round to it, not to 8. + return align_to(len, key_align); } } - align(self.key.aligned_size() + self.link.aligned_size()) + align(len) } } diff --git a/src/page/index/page.rs b/src/page/index/page.rs index d2f0b24..cccb1cf 100644 --- a/src/page/index/page.rs +++ b/src/page/index/page.rs @@ -18,9 +18,10 @@ use tokio::fs::File; use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt}; use crate::page::index::IndexPageUtility; -use crate::page::{IndexValue, PageId}; +use crate::page::{IndexValue, PageId, PageOverflowError}; use crate::{ align, align8, seek_to_page_start, Link, Persistable, SizeMeasurable, GENERAL_HEADER_SIZE, + INNER_PAGE_SIZE, }; pub fn get_index_page_size_from_data_length(length: usize) -> usize @@ -202,6 +203,27 @@ impl IndexPage { Self::read_value(file).await } + /// Rejects a slot write whose byte range would leave the page, so a bad + /// `value_index` (or an oversized serialized value) corrupts nothing. + /// + /// `offset` is relative to the page start and already includes the + /// general header. + fn check_value_write_bounds( + page_id: PageId, + offset: usize, + value_length: usize, + ) -> Result<(), PageOverflowError> { + let write_end_in_slot = offset + value_length - GENERAL_HEADER_SIZE; + if write_end_in_slot > INNER_PAGE_SIZE { + return Err(PageOverflowError { + page_id, + data_length: write_end_in_slot, + capacity: INNER_PAGE_SIZE, + }); + } + Ok(()) + } + fn get_value_offset(size: usize, value_index: usize) -> usize where T: Default + SizeMeasurable, @@ -238,11 +260,11 @@ impl IndexPage { rkyv::api::high::HighValidator<'a, rkyv::rancor::Error>, >, { - seek_to_page_start(file, page_id.0).await?; - let offset = Self::get_value_offset(size, value_index as usize); - file.seek(SeekFrom::Current(offset as i64)).await?; let bytes = rkyv::to_bytes::(&value)?; + Self::check_value_write_bounds(page_id, offset, bytes.len())?; + seek_to_page_start(file, page_id.0).await?; + file.seek(SeekFrom::Current(offset as i64)).await?; file.write_all(bytes.as_slice()).await?; if value_index != size as u16 - 1 { @@ -278,12 +300,12 @@ impl IndexPage { rkyv::api::high::HighValidator<'a, rkyv::rancor::Error>, >, { - seek_to_page_start(file, page_id.0).await?; - let offset = Self::get_value_offset(size, value_index as usize); - file.seek(SeekFrom::Current(offset as i64)).await?; let value = IndexValue::::default(); let bytes = rkyv::to_bytes::(&value)?; + Self::check_value_write_bounds(page_id, offset, bytes.len())?; + seek_to_page_start(file, page_id.0).await?; + file.seek(SeekFrom::Current(offset as i64)).await?; file.write_all(bytes.as_slice()).await?; Ok(()) @@ -391,6 +413,84 @@ mod tests { assert_eq!(new_page.index_values, page.index_values); } + #[tokio::test] + async fn persist_and_remove_value_reject_writes_past_the_page_slot() { + use super::{IndexPageUtility, PageOverflowError, SizedIndexPageUtility}; + + let path = std::env::temp_dir().join(format!( + "data_bucket_slot_write_bounds_{}.wt", + std::process::id() + )); + let mut file = tokio::fs::OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open(&path) + .await + .unwrap(); + + let value = IndexValue:: { + key: 7, + link: Default::default(), + }; + + // An in-bounds slot write works. + IndexPage::::persist_value(&mut file, 1.into(), 4, value.clone(), 3) + .await + .unwrap(); + // tokio's File buffers writes; flush so metadata() sees them. + tokio::io::AsyncWriteExt::flush(&mut file).await.unwrap(); + let length_after_valid_write = file.metadata().await.unwrap().len(); + + // A value index whose slot lies past the page must be rejected + // before anything is written. + let err = IndexPage::::persist_value(&mut file, 1.into(), 4, value, 2000) + .await + .unwrap_err(); + assert!( + err.downcast_ref::().is_some(), + "expected PageOverflowError, got: {err}" + ); + + let err = IndexPage::::remove_value(&mut file, 1.into(), 4, 2000) + .await + .unwrap_err(); + assert!( + err.downcast_ref::().is_some(), + "expected PageOverflowError, got: {err}" + ); + + // A utility larger than the page slot must be rejected too. + let utility = SizedIndexPageUtility:: { + size: 0, + node_id: IndexValue::default(), + current_index: 0, + current_length: 0, + slots: vec![0u16; 20_000], + }; + let err = as IndexPageUtility>::persist_index_page_utility( + &mut file, + 1.into(), + utility, + ) + .await + .unwrap_err(); + assert!( + err.downcast_ref::().is_some(), + "expected PageOverflowError, got: {err}" + ); + + // Nothing was written by the rejected operations. + assert_eq!( + file.metadata().await.unwrap().len(), + length_after_valid_write + ); + + drop(file); + std::fs::remove_file(&path).unwrap(); + } + #[test] fn test_split() { let mut page = IndexPage::::new( diff --git a/src/page/index/page_for_unsized.rs b/src/page/index/page_for_unsized.rs index 573bc15..de25d9d 100644 --- a/src/page/index/page_for_unsized.rs +++ b/src/page/index/page_for_unsized.rs @@ -14,9 +14,9 @@ use tokio::fs::File; use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt}; use crate::page::index::IndexPageUtility; -use crate::page::PageId; +use crate::page::{PageId, PageOverflowError}; use crate::{align8, VariableSizeMeasurable}; -use crate::{seek_to_page_start, IndexValue, SizeMeasurable, GENERAL_HEADER_SIZE}; +use crate::{seek_to_page_start, IndexValue, SizeMeasurable, GENERAL_HEADER_SIZE, INNER_PAGE_SIZE}; use crate::{Link, Persistable}; #[derive(Archive, Clone, Deserialize, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)] @@ -212,15 +212,25 @@ where rkyv::api::high::HighValidator<'a, rkyv::rancor::Error>, >, { + let bytes = rkyv::to_bytes::(&value)?; + let offset = current_offset as u64 + bytes.len() as u64; + // Values fill the page tail-first, so `offset` counts back from the + // page end: once it passes the inner-page budget the write would + // land in this page's header, or before it in the previous page. + if offset > INNER_PAGE_SIZE as u64 { + return Err(eyre::Report::new(PageOverflowError { + page_id, + data_length: offset as usize, + capacity: INNER_PAGE_SIZE, + })); + } + // We seek to page's end and will write values from tail. seek_to_page_start(file, page_id.0 + 1).await?; - - let bytes = rkyv::to_bytes::(&value)?; - let offset = current_offset + bytes.len() as u32; file.seek(SeekFrom::Current(-(offset as i64))).await?; file.write_all(bytes.as_slice()).await?; - Ok(offset) + Ok(offset as u32) } async fn read_value(file: &mut File, len: u16) -> eyre::Result> @@ -376,7 +386,76 @@ where #[cfg(test)] mod test { - use crate::{IndexValue, Link, Persistable, UnsizedIndexPage}; + use crate::page::PageOverflowError; + use crate::{IndexValue, Link, Persistable, UnsizedIndexPage, INNER_PAGE_SIZE}; + + #[tokio::test] + async fn persist_value_rejects_writes_leaving_the_page_slot() { + let path = std::env::temp_dir().join(format!( + "data_bucket_unsized_write_bounds_{}.wt", + std::process::id() + )); + let mut file = tokio::fs::OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open(&path) + .await + .unwrap(); + + let value = IndexValue:: { + key: "tail_first_value".to_string(), + link: Default::default(), + }; + + // An in-bounds tail-first write works. + UnsizedIndexPage::::persist_value(&mut file, 1.into(), 0, value.clone()) + .await + .unwrap(); + // tokio's File buffers writes; flush so metadata() sees them. + tokio::io::AsyncWriteExt::flush(&mut file).await.unwrap(); + let length_after_valid_write = file.metadata().await.unwrap().len(); + + // A current offset at the inner budget leaves no room: the write + // would land in the page header (or the previous page). + let err = UnsizedIndexPage::::persist_value( + &mut file, + 1.into(), + INNER_PAGE_SIZE as u32, + value.clone(), + ) + .await + .unwrap_err(); + assert!( + err.downcast_ref::().is_some(), + "expected PageOverflowError, got: {err}" + ); + + // A huge current offset used to wrap the u32 arithmetic and seek + // far outside the page; it must be rejected the same way. + let err = UnsizedIndexPage::::persist_value( + &mut file, + 1.into(), + u32::MAX - 4, + value, + ) + .await + .unwrap_err(); + assert!( + err.downcast_ref::().is_some(), + "expected PageOverflowError, got: {err}" + ); + + // Nothing was written by the rejected operations. + assert_eq!( + file.metadata().await.unwrap().len(), + length_after_valid_write + ); + + drop(file); + std::fs::remove_file(&path).unwrap(); + } #[test] fn to_bytes_and_back() { diff --git a/src/page/index/table_of_contents_page.rs b/src/page/index/table_of_contents_page.rs index f373f4b..0be9541 100644 --- a/src/page/index/table_of_contents_page.rs +++ b/src/page/index/table_of_contents_page.rs @@ -1,9 +1,61 @@ use rkyv::{Archive, Deserialize, Serialize}; +use std::collections::btree_map::Entry; use std::collections::BTreeMap; use std::fmt::Debug; use crate::page::PageId; -use crate::{align, Persistable, SizeMeasurable}; +use crate::{align, align_to, Persistable, SizeMeasurable, INNER_PAGE_SIZE}; + +/// Serialized size of a [`TableOfContentsPage`] with no records and no +/// empty pages: the `estimated_size` field itself plus the two empty +/// vectors. +pub const EMPTY_TABLE_OF_CONTENTS_PAGE_SIZE: usize = std::mem::size_of::() + 12; + +/// Error returned by the capacity-checked mutators of +/// [`TableOfContentsPage`] when adding a record would push the page's +/// serialized form past its page slot. +/// +/// The rejected key is handed back in [`Self::into_key`] so the caller can +/// place it on another page; [`Self::fits_empty_page`] tells whether such a +/// relocation can ever succeed. +#[derive(Debug)] +pub struct TableOfContentsOverflowError { + /// The key that was not added. + pub key: T, + /// Serialized size of the rejected record. + pub record_size: usize, + /// The page's estimated serialized size at the time of rejection. + pub estimated_size: usize, + /// The serialized-size budget of the page slot. + pub capacity: usize, +} + +impl TableOfContentsOverflowError { + /// Returns the rejected key so it can be placed on another page. + pub fn into_key(self) -> T { + self.key + } + + /// `true` when the record fits an empty page, so the caller can + /// relocate it to a fresh table-of-contents page. `false` means the + /// record can never fit any page slot and must be rejected upstream. + pub fn fits_empty_page(&self) -> bool { + EMPTY_TABLE_OF_CONTENTS_PAGE_SIZE + self.record_size <= self.capacity + } +} + +impl std::fmt::Display for TableOfContentsOverflowError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "table of contents record of {} bytes does not fit the page \ + (estimated size {} of {} bytes)", + self.record_size, self.estimated_size, self.capacity + ) + } +} + +impl std::error::Error for TableOfContentsOverflowError {} #[derive(Archive, Clone, Deserialize, Debug, Serialize)] pub struct TableOfContentsPage { @@ -21,7 +73,7 @@ where Self { records: BTreeMap::new(), empty_pages: vec![], - estimated_size: ::default_aligned_size() + 12, + estimated_size: EMPTY_TABLE_OF_CONTENTS_PAGE_SIZE, } } } @@ -36,6 +88,7 @@ struct TableOfContentsPagePersisted { impl Persistable for TableOfContentsPage where T: Clone + + SizeMeasurable + rkyv::Archive + for<'a> rkyv::Serialize< rkyv::rancor::Strategy< @@ -76,14 +129,58 @@ where let model: TableOfContentsPagePersisted = rkyv::deserialize::<_, rkyv::rancor::Error>(archived).expect("data should be valid"); let records = BTreeMap::from_iter(model.records); + // Recompute the accounting from the actual contents instead of + // trusting the stored field: pages written by versions with + // accounting bugs (0.5.2 and earlier) would otherwise carry their + // wrong estimated_size across the upgrade and mislead the capacity + // checks. + let estimated_size = Self::recompute_estimated_size(&records, &model.empty_pages); Self { records, - estimated_size: model.estimated_size, + estimated_size, empty_pages: model.empty_pages, } } } +impl TableOfContentsPage { + /// Serialized size of one `(key, PageId)` record. + /// + /// Mirrors `<(T, PageId) as SizeMeasurable>::aligned_size` without + /// needing to clone the key, so insertion and removal account records + /// with the same formula. + fn record_size(key: &T) -> usize + where + T: SizeMeasurable, + { + let len = key.aligned_size() + PageId::default().0.aligned_size(); + if let Some(key_align) = T::align() { + if key_align % 8 == 0 { + // rkyv pads the archived record out to the key's real + // alignment (16 for u128-likes), so round to it, not to 8. + return align_to(len, key_align); + } + } + align(len) + } + + /// Size accounting derived from the actual contents, used when loading + /// a persisted page: the stored `estimated_size` may carry accounting + /// bugs of the version that wrote it, and trusting it would let such an + /// error survive upgrades. + fn recompute_estimated_size(records: &BTreeMap, empty_pages: &[PageId]) -> usize + where + T: SizeMeasurable, + { + let mut estimated_size = EMPTY_TABLE_OF_CONTENTS_PAGE_SIZE; + for key in records.keys() { + estimated_size += Self::record_size(key); + } + estimated_size += empty_pages.len() * PageId::default().0.aligned_size(); + estimated_size + } +} + impl TableOfContentsPage where T: Debug + Ord + Eq, @@ -94,10 +191,64 @@ where pub fn insert(&mut self, val: T, page_id: PageId) where - T: SizeMeasurable + Clone, + T: SizeMeasurable, + { + // One map traversal via the entry API; the size arithmetic is done + // up front and nothing is cloned or serialized. + let record_size = Self::record_size(&val); + match self.records.entry(val) { + // Inserting over an existing key replaces its PageId in place, + // so the serialized page does not grow. + Entry::Occupied(mut entry) => { + entry.insert(page_id); + } + Entry::Vacant(entry) => { + entry.insert(page_id); + self.estimated_size += record_size; + } + } + } + + /// Capacity-checked [`Self::insert`]. + /// + /// Adds the record only when the page's serialized form stays within + /// its page slot ([`INNER_PAGE_SIZE`]). On overflow the page is left + /// untouched and the key is handed back in the error, so the caller can + /// place it on another page; [`TableOfContentsOverflowError::fits_empty_page`] + /// tells whether a fresh page can ever hold it. + /// + /// One map traversal; the capacity check is pure arithmetic and runs + /// before the map gains the record. + pub fn try_insert( + &mut self, + val: T, + page_id: PageId, + ) -> Result<(), TableOfContentsOverflowError> + where + T: SizeMeasurable, { - self.estimated_size += (val.clone(), page_id).aligned_size(); - let _ = self.records.insert(val, page_id); + let record_size = Self::record_size(&val); + match self.records.entry(val) { + // Replacing an existing key does not grow the page, so it + // always fits. + Entry::Occupied(mut entry) => { + entry.insert(page_id); + Ok(()) + } + Entry::Vacant(entry) => { + if self.estimated_size + record_size > INNER_PAGE_SIZE { + return Err(TableOfContentsOverflowError { + key: entry.into_key(), + record_size, + estimated_size: self.estimated_size, + capacity: INNER_PAGE_SIZE, + }); + } + entry.insert(page_id); + self.estimated_size += record_size; + Ok(()) + } + } } pub fn pop_empty_page(&mut self) -> Option @@ -125,6 +276,9 @@ where T: SizeMeasurable, { let id = self.remove_without_record(val); + // The removed page is recorded as empty, which grows the serialized + // `empty_pages` list by one `PageId`. + self.estimated_size += id.aligned_size(); self.empty_pages.push(id); id } @@ -133,20 +287,100 @@ where where T: SizeMeasurable, { - self.estimated_size -= align(val.aligned_size() + PageId::default().0.aligned_size()); - self.estimated_size += PageId::default().0.aligned_size(); + self.estimated_size -= Self::record_size(val); self.records .remove(val) .expect("value should be available if remove is called") } - pub fn update_key(&mut self, old_key: &T, new_key: T) -> Option<()> { - if let Some(id) = self.records.remove(old_key) { - self.records.insert(new_key, id); - return Some(()); + /// Re-keys the record at `old_key` to `new_key`, maintaining + /// `estimated_size` for the size difference between the two keys. + /// + /// Returns [`None`] (leaving the page untouched) when `old_key` is not + /// present. + /// + /// This variant performs no capacity check: a key that grows past the + /// page slot is accepted and only reported through `estimated_size`. + /// Use [`Self::try_update_key`] when the caller can relocate instead. + /// Two map traversals (removal of the old key, entry at the new key), + /// the minimum for touching two distinct keys. + pub fn update_key(&mut self, old_key: &T, new_key: T) -> Option<()> + where + T: SizeMeasurable, + { + let id = self.records.remove(old_key)?; + self.estimated_size -= Self::record_size(old_key); + let new_record_size = Self::record_size(&new_key); + match self.records.entry(new_key) { + // If `new_key` was already present, its record is replaced in + // place and the serialized page does not grow. + Entry::Occupied(mut entry) => { + entry.insert(id); + } + Entry::Vacant(entry) => { + entry.insert(id); + self.estimated_size += new_record_size; + } + } + Some(()) + } + + /// Capacity-checked [`Self::update_key`]. + /// + /// Applies the re-keying only when the page's serialized form stays + /// within its page slot ([`INNER_PAGE_SIZE`]); on overflow the page is + /// left untouched and `new_key` is handed back in the error, so the + /// caller can relocate the record instead. + /// + /// Returns `Ok(true)` when the key was updated and `Ok(false)` when + /// `old_key` is not present (the page stays untouched). + /// + /// Two map traversals on the accept and missing-key paths (removal of + /// the old key, entry at the new key); a rejected update restores the + /// removed record with a third. The capacity check is pure arithmetic + /// and runs before the map gains the new record. + pub fn try_update_key( + &mut self, + old_key: &T, + new_key: T, + ) -> Result> + where + T: SizeMeasurable, + { + let Some((old_entry_key, id)) = self.records.remove_entry(old_key) else { + return Ok(false); + }; + let old_record_size = Self::record_size(old_key); + let new_record_size = Self::record_size(&new_key); + match self.records.entry(new_key) { + // Re-keying onto an already existing key replaces that record + // in place: the page shrinks by the old record, so it always + // fits. + Entry::Occupied(mut entry) => { + entry.insert(id); + self.estimated_size -= old_record_size; + Ok(true) + } + Entry::Vacant(entry) => { + let projected = self.estimated_size - old_record_size + new_record_size; + if projected > INNER_PAGE_SIZE { + let key = entry.into_key(); + // Restore the removed record: a rejected update leaves + // the page as it was. + self.records.insert(old_entry_key, id); + return Err(TableOfContentsOverflowError { + key, + record_size: new_record_size, + estimated_size: self.estimated_size, + capacity: INNER_PAGE_SIZE, + }); + } + entry.insert(id); + self.estimated_size = projected; + Ok(true) + } } - None } pub fn contains(&self, val: &T) -> bool { @@ -172,7 +406,456 @@ where #[cfg(test)] mod test { - use crate::{Link, Persistable, TableOfContentsPage}; + use crate::{Link, Persistable, TableOfContentsPage, INNER_PAGE_SIZE}; + + fn link(offset: u32) -> Link { + Link { + page_id: 1.into(), + offset, + length: 32, + } + } + + #[test] + fn test_from_bytes_recomputes_stale_persisted_estimated_size() { + // Simulate a page persisted by 0.5.2, whose accounting bugs stored + // a wrong estimated_size: the load must recompute it from the + // actual records instead of trusting the stored field. + for bogus_estimated in [0usize, 3, 100_000] { + let stale = super::TableOfContentsPagePersisted { + records: vec![ + ((1u64, link(0)), crate::page::PageId::from(6)), + ((2u64, link(64)), crate::page::PageId::from(7)), + ], + empty_pages: vec![9.into()], + estimated_size: bogus_estimated, + }; + let bytes = rkyv::to_bytes::(&stale).unwrap(); + + let loaded = TableOfContentsPage::<(u64, Link)>::from_bytes(&bytes, 0); + assert_ne!(loaded.estimated_size(), bogus_estimated); + // The recomputed accounting matches what the page really + // serializes to. + assert_eq!(loaded.estimated_size(), loaded.as_bytes().as_ref().len()); + + // And it matches what fresh inserts would have produced. + let mut fresh = TableOfContentsPage::<(u64, Link)>::default(); + fresh.insert((1, link(0)), 6.into()); + fresh.insert((2, link(64)), 7.into()); + fresh.insert((3, link(128)), 9.into()); + fresh.remove(&(3, link(128))); + assert_eq!(loaded.estimated_size(), fresh.estimated_size()); + } + } + + #[test] + fn test_try_insert_bounds_real_serialized_size_for_16_aligned_keys() { + // Regression for the review blocker: with u128-family keys the size + // model under-counted every record by 8 bytes, so try_insert kept + // accepting records until the real archive was hundreds of bytes + // past the page slot. + let mut toc_page = TableOfContentsPage::::default(); + let mut i = 0u128; + while toc_page.try_insert(i, 1.into()).is_ok() { + i += 1; + assert!(i < 4096, "the page must eventually report itself full"); + } + let serialized = toc_page.as_bytes().as_ref().len(); + assert_eq!(serialized, toc_page.estimated_size()); + assert!(serialized <= INNER_PAGE_SIZE); + + let mut toc_page = TableOfContentsPage::<(u128, Link)>::default(); + let mut i = 0u128; + while toc_page.try_insert((i, link(i as u32)), 1.into()).is_ok() { + i += 1; + assert!(i < 4096, "the page must eventually report itself full"); + } + let serialized = toc_page.as_bytes().as_ref().len(); + assert_eq!(serialized, toc_page.estimated_size()); + assert!(serialized <= INNER_PAGE_SIZE); + } + + #[test] + fn test_estimated_size_is_an_upper_bound_under_adversarial_churn() { + // Property-style sweep with a deterministic LCG: string keys of + // adversarial lengths (the String model rounds to 4 bytes, so it may + // over-estimate, never under), mixed inserts, removals and key + // updates. The safety invariant of the capacity checks is + // serialized <= estimated at every step; the approximation slack + // stays below 4 bytes per record. + let mut state = 0x2545F4914F6CDD1Du64; + let mut next = move || { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + (state >> 33) as u32 + }; + + let mut toc_page = TableOfContentsPage::<(String, Link)>::default(); + let mut keys: Vec<(String, Link)> = vec![]; + let mut record_count = 0usize; + for step in 0..600u32 { + let op = next() % 4; + if op < 2 || keys.is_empty() { + let len = (next() % 48) as usize; + let key = (format!("{step:0>len$}"), link(step)); + if toc_page.try_insert(key.clone(), (step + 1).into()).is_ok() { + keys.push(key); + record_count += 1; + } + } else if op == 2 { + let key = keys.swap_remove((next() as usize) % keys.len()); + toc_page.remove(&key); + record_count -= 1; + } else { + let old = keys.swap_remove((next() as usize) % keys.len()); + let len = (next() % 48) as usize; + let new_key = (format!("{step:0>len$}"), link(step)); + match toc_page.try_update_key(&old, new_key.clone()) { + Ok(true) => keys.push(new_key), + Ok(false) => unreachable!("old key is always present"), + Err(_) => keys.push(old), + } + } + + let serialized = toc_page.as_bytes().as_ref().len(); + let estimated = toc_page.estimated_size(); + assert!( + serialized <= estimated, + "under-estimate at step {step}: serialized {serialized} > estimated {estimated}" + ); + assert!( + estimated - serialized <= 4 * (record_count + 2), + "slack blew past the documented bound at step {step}: \ + serialized {serialized}, estimated {estimated}, records {record_count}" + ); + } + } + + #[test] + fn test_try_update_key_onto_existing_key_replaces_in_place() { + let mut toc_page = TableOfContentsPage::<(u32, Link)>::default(); + toc_page.insert((1, link(0)), 6.into()); + toc_page.insert((2, link(64)), 7.into()); + + // Re-keying onto an existing key keeps the accounting exact and + // repoints the surviving record. + assert!(toc_page + .try_update_key(&(1, link(0)), (2, link(64))) + .unwrap()); + assert!(!toc_page.contains(&(1, link(0)))); + assert_eq!(toc_page.get(&(2, link(64))), Some(6.into())); + assert_eq!( + toc_page.as_bytes().as_ref().len(), + toc_page.estimated_size() + ); + + // Re-keying a record onto itself is a no-op that still succeeds. + assert!(toc_page + .try_update_key(&(2, link(64)), (2, link(64))) + .unwrap()); + assert_eq!(toc_page.get(&(2, link(64))), Some(6.into())); + assert_eq!( + toc_page.as_bytes().as_ref().len(), + toc_page.estimated_size() + ); + } + + #[test] + fn test_remove_without_record_keeps_estimated_size_exact() { + let mut toc_page = TableOfContentsPage::<(u32, Link)>::default(); + let empty_size = toc_page.as_bytes().as_ref().len(); + assert_eq!(empty_size, toc_page.estimated_size()); + + toc_page.insert((1, link(0)), 6.into()); + assert_eq!( + toc_page.as_bytes().as_ref().len(), + toc_page.estimated_size() + ); + + // No empty-page record is produced, so the size must return exactly + // to the empty-page baseline. It used to stay one PageId too big. + toc_page.remove_without_record(&(1, link(0))); + assert_eq!( + toc_page.as_bytes().as_ref().len(), + toc_page.estimated_size() + ); + assert_eq!(toc_page.estimated_size(), empty_size); + } + + #[test] + fn test_remove_keeps_estimated_size_exact() { + let mut toc_page = TableOfContentsPage::<(u32, Link)>::default(); + toc_page.insert((1, link(0)), 6.into()); + toc_page.insert((2, link(64)), 7.into()); + + // `remove` records the freed page in `empty_pages`, which itself + // takes serialized space. + toc_page.remove(&(1, link(0))); + assert_eq!( + toc_page.as_bytes().as_ref().len(), + toc_page.estimated_size() + ); + + // Reusing the empty page gives that space back. + assert_eq!(toc_page.pop_empty_page(), Some(6.into())); + assert_eq!( + toc_page.as_bytes().as_ref().len(), + toc_page.estimated_size() + ); + } + + #[test] + fn test_insert_over_existing_key_keeps_estimated_size_exact() { + let mut toc_page = TableOfContentsPage::<(u32, Link)>::default(); + toc_page.insert((1, link(0)), 6.into()); + let size_after_first = toc_page.estimated_size(); + + // Re-pointing the same key to another page replaces the record in + // place; it used to be counted as a second record. + toc_page.insert((1, link(0)), 7.into()); + assert_eq!(toc_page.get(&(1, link(0))), Some(7.into())); + assert_eq!(toc_page.estimated_size(), size_after_first); + assert_eq!( + toc_page.as_bytes().as_ref().len(), + toc_page.estimated_size() + ); + } + + #[test] + fn test_remove_accounts_8_aligned_keys_like_insert() { + // (u64, Link) records are 8-aligned, so their serialized record size + // is align8-rounded. Removal used to subtract the 4-aligned size, + // leaving estimated_size drifting upward on every insert/remove + // cycle. + let mut toc_page = TableOfContentsPage::<(u64, Link)>::default(); + let empty_size = toc_page.as_bytes().as_ref().len(); + + toc_page.insert((128, link(0)), 6.into()); + assert_eq!( + toc_page.as_bytes().as_ref().len(), + toc_page.estimated_size() + ); + + toc_page.remove_without_record(&(128, link(0))); + assert_eq!( + toc_page.as_bytes().as_ref().len(), + toc_page.estimated_size() + ); + assert_eq!(toc_page.estimated_size(), empty_size); + } + + #[test] + fn test_update_key_keeps_estimated_size_exact() { + fn assert_exact(toc_page: &TableOfContentsPage<(String, Link)>) { + assert_eq!( + toc_page.as_bytes().as_ref().len(), + toc_page.estimated_size() + ); + } + + let mut toc_page = TableOfContentsPage::<(String, Link)>::default(); + toc_page.insert(("key_0001".to_string(), link(0)), 6.into()); + assert_exact(&toc_page); + + // Growing the key must grow estimated_size with it; it used not to + // be accounted at all. + let grown = "key_0001_grown_by_a_long_suffix_0001".to_string(); + toc_page + .update_key(&("key_0001".to_string(), link(0)), (grown.clone(), link(0))) + .unwrap(); + assert_exact(&toc_page); + + // Shrinking it must give the space back. + toc_page + .update_key(&(grown, link(0)), ("key_0001".to_string(), link(0))) + .unwrap(); + assert_exact(&toc_page); + + // Updating onto an already existing key replaces that record in + // place. + toc_page.insert(("key_0002".to_string(), link(64)), 7.into()); + toc_page + .update_key( + &("key_0001".to_string(), link(0)), + ("key_0002".to_string(), link(64)), + ) + .unwrap(); + assert_exact(&toc_page); + assert_eq!( + toc_page.get(&("key_0002".to_string(), link(64))), + Some(6.into()) + ); + + // A missing old key leaves the page untouched. + let before = toc_page.estimated_size(); + assert!(toc_page + .update_key( + &("missing".to_string(), link(0)), + ("whatever".to_string(), link(0)), + ) + .is_none()); + assert_eq!(toc_page.estimated_size(), before); + assert_exact(&toc_page); + } + + #[test] + // Key strings are chosen with len % 4 == 0 (or <= 8): the String size + // model in SizeMeasurable rounds each string to 4 bytes, which is the + // documented accuracy bound of estimated_size for other lengths. + fn test_estimated_size_stays_exact_across_churn() { + fn assert_exact(toc_page: &TableOfContentsPage<(String, Link)>) { + assert_eq!( + toc_page.as_bytes().as_ref().len(), + toc_page.estimated_size() + ); + } + + let mut toc_page = TableOfContentsPage::<(String, Link)>::default(); + + for i in 0..32u32 { + toc_page.insert((format!("key_{i:04}"), link(i)), (i + 1).into()); + assert_exact(&toc_page); + } + + // Grow half of the keys through updates. + for i in 0..16u32 { + toc_page + .update_key( + &(format!("key_{i:04}"), link(i)), + (format!("key_{i:04}_grown_by_a_long_suffix_"), link(i)), + ) + .unwrap(); + assert_exact(&toc_page); + } + + // Shrink them back. + for i in 0..16u32 { + toc_page + .update_key( + &(format!("key_{i:04}_grown_by_a_long_suffix_"), link(i)), + (format!("key_{i:04}"), link(i)), + ) + .unwrap(); + assert_exact(&toc_page); + } + + // Remove a mix, with and without empty-page records. + for i in 0..8u32 { + toc_page.remove(&(format!("key_{i:04}"), link(i))); + assert_exact(&toc_page); + } + for i in 8..16u32 { + toc_page.remove_without_record(&(format!("key_{i:04}"), link(i))); + assert_exact(&toc_page); + } + while toc_page.pop_empty_page().is_some() { + assert_exact(&toc_page); + } + } + + #[test] + fn test_try_insert_never_reports_success_past_the_page_slot() { + let mut toc_page = TableOfContentsPage::<(u32, Link)>::default(); + + // Fill until the page reports itself full. + let mut rejected_at = None; + for i in 0..2048u32 { + match toc_page.try_insert((i, link(i)), (i + 1).into()) { + Ok(()) => { + assert!(toc_page.estimated_size() <= INNER_PAGE_SIZE); + } + Err(err) => { + rejected_at = Some((i, err)); + break; + } + } + } + let (i, err) = rejected_at.expect("the page must eventually report itself full"); + + // The serialized page still fits its slot exactly at the point of + // rejection. + let serialized = toc_page.as_bytes().as_ref().len(); + assert_eq!(serialized, toc_page.estimated_size()); + assert!(serialized <= INNER_PAGE_SIZE); + + // The page was left untouched by the rejected insert, and the key + // came back for relocation. + assert!(!toc_page.contains(&(i, link(i)))); + assert!(err.fits_empty_page()); + assert_eq!(err.into_key(), (i, link(i))); + + // Replacing an existing key does not grow the page, so it is still + // accepted on a full page. + let size_before = toc_page.estimated_size(); + toc_page + .try_insert((0, link(0)), 999.into()) + .expect("replacement must fit"); + assert_eq!(toc_page.estimated_size(), size_before); + assert_eq!(toc_page.get(&(0, link(0))), Some(999.into())); + } + + #[test] + fn test_try_insert_rejects_record_that_can_never_fit() { + let mut toc_page = TableOfContentsPage::<(String, Link)>::default(); + let oversized = "x".repeat(INNER_PAGE_SIZE); + + let err = toc_page + .try_insert((oversized.clone(), link(0)), 6.into()) + .expect_err("a record larger than the page slot must be rejected"); + // No fresh page can hold it either: the caller must fail upstream + // instead of relocating forever. + assert!(!err.fits_empty_page()); + assert!(!toc_page.contains(&(oversized, link(0)))); + assert_eq!( + toc_page.as_bytes().as_ref().len(), + toc_page.estimated_size() + ); + } + + #[test] + fn test_try_update_key_rejects_growth_past_the_page_slot() { + let mut toc_page = TableOfContentsPage::<(String, Link)>::default(); + toc_page.insert(("key_0001".to_string(), link(0)), 6.into()); + + let grown = "x".repeat(INNER_PAGE_SIZE); + let err = toc_page + .try_update_key(&("key_0001".to_string(), link(0)), (grown.clone(), link(0))) + .expect_err("growth past the page slot must be rejected"); + assert!(!err.fits_empty_page()); + assert_eq!(err.into_key(), (grown, link(0))); + + // The page is untouched: the old record is still there and the + // accounting is still exact. + assert_eq!( + toc_page.get(&("key_0001".to_string(), link(0))), + Some(6.into()) + ); + assert_eq!( + toc_page.as_bytes().as_ref().len(), + toc_page.estimated_size() + ); + + // A fitting update through the checked variant still works. + assert!(toc_page + .try_update_key( + &("key_0001".to_string(), link(0)), + ("key_0002".to_string(), link(0)), + ) + .unwrap()); + assert_eq!( + toc_page.as_bytes().as_ref().len(), + toc_page.estimated_size() + ); + + // A missing old key reports Ok(false) and changes nothing. + assert!(!toc_page + .try_update_key( + &("missing".to_string(), link(0)), + ("whatever".to_string(), link(0)), + ) + .unwrap()); + } #[test] fn test_sizes() { diff --git a/src/page/mod.rs b/src/page/mod.rs index cfa2bfd..17cfec1 100644 --- a/src/page/mod.rs +++ b/src/page/mod.rs @@ -16,7 +16,8 @@ pub use data::DataPage; pub use header::{GeneralHeader, DATA_VERSION}; pub use index::{ get_index_page_size_from_data_length, IndexPage, IndexPageUtility, IndexValue, - TableOfContentsPage, UnsizedIndexPage, UnsizedIndexPageUtility, + TableOfContentsOverflowError, TableOfContentsPage, UnsizedIndexPage, UnsizedIndexPageUtility, + EMPTY_TABLE_OF_CONTENTS_PAGE_SIZE, }; //pub use iterators::{DataIterator, LinksIterator}; pub use space_info::{Interval, SpaceInfoPage}; @@ -24,7 +25,7 @@ pub use ty::PageType; pub use util::{ map_data_pages_to_general, parse_data_page, parse_data_pages_batch, parse_general_header_by_index, parse_page, parse_pages_batch, parse_space_info, persist_page, - persist_pages_batch, seek_by_link, seek_to_page_start, update_at, + persist_pages_batch, seek_by_link, seek_to_page_start, update_at, PageOverflowError, }; // TODO: Move to config diff --git a/src/page/util.rs b/src/page/util.rs index 617cf23..af1dd77 100644 --- a/src/page/util.rs +++ b/src/page/util.rs @@ -8,7 +8,36 @@ use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt}; use super::SpaceInfoPage; use crate::page::header::GeneralHeader; use crate::page::ty::PageType; -use crate::{DataPage, GeneralPage, Link, Persistable, GENERAL_HEADER_SIZE, PAGE_SIZE}; +use crate::page::PageId; +use crate::{ + DataPage, GeneralPage, Link, Persistable, GENERAL_HEADER_SIZE, INNER_PAGE_SIZE, PAGE_SIZE, +}; + +/// Returned when a write into a page would not fit the page slot: letting +/// it through would spill past a [`PAGE_SIZE`] boundary and corrupt a +/// neighboring page (or, for tail-first writes, this page's own header). +#[derive(Debug)] +pub struct PageOverflowError { + /// The page whose write is over budget. + pub page_id: PageId, + /// Inner-page bytes the write needs (for slot writes, where the write + /// would end within the slot). + pub data_length: usize, + /// The slot budget for inner data ([`INNER_PAGE_SIZE`]). + pub capacity: usize, +} + +impl std::fmt::Display for PageOverflowError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "page {} write needs {} bytes, exceeding the {}-byte page slot", + self.page_id, self.data_length, self.capacity + ) + } +} + +impl std::error::Error for PageOverflowError {} pub fn map_data_pages_to_general( pages: Vec>, @@ -68,7 +97,17 @@ where T: Persistable + Send + Sync, { let inner_bytes = page.inner.as_bytes(); - page.header.data_length = inner_bytes.as_ref().len() as u32; + let inner_length = inner_bytes.as_ref().len(); + // An over-budget page must fail here, in its own persist, instead of + // silently corrupting the neighboring page. + if inner_length > INNER_PAGE_SIZE { + return Err(eyre::Report::new(PageOverflowError { + page_id: page.header.page_id, + data_length: inner_length, + capacity: INNER_PAGE_SIZE, + })); + } + page.header.data_length = inner_length as u32; file.write_all(page.header.as_bytes().as_ref()).await?; file.write_all(inner_bytes.as_ref()).await?; Ok(()) @@ -94,16 +133,25 @@ where } } +/// Byte offset of the page with the given index, computed in `u64`. +/// +/// The arithmetic must never be done in `u32`: with the default 16 KiB +/// [`PAGE_SIZE`], any index past 262 143 puts the page start beyond 4 GiB, +/// and a `u32` multiply silently wraps the offset back into the start of +/// the file. +pub(crate) fn page_start_offset(index: u32) -> u64 { + index as u64 * PAGE_SIZE as u64 +} + pub async fn seek_to_page_start(file: &mut File, index: u32) -> eyre::Result<()> { - file.seek(SeekFrom::Start(index as u64 * PAGE_SIZE as u64)) - .await?; + file.seek(SeekFrom::Start(page_start_offset(index))).await?; Ok(()) } async fn seek_to_page_start_relatively(file: &mut File, index: u32) -> eyre::Result<()> { let curr_position = file.stream_position().await?; file.seek(SeekFrom::Current( - (index * PAGE_SIZE as u32) as i64 - curr_position as i64, + page_start_offset(index) as i64 - curr_position as i64, )) .await?; Ok(()) @@ -131,7 +179,9 @@ pub async fn update_at( )); } - if (link.offset + link.length) > DATA_LENGTH { + // Sum in u64: `offset + length` in u32 can wrap past 4 GiB and slip + // under the bound, letting the write land outside the page. + if link.offset as u64 + link.length as u64 > DATA_LENGTH as u64 { return Err(eyre!( "Link range (offset: {}, length: {}) exceeds data bounds ({})", link.offset, @@ -409,6 +459,204 @@ pub async fn parse_space_info( // Ok(result) // } +#[cfg(test)] +mod tests { + use super::{page_start_offset, parse_data_pages_batch, persist_pages_batch}; + use crate::page::header::GeneralHeader; + use crate::page::ty::PageType; + use crate::{DataPage, GeneralPage, DATA_VERSION, INNER_PAGE_SIZE, PAGE_SIZE}; + + /// First page index whose start offset no longer fits in `u32`. + const FIRST_PAGE_PAST_4_GIB: u32 = (u32::MAX / PAGE_SIZE as u32) + 1; + + #[test] + fn page_start_offset_is_computed_in_u64() { + assert_eq!( + page_start_offset(FIRST_PAGE_PAST_4_GIB), + FIRST_PAGE_PAST_4_GIB as u64 * PAGE_SIZE as u64 + ); + assert!(page_start_offset(FIRST_PAGE_PAST_4_GIB) > u32::MAX as u64); + // The largest possible page id must be addressable too. + assert_eq!( + page_start_offset(u32::MAX), + u32::MAX as u64 * PAGE_SIZE as u64 + ); + // The old `u32` arithmetic wrapped this offset back into the first + // pages of the file. + assert_ne!( + page_start_offset(FIRST_PAGE_PAST_4_GIB), + FIRST_PAGE_PAST_4_GIB.wrapping_mul(PAGE_SIZE as u32) as u64 + ); + } + + fn data_page_with_marker(marker: &[u8]) -> DataPage { + let mut data = [0u8; INNER_PAGE_SIZE]; + data[..marker.len()].copy_from_slice(marker); + DataPage { + length: marker.len() as u32, + data, + } + } + + #[tokio::test] + async fn persist_page_rejects_inner_data_past_the_page_slot() { + // A data page whose buffer is larger than the slot budget can hand + // persist more bytes than fit between two page starts. + const OVERSIZED: usize = crate::PAGE_SIZE + 128; + + let path = std::env::temp_dir().join(format!( + "data_bucket_persist_overflow_{}.wt", + std::process::id() + )); + let mut file = tokio::fs::OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open(&path) + .await + .unwrap(); + + let mut page = GeneralPage { + header: GeneralHeader { + data_version: DATA_VERSION, + space_id: 1.into(), + page_id: 1.into(), + previous_id: 0.into(), + next_id: 0.into(), + page_type: PageType::Data, + data_length: 0, + }, + inner: DataPage { + length: OVERSIZED as u32, + data: [7u8; OVERSIZED], + }, + }; + + let err = super::persist_page(&mut page, &mut file).await.unwrap_err(); + assert!( + err.downcast_ref::().is_some(), + "expected PageOverflowError, got: {err}" + ); + + // Nothing may have been written: the neighboring page is the one an + // unchecked write would have corrupted. + assert_eq!(file.metadata().await.unwrap().len(), 0); + + // A page that fits its slot still persists. + page.inner.length = 64; + super::persist_page(&mut page, &mut file).await.unwrap(); + + drop(file); + std::fs::remove_file(&path).unwrap(); + } + + #[tokio::test] + async fn update_at_rejects_offset_plus_length_wrapping_u32() { + let path = std::env::temp_dir().join(format!( + "data_bucket_update_at_wrap_{}.wt", + std::process::id() + )); + let mut file = tokio::fs::OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open(&path) + .await + .unwrap(); + + // In u32, offset + length wraps to 5 and used to pass the bounds + // check, sending the write far outside the page. + let link = crate::Link { + page_id: 1.into(), + offset: u32::MAX - 2, + length: 8, + }; + let err = super::update_at::<100>(&mut file, link, &[1, 2, 3, 4, 5, 6, 7, 8]) + .await + .unwrap_err(); + assert!(err.to_string().contains("exceeds data bounds")); + + drop(file); + std::fs::remove_file(&path).unwrap(); + } + + #[tokio::test] + async fn batch_persist_and_parse_address_pages_past_4_gib() { + const FIRST_MARKER: &[u8] = b"FIRSTPG!"; + const BOUNDARY_MARKER: &[u8] = b"BOUNDARY"; + + let path = std::env::temp_dir().join(format!( + "data_bucket_seek_past_4gib_{}.wt", + std::process::id() + )); + let mut file = tokio::fs::OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open(&path) + .await + .unwrap(); + + let first_page = GeneralPage { + header: GeneralHeader { + data_version: DATA_VERSION, + space_id: 1.into(), + page_id: 1.into(), + previous_id: 0.into(), + next_id: 2.into(), + page_type: PageType::Data, + data_length: 0, + }, + inner: data_page_with_marker(FIRST_MARKER), + }; + let boundary_page = GeneralPage { + header: GeneralHeader { + data_version: DATA_VERSION, + space_id: 1.into(), + page_id: FIRST_PAGE_PAST_4_GIB.into(), + previous_id: 1.into(), + next_id: 0.into(), + page_type: PageType::Data, + data_length: 0, + }, + inner: data_page_with_marker(BOUNDARY_MARKER), + }; + + persist_pages_batch(vec![first_page, boundary_page], &mut file) + .await + .unwrap(); + + // The boundary page must have been written past 4 GiB (the file is + // sparse, so this stays cheap), not wrapped back onto the first pages. + // tokio's File buffers writes; flush so metadata() sees them. + tokio::io::AsyncWriteExt::flush(&mut file).await.unwrap(); + let file_length = file.metadata().await.unwrap().len(); + assert!(file_length > page_start_offset(FIRST_PAGE_PAST_4_GIB)); + + let pages = parse_data_pages_batch::<{ PAGE_SIZE as u32 }, INNER_PAGE_SIZE>( + &mut file, + vec![1, FIRST_PAGE_PAST_4_GIB], + ) + .await + .unwrap(); + + assert_eq!(pages.len(), 2); + assert_eq!(pages[0].header.page_id, 1.into()); + assert_eq!(&pages[0].inner.data[..FIRST_MARKER.len()], FIRST_MARKER); + assert_eq!(pages[1].header.page_id, FIRST_PAGE_PAST_4_GIB.into()); + assert_eq!( + &pages[1].inner.data[..BOUNDARY_MARKER.len()], + BOUNDARY_MARKER + ); + + drop(file); + std::fs::remove_file(&path).unwrap(); + } +} + // #[cfg(test)] // pub mod test { // use std::collections::HashMap; diff --git a/src/util/mod.rs b/src/util/mod.rs index 2b05af2..b5be972 100644 --- a/src/util/mod.rs +++ b/src/util/mod.rs @@ -2,4 +2,4 @@ mod persistable; mod sized; pub use persistable::{access_archived, Persistable}; -pub use sized::{align, align8, align_vec, SizeMeasurable, VariableSizeMeasurable}; +pub use sized::{align, align8, align_to, align_vec, SizeMeasurable, VariableSizeMeasurable}; diff --git a/src/util/sized.rs b/src/util/sized.rs index 8a042ae..444fca6 100644 --- a/src/util/sized.rs +++ b/src/util/sized.rs @@ -22,6 +22,36 @@ pub const fn align8(len: usize) -> usize { } } +/// Rounds `len` up to a multiple of `alignment`. +pub const fn align_to(len: usize, alignment: usize) -> usize { + if len.is_multiple_of(alignment) { + len + } else { + (len / alignment + 1) * alignment + } +} + +/// The widest 8-or-more-byte alignment claimed by the two member types, if +/// any. +/// +/// rkyv pads a compound value out to its widest member alignment (8 for +/// `u64`-likes, 16 for `u128`/`i128`), so compound size models must round +/// to the real member alignment. Rounding a 16-aligned member to only 8 +/// under-counts every archived record by up to 8 bytes. +fn wide_member_align() -> Option +where + T1: SizeMeasurable, + T2: SizeMeasurable, +{ + let mut wide = None; + for member_align in [T1::align(), T2::align()].into_iter().flatten() { + if member_align % 8 == 0 && Some(member_align) > wide { + wide = Some(member_align); + } + } + wide +} + pub fn align_vec(mut v: AlignedVec) -> AlignedVec { if v.len() != align(v.len()) { let count = align(v.len()) - v.len(); @@ -123,31 +153,15 @@ where T2: SizeMeasurable, { fn aligned_size(&self) -> usize { - if let Some(align) = T1::align() { - if align % 8 == 0 { - return align8(self.0.aligned_size() + self.1.aligned_size()); - } - } - if let Some(align) = T2::align() { - if align % 8 == 0 { - return align8(self.0.aligned_size() + self.1.aligned_size()); - } + let len = self.0.aligned_size() + self.1.aligned_size(); + match wide_member_align::() { + Some(member_align) => align_to(len, member_align), + None => align(len), } - align(self.0.aligned_size() + self.1.aligned_size()) } fn align() -> Option { - if let Some(align) = T1::align() { - if align % 8 == 0 { - return Some(8); - } - } - if let Some(align) = T2::align() { - if align % 8 == 0 { - return Some(8); - } - } - None + wide_member_align::() } } @@ -358,6 +372,58 @@ mod test { ) } + #[test] + fn test_16_aligned_members_measure_like_rkyv() { + // rkyv pads compounds out to 16 for u128/i128 members; the model + // used to round to 8 and under-count every archived record. + let t = (u128::MAX, Link::default()); + assert_eq!( + t.aligned_size(), + to_bytes::(&t).unwrap().len() + ); + let t = (u128::MAX, u64::MAX); + assert_eq!( + t.aligned_size(), + to_bytes::(&t).unwrap().len() + ); + let t = (i128::MAX, Link::default()); + assert_eq!( + t.aligned_size(), + to_bytes::(&t).unwrap().len() + ); + assert_eq!(<(u128, Link) as SizeMeasurable>::align(), Some(16)); + assert_eq!(<(u128, u64) as SizeMeasurable>::align(), Some(16)); + assert_eq!(<(u64, Link) as SizeMeasurable>::align(), Some(8)); + + let v = IndexValue { + key: u128::MAX, + link: Link::default(), + }; + assert_eq!( + v.aligned_size(), + to_bytes::(&v).unwrap().len() + ); + let v = IndexValue { + key: (u128::MAX, u64::MAX), + link: Link::default(), + }; + assert_eq!( + v.aligned_size(), + to_bytes::(&v).unwrap().len() + ); + + // The archived element stride of a vec must match too, so per-record + // accounting stays exact for pages of such records. + let vec = vec![(u128::MAX, Link::default()); 100]; + let one = (u128::MAX, Link::default()).aligned_size(); + let bytes = to_bytes::(&vec).unwrap().len(); + let bytes_two_hundred = + to_bytes::(&vec![(u128::MAX, Link::default()); 200]) + .unwrap() + .len(); + assert_eq!(bytes_two_hundred - bytes, 100 * one); + } + #[test] fn test_tuple() { let t = (u64::MAX, Link::default());