From de6c3c3bc7f2f397386912b4c135c50c8b05e5d3 Mon Sep 17 00:00:00 2001 From: Byron Date: Sun, 30 Aug 2026 19:50:06 +0200 Subject: [PATCH] perf: preallocate snapshot v2 imports MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Snapshot V2 prefixes the record stream with the total node count and native path-name bytes. Interactive imports use them for exact reservations while V1 remains readable and declared requirements are validated. ## `~/dev` benchmark `dua i --import SNAPSHOT --once`, release builds; one frozen V1 traversal converted to V2 for identical contents. | Metric | V1 | V2 | Change | | --- | ---: | ---: | ---: | | Median peak RSS (5 runs) | 369,868,800 B | 230,473,728 B | -37.7% | | Mean wall time (20 runs) | 430.5 ± 1.8 ms | 438.1 ± 2.8 ms | +1.8% | | Compressed snapshot size | 38,215,358 B | 38,215,404 B | +46 B | Assisted-by: GPT 5.6 Co-authored-by: GPT 5.6 --- src/snapshot/mod.rs | 113 ++++++++++++++++++++++++++++++++---------- src/snapshot/tests.rs | 21 +++++++- src/traverse.rs | 23 +++++++++ tests/snapshot.rs | 27 ++++++++-- 4 files changed, 154 insertions(+), 30 deletions(-) diff --git a/src/snapshot/mod.rs b/src/snapshot/mod.rs index 0c03bc9e..c494899d 100644 --- a/src/snapshot/mod.rs +++ b/src/snapshot/mod.rs @@ -11,7 +11,8 @@ use std::{ }; const MAGIC: &[u8; 8] = b"DUASNAP\0"; -const VERSION: u16 = 1; +const VERSION: u16 = 2; +const VERSION_1: u16 = 1; const HEADER_LEN: usize = 12; const MAX_NAME_LEN: usize = 1024 * 1024; const MAX_RECORD_LEN: usize = 2 * 1024 * 1024; @@ -79,10 +80,12 @@ enum SnapshotReader { struct Decoder { reader: HashingReader>>, + allocation: Option<(usize, usize)>, open_nodes: Vec, sibling_names: Vec>, record: Vec, node_count: u64, + name_bytes: u64, total_size: u128, total_entries: u64, summary: Option, @@ -285,7 +288,7 @@ impl Replay { } } -/// Write `traversal` as a deterministic version-1 snapshot, optionally compressed with zlib. +/// Write `traversal` as a deterministic version-2 snapshot, optionally compressed with zlib. /// /// `roots` must contain the traversal's top-level nodes in original input order. pub fn write( @@ -305,13 +308,6 @@ pub fn write( } fn write_raw(writer: impl Write, traversal: &Traversal, roots: &[TreeIndex]) -> Result<()> { - let mut writer = gix::hash::io::Write::new(BufWriter::new(writer), gix::hash::Kind::Sha256); - let mut header = [0; HEADER_LEN]; - header[..MAGIC.len()].copy_from_slice(MAGIC); - header[8..10].copy_from_slice(&VERSION.to_le_bytes()); - header[10] = PATH_ENCODING; - writer.write_all(&header)?; - let mut seen_roots = HashSet::with_capacity(roots.len()); for &root in roots { if !seen_roots.insert(root) { @@ -320,6 +316,32 @@ fn write_raw(writer: impl Write, traversal: &Traversal, roots: &[TreeIndex]) -> } drop(seen_roots); + let mut stack = roots.to_vec(); + let mut stored_node_count = 0u64; + let mut stored_name_bytes = 0u64; + while let Some(index) = stack.pop() { + let name = traversal + .tree + .native_name(index) + .ok_or_else(|| anyhow!("snapshot node {index:?} does not exist"))?; + stored_node_count = stored_node_count + .checked_add(1) + .context("snapshot contains too many nodes")?; + stored_name_bytes = stored_name_bytes + .checked_add(u64::try_from(name.len()).context("snapshot names are too long")?) + .context("snapshot names are too long")?; + stack.extend(traversal.tree.children(index)); + } + + let mut writer = gix::hash::io::Write::new(BufWriter::new(writer), gix::hash::Kind::Sha256); + let mut header = [0; HEADER_LEN]; + header[..MAGIC.len()].copy_from_slice(MAGIC); + header[8..10].copy_from_slice(&VERSION.to_le_bytes()); + header[10] = PATH_ENCODING; + writer.write_all(&header)?; + write_uleb128(&mut writer, u128::from(stored_node_count))?; + write_uleb128(&mut writer, u128::from(stored_name_bytes))?; + let mut stack = Vec::new(); let mut record = Vec::new(); let mut node_count = 0u64; @@ -414,13 +436,20 @@ fn write_raw(writer: impl Write, traversal: &Traversal, roots: &[TreeIndex]) -> Ok(()) } -/// Read and fully verify a version-1 snapshot before returning its traversal. +/// Read and fully verify a version-1 or version-2 snapshot before returning its traversal. pub fn read(reader: impl Read) -> Result { + let mut decoder = Decoder::new(reader)?; let mut traversal = Traversal::new(); traversal.cost = Some(Duration::ZERO); + if let Some((node_count, name_bytes)) = decoder.allocation { + traversal + .tree + .try_reserve_exact(node_count, name_bytes) + .map_err(|err| anyhow!("could not preallocate snapshot tree: {err}"))?; + } let mut parents = Vec::new(); let mut roots = Vec::new(); - let summary = decode(reader, |entry| { + while let Some(entry) = decoder.next_entry(None)? { parents.truncate(entry.depth); let parent = parents.last().copied().unwrap_or(traversal.root_index); let node = traversal @@ -437,8 +466,10 @@ pub fn read(reader: impl Read) -> Result { .try_reserve(1) .context("could not grow snapshot ancestor stack")?; parents.push(node); - Ok(()) - })?; + } + let summary = decoder + .summary + .context("snapshot ended without a summary")?; traversal .tree @@ -450,17 +481,6 @@ pub fn read(reader: impl Read) -> Result { Ok(Snapshot { traversal, roots }) } -fn decode( - reader: impl Read, - mut on_entry: impl for<'entry> FnMut(DecodedEntry<'entry>) -> Result<()>, -) -> Result { - let mut decoder = Decoder::new(reader)?; - while let Some(entry) = decoder.next_entry(None)? { - on_entry(entry)?; - } - Ok(decoder.summary.expect("end of snapshot stores its summary")) -} - impl Decoder { fn new(reader: R) -> Result { let mut reader = HashingReader::new(BufReader::new(SnapshotReader::new(reader)?)); @@ -470,7 +490,7 @@ impl Decoder { bail!("invalid snapshot at byte 0: bad magic"); } let version = u16::from_le_bytes([header[8], header[9]]); - if version != VERSION { + if version != VERSION_1 && version != VERSION { bail!("invalid snapshot at byte 8: unsupported version {version}"); } if header[10] != PATH_ENCODING { @@ -483,12 +503,41 @@ impl Decoder { bail!("invalid snapshot at byte 11: unknown header flags"); } + let allocation = match version { + VERSION_1 => None, + VERSION => { + let node_count_offset = reader.offset; + let node_count = read_u64(&mut reader).map_err(|err| { + anyhow!("invalid snapshot integer at byte {node_count_offset}: {err}") + })?; + if node_count >= u64::from(u32::MAX) { + bail!("snapshot exceeds the tree node-index limit"); + } + let name_bytes_offset = reader.offset; + let name_bytes = read_u64(&mut reader).map_err(|err| { + anyhow!("invalid snapshot integer at byte {name_bytes_offset}: {err}") + })?; + if name_bytes > u64::from(u32::MAX) { + bail!("snapshot exceeds the tree name-storage limit"); + } + Some(( + usize::try_from(node_count) + .context("snapshot node count exceeds this address space")?, + usize::try_from(name_bytes) + .context("snapshot name storage exceeds this address space")?, + )) + } + _ => unreachable!("snapshot version was checked above"), + }; + Ok(Self { reader, + allocation, open_nodes: Vec::new(), sibling_names: Vec::new(), record: Vec::new(), node_count: 0, + name_bytes: 0, total_size: 0, total_entries: 0, summary: None, @@ -546,6 +595,15 @@ impl Decoder { if node_id >= u64::from(u32::MAX) { bail!("snapshot exceeds the tree node-index limit"); } + self.name_bytes = self + .name_bytes + .checked_add(u64::try_from(native_name.len()).context("snapshot names are too long")?) + .context("snapshot names are too long")?; + if self.allocation.is_some_and(|(node_count, name_bytes)| { + node_id > node_count as u64 || self.name_bytes > name_bytes as u64 + }) { + bail!("snapshot exceeds its declared storage requirements"); + } let sibling_ordinal = if parent_id == 0 { self.open_nodes.clear(); @@ -635,6 +693,11 @@ impl Decoder { self.node_count ); } + if self.allocation.is_some_and(|(node_count, name_bytes)| { + self.node_count != node_count as u64 || self.name_bytes != name_bytes as u64 + }) { + bail!("snapshot does not match its declared storage requirements"); + } let actual = self.reader.hash.clone().try_finalize()?; let mut expected = [0; DIGEST_LEN]; diff --git a/src/snapshot/tests.rs b/src/snapshot/tests.rs index 5cdef767..91c4dffc 100644 --- a/src/snapshot/tests.rs +++ b/src/snapshot/tests.rs @@ -10,7 +10,7 @@ fn replay_rejects_source_changes_after_validation() { let mut replay = Replay::new(Cursor::new(bytes.clone())).unwrap(); let mut changed = bytes[..bytes.len() - DIGEST_LEN].to_vec(); - changed[16] = b'b'; + changed[18] = b'b'; *replay.reader.get_mut() = rehash(changed); assert!( @@ -251,6 +251,25 @@ fn rejects_header_record_parent_and_name_errors() { ); } +#[test] +fn rejects_incorrect_v2_storage_requirements() { + let mut traversal = Traversal::new(); + let root_index = traversal.root_index; + let root = add(&mut traversal, root_index, "a", EntryData::default()); + let bytes = encoded(&traversal, &[root]); + + for offset in [HEADER_LEN, HEADER_LEN + 1] { + let mut changed = bytes[..bytes.len() - DIGEST_LEN].to_vec(); + changed[offset] += 1; + let err = read(Cursor::new(rehash(changed))).unwrap_err(); + assert!( + err.to_string() + .contains("does not match its declared storage requirements"), + "{err:#}" + ); + } +} + #[test] fn rejects_mismatched_record_length_footer_and_timestamp() { let root = record( diff --git a/src/traverse.rs b/src/traverse.rs index 19f487b2..a82c4131 100644 --- a/src/traverse.rs +++ b/src/traverse.rs @@ -267,6 +267,29 @@ impl Tree { } } + pub(crate) fn try_reserve_exact( + &mut self, + additional_nodes: usize, + additional_name_bytes: usize, + ) -> Result<(), TreeError> { + let node_count = self + .nodes + .len() + .checked_add(additional_nodes) + .ok_or(TreeError::Capacity)?; + let name_bytes = self + .names + .len() + .checked_add(additional_name_bytes) + .ok_or(TreeError::Capacity)?; + if node_count > u32::MAX as usize || name_bytes > u32::MAX as usize { + return Err(TreeError::Capacity); + } + self.nodes.try_reserve_exact(additional_nodes)?; + self.names.try_reserve_exact(additional_name_bytes)?; + Ok(()) + } + /// Add a parentless node. /// /// # Panics diff --git a/tests/snapshot.rs b/tests/snapshot.rs index b6fd89a2..34d4a084 100644 --- a/tests/snapshot.rs +++ b/tests/snapshot.rs @@ -8,13 +8,21 @@ use std::{ }; #[cfg(unix)] -const ONE_FILE: &[u8] = &[ +const ONE_FILE_V1: &[u8] = &[ 0x44, 0x55, 0x41, 0x53, 0x4e, 0x41, 0x50, 0x00, 0x01, 0x00, 0x00, 0x00, 0x07, 0x01, 0x00, 0x01, 0x61, 0x01, 0x00, 0x00, 0x00, 0x01, 0x69, 0x58, 0x7d, 0x2b, 0xa3, 0xd1, 0xef, 0x5b, 0x1d, 0x3d, 0x5b, 0xef, 0xc7, 0x03, 0x6e, 0x80, 0xfb, 0x3a, 0xf4, 0x1a, 0xea, 0xa0, 0xd0, 0x8f, 0x12, 0xcf, 0x15, 0x79, 0xc3, 0xb1, 0xa4, 0x7b, ]; +#[cfg(unix)] +const ONE_FILE_V2: &[u8] = &[ + 0x44, 0x55, 0x41, 0x53, 0x4e, 0x41, 0x50, 0x00, 0x02, 0x00, 0x00, 0x00, 0x01, 0x01, 0x07, 0x01, + 0x00, 0x01, 0x61, 0x01, 0x00, 0x00, 0x00, 0x01, 0x07, 0x2c, 0x93, 0x13, 0x2c, 0xe6, 0xc5, 0xed, + 0x37, 0x7b, 0xa8, 0xe3, 0x21, 0x9b, 0x9d, 0x50, 0x97, 0xd4, 0x4a, 0xa6, 0x96, 0x1a, 0xfd, 0x26, + 0x57, 0x54, 0xb8, 0x59, 0x5f, 0x0a, 0xad, 0x61, +]; + fn add( traversal: &mut Traversal, parent: TreeIndex, @@ -52,7 +60,7 @@ fn one_file_matches_the_golden_stream() { ); let bytes = encoded(&traversal, &[root]); - assert_eq!(bytes, ONE_FILE); + assert_eq!(bytes, ONE_FILE_V2); let snapshot = read(Cursor::new(bytes)).unwrap(); assert_eq!(snapshot.roots.len(), 1); @@ -79,7 +87,18 @@ fn one_file_matches_the_golden_stream() { Some(1) ); assert_eq!(snapshot.traversal.cost, Some(Duration::ZERO)); - assert_eq!(encoded(&snapshot.traversal, &snapshot.roots), ONE_FILE); + assert_eq!(encoded(&snapshot.traversal, &snapshot.roots), ONE_FILE_V2); +} + +#[cfg(unix)] +#[test] +fn version_1_remains_readable() { + let snapshot = read(Cursor::new(ONE_FILE_V1)).unwrap(); + assert_eq!( + snapshot.traversal.tree.name(snapshot.roots[0]).unwrap(), + Path::new("a") + ); + assert_eq!(encoded(&snapshot.traversal, &snapshot.roots), ONE_FILE_V2); } #[test] @@ -316,7 +335,7 @@ fn rejects_corruption_bad_digest_truncation_and_trailing_data() { let bytes = encoded(&traversal, &[root]); let mut corrupt = bytes.clone(); - corrupt[16] = b'b'; + corrupt[18] = b'b'; assert!( read(Cursor::new(corrupt)) .unwrap_err()