diff --git a/lore-storage/src/fragment_engine.rs b/lore-storage/src/fragment_engine.rs index 07b3c434..ef6c9de4 100644 --- a/lore-storage/src/fragment_engine.rs +++ b/lore-storage/src/fragment_engine.rs @@ -24,12 +24,62 @@ use crate::types::FragmentReference; use crate::types::Partition; use crate::write::store_fragment; -/// Splits content into fragments using `FastCDC` and stores them via -/// [`store_fragment`]. +/// Figure out where to cut `buffer` into chunks, all in one go. +async fn chunk_boundaries( + buffer: Bytes, + fixed_size_chunk: usize, +) -> Result, StorageError> { + let size = buffer.len(); + if fixed_size_chunk > 0 { + let step = fixed_size_chunk; + Ok((0..size) + .step_by(step) + .map(|offset| (offset, (offset + step).min(size))) + .collect()) + } else { + let chunker = { + // SAFETY: The chunker borrows `buffer` via a forged `'static` slice (see + // `extend_lifetime`). The compute-pool task is detached and is not + // cancelled if this future is dropped at the `rx.await` below, so + // we must keep the buffer allocation alive for the whole task. + // `Bytes::clone` bumps the refcount with a stable data pointer, + // guaranteeing the slice stays valid until the task finishes. + let slice: &[u8] = unsafe { extend_lifetime(buffer.as_ref()) }; + fastcdc::v2020::FastCDC::with_level( + slice, + FRAGMENT_SIZE_MINIMUM as u32, + FRAGMENT_SIZE_EXPECTED as u32, + FRAGMENT_SIZE_THRESHOLD as u32, + fastcdc::v2020::Normalization::Level1, + ) + }; + let buffer_guard = buffer.clone(); + let (tx, rx) = tokio::sync::oneshot::channel(); + lore_base::runtime::compute_pool().spawn(move || { + let _ = tx.send( + chunker + .map(|c| (c.offset, c.offset + c.length)) + .collect::>(), + ); + drop(buffer_guard); + }); + let boundaries = rx + .await + .map_err(|e| StorageError::internal_with_context(e, "chunker task failed"))?; + Ok(boundaries) + } +} + +/// Cuts `buffer` into chunks and stores each one via [`store_fragment`]. +/// +/// Uses `FastCDC` (content-defined chunking) when `flags.fixed_size_chunk` is 0, +/// or fixed-size chunking when it is >0. Each chunk is hashed, stored as a +/// content-addressed fragment in the immutable `store`, and assembled into a +/// fragment list (Merklized). Returns the root address and fragment. /// -/// For each chunk the function calls [`store_fragment`] directly with the -/// provided store, partition, and optional remote session. This replaces the -/// previous closure-based `store_fn` pattern. +/// When the entire buffer fits in a single chunk, the single-fragment fast path +/// is used and no fragment list is created. In `hash_only` mode, fragments are +/// not actually stored — only their hashes are computed. #[allow(clippy::too_many_arguments)] pub async fn write_fragmented( store: Arc, @@ -41,56 +91,18 @@ pub async fn write_fragmented( remote_session: Option>, tracker: Option>, ) -> Result<(Address, Fragment), StorageError> { - // Raw data, use content defined chunking with FastCDC let size = buffer.len(); - let mut chunk_offset = 0usize; let mut tasks = JoinSet::>::new(); - let mut chunk_index = 0usize; - - let chunker = Arc::new({ - // SAFETY: we await on all spawned tasks using this chunker before buffer - // is dropped, so the lifetime will not escape the scope of the buffer. - let buffer: &[u8] = unsafe { extend_lifetime(buffer.as_ref()) }; - fastcdc::v2020::FastCDC::with_level( - buffer, - FRAGMENT_SIZE_MINIMUM as u32, - FRAGMENT_SIZE_EXPECTED as u32, - FRAGMENT_SIZE_THRESHOLD as u32, - fastcdc::v2020::Normalization::Level1, - ) - }); lore_base::lore_trace!( "Write and fragment buffer to immutable store: {size} bytes representing {size} bytes (flags {flags:?})", ); - while chunk_offset < size { - let chunk_remain = size - chunk_offset; - let chunk_end = if flags.fixed_size_chunk > 0 { - chunk_offset + std::cmp::min(flags.fixed_size_chunk, chunk_remain) - } else { - // FastCDC chunking is CPU-bound (rolling hash over the buffer). - // Run it on the shared compute pool alongside compression so it - // doesn't contend with blocking IO on tokio's blocking pool. - let chunker = chunker.clone(); - // The chunker borrows `buffer` via a forged `'static` slice (see - // `extend_lifetime` above). The rayon task is detached and is not - // cancelled if this future is dropped at the `rx.await` below, so - // we must keep the buffer allocation alive for the whole task. - // `Bytes::clone` bumps the refcount with a stable data pointer, - // guaranteeing the slice stays valid until the task finishes. - let buffer_guard = buffer.clone(); - let (tx, rx) = tokio::sync::oneshot::channel(); - lore_base::runtime::compute_pool().spawn(move || { - let _ = tx.send(chunker.cut(chunk_offset, chunk_remain)); - drop(buffer_guard); - }); - let (_, chunk_end) = rx - .await - .map_err(|e| StorageError::internal_with_context(e, "chunker task failed"))?; - chunk_end - }; + + let chunk_boundaries = chunk_boundaries(buffer.clone(), flags.fixed_size_chunk).await?; + + for (chunk_index, (chunk_offset, chunk_end)) in chunk_boundaries.into_iter().enumerate() { let chunk_size = chunk_end - chunk_offset; - let chunk_buffer = buffer.slice(chunk_offset..(chunk_offset + chunk_size)); + let chunk_buffer = buffer.slice(chunk_offset..chunk_end); let fragment = Fragment { flags: flags.into(), @@ -154,12 +166,7 @@ pub async fn write_fragmented( }; Ok((chunk_index, chunk_offset, chunk_address)) }); - - chunk_offset += chunk_size; - chunk_index += 1; } - drop(chunker); - drop(buffer); let mut list_buffer = BytesMut::with_capacity(tasks.len() * std::mem::size_of::()); @@ -407,3 +414,78 @@ where { unsafe { &*(data as *const T) } } + +#[cfg(test)] +mod tests { + use bytes::Bytes; + + use super::*; + + /// Make a buffer of random bytes so chunking has to actually cut it. + fn mixed_pattern_buffer(size: usize) -> Bytes { + use rand::Rng; + let mut data = vec![0u8; size]; + rand::rng().fill(&mut data[..]); + Bytes::from(data) + } + + #[tokio::test] + async fn fastcdc_batch_handles_buffer_smaller_than_min_chunk() { + // Tiny buffer — should be one chunk covering the whole thing. + let buffer = mixed_pattern_buffer(1024); + let actual = chunk_boundaries(buffer.clone(), 0) + .await + .expect("chunking succeeds"); + assert_eq!(actual, vec![(0, buffer.len())]); + } + + #[tokio::test] + async fn fastcdc_batch_handles_empty_buffer() { + let buffer = Bytes::new(); + let actual = chunk_boundaries(buffer, 0) + .await + .expect("chunking succeeds"); + assert!(actual.is_empty()); + } + + #[tokio::test] + async fn fastcdc_batch_handles_pathological_all_zero_buffer() { + // All-zero input — the rolling hash never matches, but the split should still be clean. + let buffer = Bytes::from(vec![0u8; 512 * 1024]); + let actual = chunk_boundaries(buffer.clone(), 0) + .await + .expect("chunking succeeds"); + let reconstructed_size: usize = actual.iter().map(|(s, e)| e - s).sum(); + assert_eq!(reconstructed_size, buffer.len()); + assert!(actual.iter().all(|&(s, e)| s < e)); + assert_eq!(actual.first().copied(), Some((0, actual[0].1))); + assert_eq!(actual.last().unwrap().1, buffer.len()); + } + + #[tokio::test] + async fn fixed_size_chunking_covers_whole_buffer() { + let buffer = mixed_pattern_buffer(10 * 1024 + 17); + let step = 1024; + let actual = chunk_boundaries(buffer.clone(), step) + .await + .expect("chunking succeeds"); + + let expected: Vec<(usize, usize)> = (0..buffer.len()) + .step_by(step) + .map(|o| (o, (o + step).min(buffer.len()))) + .collect(); + assert_eq!(actual, expected); + + let reconstructed_size: usize = actual.iter().map(|(s, e)| e - s).sum(); + assert_eq!(reconstructed_size, buffer.len()); + } + + #[tokio::test] + async fn fixed_size_chunking_handles_buffer_smaller_than_step() { + let buffer = Bytes::from(vec![1u8; 100]); + let actual = chunk_boundaries(buffer.clone(), 1024) + .await + .expect("chunking succeeds"); + assert_eq!(actual, vec![(0, buffer.len())]); + } +}