diff --git a/api_changes/update_260424_next_stable_chunk_boundary.md b/api_changes/update_260424_next_stable_chunk_boundary.md index 6f9aba3d8..696f67ce2 100644 --- a/api_changes/update_260424_next_stable_chunk_boundary.md +++ b/api_changes/update_260424_next_stable_chunk_boundary.md @@ -1,8 +1,9 @@ This update adds a new public deduplication helper for computing restart-safe chunk boundaries from existing chunk boundary metadata. What changed -- Added `xet_data::deduplication::next_stable_chunk_boundary(starting_position, chunk_boundaries) -> Option`. -- Re-exported it from `xet_data::deduplication` so downstream crates can use it directly. +- Added `next_stable_chunk_boundary(starting_position, chunk_boundaries) -> Option`. +- Canonical implementation lives in `xet_core_structures::xorb_object::constants` (alongside the chunk-size constants it uses). +- Re-exported from `xet_data::deduplication` for convenience. - The function scans forward from `starting_position` and returns the next chunk boundary that satisfies the stable-boundary condition: - two consecutive chunk sizes in `[2 * min_chunk, max_chunk - min_chunk)`, - where `min_chunk` and `max_chunk` are derived from chunking constants. @@ -10,6 +11,7 @@ What changed Why this matters - Callers that already have chunk-boundary metadata can locate a stable resume boundary without re-reading file bytes. - This enables deterministic alignment behavior for resumed/partial workflows that need chunk boundaries robust to prefix changes. +- The server-side `build_file_chunk_hashes_response` now extends dirty ranges to stable chunk boundaries before building windows, so that the client's rechunking around each dirty window is guaranteed to converge to the same chunk boundaries as the original file. Usage notes - `chunk_boundaries` should be monotonically increasing chunk-end offsets produced by the same chunking configuration. diff --git a/wasm/hf_xet_wasm/src/wasm_file_cleaner.rs b/wasm/hf_xet_wasm/src/wasm_file_cleaner.rs index 9aa478c5d..eea2dd405 100644 --- a/wasm/hf_xet_wasm/src/wasm_file_cleaner.rs +++ b/wasm/hf_xet_wasm/src/wasm_file_cleaner.rs @@ -184,7 +184,8 @@ impl SingleFileCleaner { let metadata_ext = FileMetadataExt::new(sha256); // Now finish the deduplication process. - let (file_hash, remaining_file_data, deduplication_metrics) = self.dedup_manager.finalize(Some(metadata_ext)); + let (file_hash, _chunk_hashes, remaining_file_data, deduplication_metrics) = + self.dedup_manager.finalize(Some(metadata_ext)); // Let's check some things that should be invariants { diff --git a/xet_client/src/cas_client/chunk_window_builder.rs b/xet_client/src/cas_client/chunk_window_builder.rs new file mode 100644 index 000000000..85d165e19 --- /dev/null +++ b/xet_client/src/cas_client/chunk_window_builder.rs @@ -0,0 +1,393 @@ +//! Server-side state machine for `GET /v2/file-chunk-hashes/{file_id}` (mirrored from +//! xetcas PR #987) plus a small driver helper used by simulation clients to produce a +//! [`FileChunkHashesResponse`] without routing through HTTP. + +use xet_core_structures::merklehash::{MerkleHash, MerkleHashSubtree}; +use xet_core_structures::metadata_shard::file_structs::MDBFileInfo; +use xet_core_structures::xorb_object::constants::next_stable_chunk_boundary; + +use crate::cas_types::{ChunkWindow, FileChunkHashesResponse, FileRange}; +use crate::error::{ClientError, Result}; + +pub struct ChunkWindowBuilder<'a> { + dirty_ranges: &'a [FileRange], + dirty_idx: usize, + in_dirty_zone: bool, + gap_is_first: bool, + /// End-byte cursor: the start of the next chunk fed to the builder. + cursor: u64, + gap_chunks: Vec<(MerkleHash, u64)>, + windows: Vec, + hash_ranges: Vec>, +} + +fn next_stable_end_for_range(range_end: u64, file_size: u64, chunk_boundaries: &[usize]) -> u64 { + if range_end >= file_size { + return file_size; + } + let Ok(starting_position) = usize::try_from(range_end) else { + return file_size; + }; + next_stable_chunk_boundary(starting_position, chunk_boundaries) + .and_then(|boundary| u64::try_from(boundary).ok()) + .map(|boundary| boundary.min(file_size)) + .unwrap_or(file_size) +} + +fn extend_dirty_ranges_to_stable_windows( + mut dirty_ranges: Vec, + file_size: u64, + chunk_boundaries: &[usize], +) -> Vec { + for range in &mut dirty_ranges { + if range.end < file_size { + range.end = next_stable_end_for_range(range.end, file_size, chunk_boundaries); + } + } + + dirty_ranges.sort_by_key(|range| range.start); + let mut coalesced: Vec = Vec::with_capacity(dirty_ranges.len()); + for range in dirty_ranges { + if let Some(last) = coalesced.last_mut() + && range.start <= last.end + { + last.end = last.end.max(range.end); + continue; + } + coalesced.push(range); + } + coalesced +} + +impl<'a> ChunkWindowBuilder<'a> { + pub fn new(dirty_ranges: &'a [FileRange]) -> Self { + debug_assert!( + dirty_ranges.windows(2).all(|w| w[0].end <= w[1].start), + "dirty_ranges must be sorted and non-overlapping" + ); + Self { + dirty_ranges, + dirty_idx: 0, + in_dirty_zone: false, + gap_is_first: true, + cursor: 0, + gap_chunks: Vec::new(), + windows: Vec::with_capacity(dirty_ranges.len()), + hash_ranges: Vec::new(), + } + } + + pub fn process_chunk(&mut self, hash: MerkleHash, size: u64, byte_end: u64) { + let byte_start = self.cursor; + let overlaps_dirty = self.overlaps_current_dirty(byte_start, byte_end); + + if !self.in_dirty_zone { + if overlaps_dirty { + self.open_window(byte_end); + self.in_dirty_zone = true; + } else { + self.gap_chunks.push((hash, size)); + } + } else if overlaps_dirty { + self.windows + .last_mut() + .expect("in_dirty_zone implies a window has been opened") + .end = byte_end; + self.merge_ahead(byte_end); + } else { + self.dirty_idx += 1; + self.gap_is_first = false; + + // The first clean chunk after a dirty zone may itself overlap the next + // dirty range (back-to-back dirty ranges on adjacent chunks). + let overlaps_next = self.overlaps_current_dirty(byte_start, byte_end); + if overlaps_next { + self.open_window(byte_end); + } else { + self.in_dirty_zone = false; + self.gap_chunks.push((hash, size)); + } + } + self.cursor = byte_end; + } + + /// Returns true when the entry ending at `byte_end` (and starting at the cursor) + /// is fully contained within the current dirty range. + pub fn entry_fully_dirty(&self, byte_end: u64) -> bool { + self.dirty_idx < self.dirty_ranges.len() + && self.dirty_ranges[self.dirty_idx].start <= self.cursor + && byte_end <= self.dirty_ranges[self.dirty_idx].end + } + + /// Process a fully-dirty shard entry without iterating its individual chunks. + pub fn skip_dirty_entry(&mut self, byte_end: u64) { + if !self.in_dirty_zone { + self.open_window(byte_end); + self.in_dirty_zone = true; + } else { + self.windows + .last_mut() + .expect("in_dirty_zone implies a window has been opened") + .end = byte_end; + self.merge_ahead(byte_end); + } + self.cursor = byte_end; + } + + /// Consume the builder and return the dirty windows + N+1 gap hash ranges. + pub fn finish(mut self) -> (Vec, Vec>) { + let trailing = MerkleHashSubtree::from_chunks(self.gap_is_first, &self.gap_chunks, true); + self.hash_ranges.push(Self::to_option(trailing)); + (self.windows, self.hash_ranges) + } + + fn open_window(&mut self, byte_end: u64) { + let gap = MerkleHashSubtree::from_chunks(self.gap_is_first, &self.gap_chunks, false); + self.hash_ranges.push(Self::to_option(gap)); + self.gap_chunks.clear(); + self.windows.push(FileRange::new(self.cursor, byte_end)); + self.merge_ahead(byte_end); + } + + fn overlaps_current_dirty(&self, byte_start: u64, byte_end: u64) -> bool { + self.dirty_idx < self.dirty_ranges.len() + && byte_end > self.dirty_ranges[self.dirty_idx].start + && byte_start < self.dirty_ranges[self.dirty_idx].end + } + + fn merge_ahead(&mut self, byte_end: u64) { + while self.dirty_idx + 1 < self.dirty_ranges.len() && byte_end > self.dirty_ranges[self.dirty_idx + 1].start { + self.dirty_idx += 1; + } + } + + fn to_option(hr: MerkleHashSubtree) -> Option { + if hr.is_empty() { None } else { Some(hr) } + } +} + +/// Drive [`ChunkWindowBuilder`] over a flat list of `(chunk_hash, size)` pairs and +/// assemble the [`FileChunkHashesResponse`]. Simulation clients pre-collect the chunks for +/// the file (typically by walking segments + xorb metadata) and call this to answer +/// `Client::get_file_chunk_hashes` locally. +/// +/// Also emits `gap_verification`: for each **stable** original segment (one that lies +/// entirely outside the dirty windows), the corresponding `FileVerificationEntry` from +/// `file_info.verification` is copied into `gap_verification` in segment order. The +/// composed shard built by `upload_ranges` uses these to reconstruct its own verification +/// section without needing per-chunk hashes for the stable segments. When `file_info` +/// has no verification entries (legacy / test files), `gap_verification` is empty. +pub fn build_file_chunk_hashes_response( + file_info: &MDBFileInfo, + dirty_ranges: Vec, + chunks: impl IntoIterator, +) -> Result { + let file_size = file_info.file_size(); + let dirty_ranges: Vec = dirty_ranges + .into_iter() + .map(|r| FileRange::new(r.start, r.end.min(file_size))) + .filter(|r| r.start < r.end && r.start < file_size) + .collect(); + if dirty_ranges.is_empty() { + return Err(ClientError::Other("no valid dirty ranges".into())); + } + + let chunks: Vec<(MerkleHash, u64)> = chunks.into_iter().collect(); + let mut chunk_boundaries: Vec = Vec::with_capacity(chunks.len()); + let mut stable_ranges_supported = true; + { + let mut boundary_cursor: u64 = 0; + for (_, size) in &chunks { + boundary_cursor += *size; + let Ok(boundary) = usize::try_from(boundary_cursor) else { + stable_ranges_supported = false; + break; + }; + chunk_boundaries.push(boundary); + } + } + + let dirty_ranges = if stable_ranges_supported { + extend_dirty_ranges_to_stable_windows(dirty_ranges, file_size, &chunk_boundaries) + } else { + dirty_ranges + }; + + let total_chunks = chunks.len() as u64; + let mut builder = ChunkWindowBuilder::new(&dirty_ranges); + let mut cumulative_bytes: u64 = 0; + for (hash, size) in chunks { + cumulative_bytes += size; + builder.process_chunk(hash, size, cumulative_bytes); + } + + let (windows, hash_ranges) = builder.finish(); + if windows.is_empty() { + return Err(ClientError::Other("dirty ranges do not overlap any chunks".into())); + } + + // Emit one range hash per stable segment (= no overlap with any window). + // Segments and windows are both monotonic, so a two-pointer walk is O(S+W). + // + // Contract: `verification` is either empty (legacy / test files without verification + // entries) or 1:1 with `segments`. A partially-populated mismatch is a real bug we + // want loud here, rather than as a confusing "ran out of gap_verification entries" + // error later in `compose_mdb`. + let gap_verification = if file_info.verification.is_empty() { + Vec::new() + } else if file_info.verification.len() == file_info.segments.len() { + let mut gv = Vec::new(); + let mut acc = 0u64; + let mut wi = 0usize; + for (idx, seg) in file_info.segments.iter().enumerate() { + let seg_start = acc; + let seg_end = acc + seg.unpacked_segment_bytes as u64; + acc = seg_end; + while wi < windows.len() && windows[wi].end <= seg_start { + wi += 1; + } + let overlaps = wi < windows.len() && windows[wi].start < seg_end; + if !overlaps { + gv.push(crate::cas_types::HexMerkleHash::from(file_info.verification[idx].range_hash)); + } + } + gv + } else { + return Err(ClientError::Other(format!( + "file_info has {} verification entries but {} segments; \ + expected either zero or a 1:1 mapping", + file_info.verification.len(), + file_info.segments.len() + ))); + }; + + Ok(FileChunkHashesResponse { + total_chunks, + file_size, + windows: windows + .into_iter() + .map(|r| ChunkWindow { + dirty_byte_range: [r.start, r.end], + }) + .collect(), + hash_ranges, + gap_verification, + }) +} + +#[cfg(test)] +mod tests { + use xet_core_structures::merklehash::MerkleHash; + use xet_core_structures::metadata_shard::file_structs::{FileDataSequenceEntry, MDBFileInfo}; + use xet_core_structures::xorb_object::constants::{ + MAXIMUM_CHUNK_MULTIPLIER, MINIMUM_CHUNK_DIVISOR, TARGET_CHUNK_SIZE, + }; + + use super::*; + + fn stable_chunk_size() -> u64 { + let minimum_chunk = *TARGET_CHUNK_SIZE / *MINIMUM_CHUNK_DIVISOR; + let maximum_chunk = *TARGET_CHUNK_SIZE * *MAXIMUM_CHUNK_MULTIPLIER; + let size = 2 * minimum_chunk; + assert!(size < maximum_chunk - minimum_chunk); + size as u64 + } + + fn build_test_file_info_and_chunks(n_chunks: usize, chunk_size: u64) -> (MDBFileInfo, Vec<(MerkleHash, u64)>) { + let chunks: Vec<(MerkleHash, u64)> = (0..n_chunks) + .map(|i| (MerkleHash::random_from_seed(i as u64 + 1), chunk_size)) + .collect(); + let file_size = chunk_size * n_chunks as u64; + let file_info = MDBFileInfo { + segments: vec![FileDataSequenceEntry { + xorb_hash: MerkleHash::random_from_seed(999), + xorb_flags: 0, + unpacked_segment_bytes: file_size as u32, + chunk_index_start: 0, + chunk_index_end: n_chunks as u32, + }], + ..Default::default() + }; + (file_info, chunks) + } + + #[test] + fn test_server_extends_dirty_window_to_next_stable_boundary() { + let chunk_size = stable_chunk_size(); + let (file_info, chunks) = build_test_file_info_and_chunks(6, chunk_size); + let dirty_ranges = vec![FileRange::new(1, chunk_size + 1)]; + + let response = build_file_chunk_hashes_response(&file_info, dirty_ranges, chunks).unwrap(); + + assert_eq!(response.windows.len(), 1); + assert_eq!(response.windows[0].dirty_byte_range, [0, 4 * chunk_size]); + } + + #[test] + fn test_server_coalesces_ranges_after_stable_extension() { + let chunk_size = stable_chunk_size(); + let (file_info, chunks) = build_test_file_info_and_chunks(8, chunk_size); + let dirty_ranges = vec![ + FileRange::new(1, chunk_size + 1), + FileRange::new(3 * chunk_size + 7, 3 * chunk_size + 42), + ]; + + let response = build_file_chunk_hashes_response(&file_info, dirty_ranges, chunks).unwrap(); + + assert_eq!(response.windows.len(), 1); + assert_eq!(response.windows[0].dirty_byte_range, [0, 6 * chunk_size]); + } + + #[test] + fn test_server_no_extension_when_range_already_at_file_end() { + let chunk_size = stable_chunk_size(); + let (file_info, chunks) = build_test_file_info_and_chunks(6, chunk_size); + let file_size = chunk_size * 6; + let dirty_ranges = vec![FileRange::new(4 * chunk_size, file_size)]; + + let response = build_file_chunk_hashes_response(&file_info, dirty_ranges, chunks).unwrap(); + + assert_eq!(response.windows.len(), 1); + assert_eq!(response.windows[0].dirty_byte_range, [4 * chunk_size, file_size]); + } + + #[test] + fn test_server_separate_ranges_stay_separate_when_far_apart() { + let chunk_size = stable_chunk_size(); + let n_chunks = 20; + let (file_info, chunks) = build_test_file_info_and_chunks(n_chunks, chunk_size); + let dirty_ranges = vec![ + FileRange::new(0, chunk_size), + FileRange::new(14 * chunk_size, 15 * chunk_size), + ]; + + let response = build_file_chunk_hashes_response(&file_info, dirty_ranges, chunks).unwrap(); + + assert!( + response.windows.len() >= 2, + "far-apart ranges should remain separate, got {} window(s)", + response.windows.len() + ); + assert_eq!(response.hash_ranges.len(), response.windows.len() + 1); + } + + #[test] + fn test_server_extension_falls_through_to_file_end_when_no_stable_boundary() { + let minimum_chunk = *TARGET_CHUNK_SIZE / *MINIMUM_CHUNK_DIVISOR; + let maximum_chunk = *TARGET_CHUNK_SIZE * *MAXIMUM_CHUNK_MULTIPLIER; + let forced_size = maximum_chunk as u64; + let (file_info, chunks) = build_test_file_info_and_chunks(4, forced_size); + let file_size = forced_size * 4; + let dirty_ranges = vec![FileRange::new(0, forced_size)]; + + let response = build_file_chunk_hashes_response(&file_info, dirty_ranges, chunks).unwrap(); + + assert_eq!(response.windows.len(), 1); + let w = &response.windows[0]; + assert_eq!( + w.dirty_byte_range[1], file_size, + "when no stable boundary exists, window should extend to file end; \ + minimum_chunk={minimum_chunk}, maximum_chunk={maximum_chunk}" + ); + } +} diff --git a/xet_client/src/cas_client/interface.rs b/xet_client/src/cas_client/interface.rs index 8664336f7..2e36168a5 100644 --- a/xet_client/src/cas_client/interface.rs +++ b/xet_client/src/cas_client/interface.rs @@ -5,7 +5,9 @@ use xet_core_structures::xorb_object::SerializedXorbObject; use super::adaptive_concurrency::ConnectionPermit; use super::progress_tracked_streams::ProgressCallback; -use crate::cas_types::{BatchQueryReconstructionResponse, FileRange, HttpRange, QueryReconstructionResponseV2}; +use crate::cas_types::{ + BatchQueryReconstructionResponse, FileChunkHashesResponse, FileRange, HttpRange, QueryReconstructionResponseV2, +}; use crate::error::Result; #[async_trait::async_trait] @@ -70,4 +72,17 @@ pub trait Client: Send + Sync { progress_callback: Option, upload_permit: ConnectionPermit, ) -> Result; + + /// Compute chunk-aligned dirty windows + opaque gap [`MerkleHashSubtree`] summaries for the + /// given file, narrowed to `dirty_ranges`. + /// + /// `dirty_ranges` must be sorted and non-overlapping. Per-chunk hashes are never returned; + /// the response carries only `windows.len()` dirty windows and `windows.len() + 1` gap + /// subtrees, which the client merges with locally-recomputed window subtrees to obtain the + /// new file hash. + async fn get_file_chunk_hashes( + &self, + file_id: &MerkleHash, + dirty_ranges: Vec, + ) -> Result; } diff --git a/xet_client/src/cas_client/mod.rs b/xet_client/src/cas_client/mod.rs index 7b4b73ffa..f1820dc85 100644 --- a/xet_client/src/cas_client/mod.rs +++ b/xet_client/src/cas_client/mod.rs @@ -14,6 +14,7 @@ pub use crate::common::http_client::{Api, ResponseErrorLogger, build_auth_http_c pub mod adaptive_concurrency; pub mod auth; +pub mod chunk_window_builder; pub mod exports; mod interface; pub mod multipart; diff --git a/xet_client/src/cas_client/remote_client.rs b/xet_client/src/cas_client/remote_client.rs index 349abbd92..3a8bc6177 100644 --- a/xet_client/src/cas_client/remote_client.rs +++ b/xet_client/src/cas_client/remote_client.rs @@ -23,8 +23,9 @@ use super::progress_tracked_streams::{ use super::retry_wrapper::{RetryWrapper, RetryableReqwestError}; use super::{Client, INFORMATION_LOG_LEVEL}; use crate::cas_types::{ - BatchQueryReconstructionResponse, FileRange, HttpRange, Key, QueryReconstructionResponse, + BatchQueryReconstructionResponse, FileChunkHashesResponse, FileRange, HttpRange, Key, QueryReconstructionResponse, QueryReconstructionResponseV2, UploadShardResponse, UploadShardResponseType, UploadXorbResponse, + X_RANGE_DIRTY_HEADER, }; use crate::common::http_client::{self, Api}; use crate::error::{ClientError, Result}; @@ -748,6 +749,48 @@ impl Client for RemoteClient { Ok(n_upload_bytes) } + + #[instrument(skip_all, name = "RemoteClient::get_file_chunk_hashes", fields(file.hash = file_id.hex(), n_ranges = dirty_ranges.len()))] + async fn get_file_chunk_hashes( + &self, + file_id: &MerkleHash, + dirty_ranges: Vec, + ) -> Result { + if dirty_ranges.is_empty() { + return Err(ClientError::Other("get_file_chunk_hashes requires at least one dirty range".into())); + } + + let url = Url::parse(&format!("{}/v2/file-chunk-hashes/{}", self.endpoint, file_id.hex()))?; + + // Multi-range `bytes=A-B,C-D` value. `HttpRange` is inclusive-end and `Display`s as + // `start-end`; conversion from `FileRange` does the +1/-1 for us. + let header_value = HeaderValue::from_str(&format!( + "bytes={}", + dirty_ranges + .iter() + .copied() + .map(HttpRange::from) + .map(|r| r.to_string()) + .collect::>() + .join(",") + )) + .map_err(|err| ClientError::Other(format!("invalid X-Range-Dirty header value: {err}")))?; + + let api_tag = "cas::get_file_chunk_hashes"; + let client = self.authenticated_http_client.clone(); + + let response: FileChunkHashesResponse = RetryWrapper::new(self.ctx.clone(), api_tag) + .run_and_extract_json(move || { + client + .get(url.clone()) + .header(X_RANGE_DIRTY_HEADER, header_value.clone()) + .with_extension(Api(api_tag)) + .send() + }) + .await?; + + Ok(response) + } } #[cfg(test)] diff --git a/xet_client/src/cas_client/simulation/local_client.rs b/xet_client/src/cas_client/simulation/local_client.rs index 036a4588f..cf4dd0725 100644 --- a/xet_client/src/cas_client/simulation/local_client.rs +++ b/xet_client/src/cas_client/simulation/local_client.rs @@ -17,6 +17,7 @@ use tokio::time::{Duration, Instant}; use tracing::{error, info, warn}; use xet_core_structures::merklehash::{MerkleHash, compute_data_hash}; use xet_core_structures::metadata_shard::file_structs::{FileDataSequenceHeader, MDBFileInfo, MDBFileInfoView}; +use xet_core_structures::metadata_shard::shard_file_reconstructor::FileReconstructor; use xet_core_structures::metadata_shard::shard_format::MDB_FILE_INFO_ENTRY_SIZE; use xet_core_structures::metadata_shard::shard_in_memory::MDBInMemoryShard; use xet_core_structures::metadata_shard::streaming_shard::MDBMinimalShard; @@ -35,10 +36,12 @@ use super::direct_access_client::DirectAccessClient; use super::xorb_utils::{self, REFERENCE_INSTANT, duration_to_expiration_secs_ceil}; use crate::cas_client::Client; use crate::cas_client::adaptive_concurrency::AdaptiveConcurrencyController; +use crate::cas_client::chunk_window_builder::build_file_chunk_hashes_response; use crate::cas_client::progress_tracked_streams::ProgressCallback; use crate::cas_types::{ - BatchQueryReconstructionResponse, FileRange, HexMerkleHash, HttpRange, QueryReconstructionResponse, - QueryReconstructionResponseV2, XorbMultiRangeFetch, XorbRangeDescriptor, XorbReconstructionFetchInfo, + BatchQueryReconstructionResponse, FileChunkHashesResponse, FileRange, HexMerkleHash, HttpRange, + QueryReconstructionResponse, QueryReconstructionResponseV2, XorbMultiRangeFetch, XorbRangeDescriptor, + XorbReconstructionFetchInfo, }; use crate::error::{ClientError, Result}; @@ -1692,6 +1695,30 @@ impl Client for LocalClient { // Should not reach here, but return error if we do. Err(ClientError::PresignedUrlExpirationError) } + + async fn get_file_chunk_hashes( + &self, + file_id: &MerkleHash, + dirty_ranges: Vec, + ) -> Result { + self.apply_api_delay().await; + + let Some((file_info, _)) = self.shard_manager.get_file_reconstruction_info(file_id).await? else { + return Err(ClientError::FileNotFound(*file_id)); + }; + + let mut chunks: Vec<(MerkleHash, u64)> = Vec::new(); + for segment in &file_info.segments { + let xorb_obj = self.xorb_footer(&segment.xorb_hash).await?; + chunks.extend( + xorb_obj + .chunk_hash_sizes(segment.chunk_index_start, segment.chunk_index_end) + .map_err(|err| ClientError::Other(format!("chunk_hash_sizes error: {err}")))?, + ); + } + + build_file_chunk_hashes_response(&file_info, dirty_ranges, chunks) + } } fn map_redb_db_error(e: impl std::fmt::Debug) -> ClientError { diff --git a/xet_client/src/cas_client/simulation/local_server/server.rs b/xet_client/src/cas_client/simulation/local_server/server.rs index 21d746a9b..c1f1fc01d 100644 --- a/xet_client/src/cas_client/simulation/local_server/server.rs +++ b/xet_client/src/cas_client/simulation/local_server/server.rs @@ -493,6 +493,14 @@ impl Client for LocalTestServer { .upload_xorb(prefix, serialized_xorb_object, progress_callback, upload_permit) .await } + + async fn get_file_chunk_hashes( + &self, + file_id: &xet_core_structures::merklehash::MerkleHash, + dirty_ranges: Vec, + ) -> Result { + self.client.get_file_chunk_hashes(file_id, dirty_ranges).await + } } #[cfg(test)] diff --git a/xet_client/src/cas_client/simulation/local_server/simulation_control_client.rs b/xet_client/src/cas_client/simulation/local_server/simulation_control_client.rs index 9c7a05fa4..fc3197617 100644 --- a/xet_client/src/cas_client/simulation/local_server/simulation_control_client.rs +++ b/xet_client/src/cas_client/simulation/local_server/simulation_control_client.rs @@ -19,7 +19,9 @@ use crate::cas_client::interface::Client; use crate::cas_client::simulation::deletion_controls::ObjectTag; use crate::cas_client::simulation::xorb_utils::duration_to_expiration_secs_ceil; use crate::cas_client::simulation::{DeletionControlableClient, DirectAccessClient}; -use crate::cas_types::{FileRange, HexMerkleHash, QueryReconstructionResponseV2, XorbReconstructionFetchInfo}; +use crate::cas_types::{ + FileChunkHashesResponse, FileRange, HexMerkleHash, QueryReconstructionResponseV2, XorbReconstructionFetchInfo, +}; use crate::error::{ClientError, Result}; const CONFIG_POST_MAX_ATTEMPTS: usize = 4; @@ -243,6 +245,14 @@ impl Client for SimulationControlClient { .upload_xorb(prefix, serialized_xorb_object, progress_callback, upload_permit) .await } + + async fn get_file_chunk_hashes( + &self, + file_id: &MerkleHash, + dirty_ranges: Vec, + ) -> Result { + self.remote_client.get_file_chunk_hashes(file_id, dirty_ranges).await + } } #[async_trait] diff --git a/xet_client/src/cas_client/simulation/memory_client.rs b/xet_client/src/cas_client/simulation/memory_client.rs index 669896ef4..4b1237b07 100644 --- a/xet_client/src/cas_client/simulation/memory_client.rs +++ b/xet_client/src/cas_client/simulation/memory_client.rs @@ -30,9 +30,11 @@ use super::deletion_controls::ObjectTag; use super::direct_access_client::DirectAccessClient; use super::random_xorb::RandomXorb; use super::xorb_utils::{self, REFERENCE_INSTANT, duration_to_expiration_secs_ceil}; +use crate::cas_client::chunk_window_builder::build_file_chunk_hashes_response; use crate::cas_types::{ - BatchQueryReconstructionResponse, FileRange, HexMerkleHash, HttpRange, QueryReconstructionResponse, - QueryReconstructionResponseV2, XorbMultiRangeFetch, XorbRangeDescriptor, XorbReconstructionFetchInfo, + BatchQueryReconstructionResponse, FileChunkHashesResponse, FileRange, HexMerkleHash, HttpRange, + QueryReconstructionResponse, QueryReconstructionResponseV2, XorbMultiRangeFetch, XorbRangeDescriptor, + XorbReconstructionFetchInfo, }; use crate::error::{ClientError, Result}; @@ -950,6 +952,40 @@ impl Client for MemoryClient { } Ok((Bytes::from(all_decompressed), all_chunk_indices)) } + + async fn get_file_chunk_hashes( + &self, + file_id: &MerkleHash, + dirty_ranges: Vec, + ) -> Result { + self.apply_api_delay().await; + + let file_info = { + let shard = self.shard.read().await; + shard + .get_file_reconstruction_info(file_id) + .ok_or(ClientError::FileNotFound(*file_id))? + }; + + let xorbs = self.xorbs.read().await; + let mut chunks: Vec<(MerkleHash, u64)> = Vec::new(); + for segment in &file_info.segments { + let storage = xorbs + .get(&segment.xorb_hash) + .ok_or(ClientError::XORBNotFound(segment.xorb_hash))?; + let xorb_obj = match storage { + XorbStorage::Materialized { entry, .. } => std::borrow::Cow::Borrowed(&entry.xorb_object), + XorbStorage::Random { xorb, .. } => std::borrow::Cow::Owned(xorb.get_xorb_object()), + }; + chunks.extend( + xorb_obj + .chunk_hash_sizes(segment.chunk_index_start, segment.chunk_index_end) + .map_err(|err| ClientError::Other(format!("chunk_hash_sizes error: {err}")))?, + ); + } + + build_file_chunk_hashes_response(&file_info, dirty_ranges, chunks) + } } #[cfg(not(target_family = "wasm"))] diff --git a/xet_client/src/cas_client/simulation/simulation_client.rs b/xet_client/src/cas_client/simulation/simulation_client.rs index 49b878839..4deab2b0f 100644 --- a/xet_client/src/cas_client/simulation/simulation_client.rs +++ b/xet_client/src/cas_client/simulation/simulation_client.rs @@ -234,4 +234,12 @@ impl Client for RemoteSimulationClient { .upload_xorb(prefix, serialized_xorb_object, progress_callback, upload_permit) .await } + + async fn get_file_chunk_hashes( + &self, + file_id: &xet_core_structures::merklehash::MerkleHash, + dirty_ranges: Vec, + ) -> Result { + self.inner.get_file_chunk_hashes(file_id, dirty_ranges).await + } } diff --git a/xet_client/src/cas_client/simulation/simulation_server.rs b/xet_client/src/cas_client/simulation/simulation_server.rs index 61e2ec46f..a0d4fa642 100644 --- a/xet_client/src/cas_client/simulation/simulation_server.rs +++ b/xet_client/src/cas_client/simulation/simulation_server.rs @@ -541,6 +541,14 @@ impl Client for LocalTestServer { .upload_xorb(prefix, serialized_xorb_object, progress_callback, upload_permit) .await } + + async fn get_file_chunk_hashes( + &self, + file_id: &xet_core_structures::merklehash::MerkleHash, + dirty_ranges: Vec, + ) -> Result { + self.client.get_file_chunk_hashes(file_id, dirty_ranges).await + } } #[async_trait] diff --git a/xet_client/src/cas_types/mod.rs b/xet_client/src/cas_types/mod.rs index db05f6af1..fdbcfa983 100644 --- a/xet_client/src/cas_types/mod.rs +++ b/xet_client/src/cas_types/mod.rs @@ -7,7 +7,7 @@ use std::str::FromStr; use serde::{Deserialize, Serialize}; use serde_repr::{Deserialize_repr, Serialize_repr}; use thiserror::Error; -use xet_core_structures::merklehash::MerkleHash; +use xet_core_structures::merklehash::{MerkleHash, MerkleHashSubtree}; mod key; pub use key::*; @@ -311,6 +311,42 @@ pub struct QueryChunkResponse { pub shard: MerkleHash, } +/// HTTP header carrying the dirty byte ranges to feed to `GET /v2/file-chunk-hashes/{file_id}`. +/// +/// Distinct from the standard `Range` header (which scopes the response body): this header tags +/// regions that the client intends to re-chunk, and the response covers the whole file (windows + +/// gap subtrees). Value uses the same `bytes=A-B,C-D` syntax as `Range`. +pub const X_RANGE_DIRTY_HEADER: &str = "X-Range-Dirty"; + +/// One chunk-aligned dirty window of a file, returned by `GET /v2/file-chunk-hashes/{file_id}`. +/// +/// `dirty_byte_range` is `[start, end)` and is expanded outward to the chunk boundaries that +/// fully contain the requested dirty range, so the client must re-chunk the entire span. +#[derive(Debug, Serialize, Deserialize, Clone)] +#[serde(rename_all = "camelCase")] +pub struct ChunkWindow { + pub dirty_byte_range: [u64; 2], +} + +/// Response shape for `GET /v2/file-chunk-hashes/{file_id}`. +/// +/// Contains `windows.len()` dirty windows interleaved with `windows.len() + 1` opaque +/// `MerkleHashSubtree` summaries for the surrounding gaps. To reconstruct the new file hash, +/// merge `[hash_ranges[0], window0_subtree, hash_ranges[1], window1_subtree, ..., hash_ranges[N]]` +/// using `MerkleHashSubtree::merge`. Per-chunk hashes are never transferred. +#[derive(Debug, Serialize, Deserialize, Clone)] +#[serde(rename_all = "camelCase")] +pub struct FileChunkHashesResponse { + pub total_chunks: u64, + pub file_size: u64, + pub windows: Vec, + pub hash_ranges: Vec>, + /// One range hash per **stable original segment** (= a segment that lies in a gap + /// between dirty windows or before/after them, in segment order). Wraps each into a + /// `FileVerificationEntry` to populate the composed shard's verification section. + pub gap_verification: Vec, +} + #[cfg(test)] mod tests { use super::*; diff --git a/xet_core_structures/src/merklehash/mod.rs b/xet_core_structures/src/merklehash/mod.rs index 71aba67d3..569d116c2 100644 --- a/xet_core_structures/src/merklehash/mod.rs +++ b/xet_core_structures/src/merklehash/mod.rs @@ -48,6 +48,9 @@ pub mod data_hash; pub use data_hash::*; pub type MerkleHash = DataHash; +/// List of (chunk_hash, chunk_uncompressed_size) pairs for a file or xorb range. +pub type ChunkHashList = Vec<(MerkleHash, u64)>; + mod aggregated_hashes; pub mod merkle_hash_subtree; pub mod passthrough_hasher; diff --git a/xet_core_structures/src/metadata_shard/shard_file_manager.rs b/xet_core_structures/src/metadata_shard/shard_file_manager.rs index e77d05dd4..3a181fca1 100644 --- a/xet_core_structures/src/metadata_shard/shard_file_manager.rs +++ b/xet_core_structures/src/metadata_shard/shard_file_manager.rs @@ -256,12 +256,7 @@ impl ShardFileManager { s.verify_shard_integrity_debug_only(); // Make sure the shard is in the shard directory - debug_assert!( - s.path.starts_with(&self.shard_directory), - "{:?} not in {:?}", - &s.path, - &self.shard_directory - ); + debug_assert!(s.path.starts_with(&self.shard_directory), "{:?} not in {:?}", s.path, self.shard_directory); if self .shard_bookkeeper diff --git a/xet_core_structures/src/metadata_shard/shard_in_memory.rs b/xet_core_structures/src/metadata_shard/shard_in_memory.rs index acff71938..1ce9deffa 100644 --- a/xet_core_structures/src/metadata_shard/shard_in_memory.rs +++ b/xet_core_structures/src/metadata_shard/shard_in_memory.rs @@ -90,14 +90,14 @@ impl MDBInMemoryShard { pub fn recalculate_shard_size(&mut self) { // Calculate the size let mut num_bytes = 0u64; - for (_, xorb_block_contents) in self.xorb_content.iter() { + for xorb_block_contents in self.xorb_content.values() { num_bytes += xorb_block_contents.num_bytes(); // The xorb lookup table num_bytes += (size_of::() + size_of::()) as u64; } - for (_, file_info) in self.file_content.iter() { + for file_info in self.file_content.values() { num_bytes += file_info.num_bytes(); num_bytes += (size_of::() + size_of::()) as u64; } diff --git a/xet_core_structures/src/xorb_object/byte_grouping/compression_stats/collect_compression_stats.rs b/xet_core_structures/src/xorb_object/byte_grouping/compression_stats/collect_compression_stats.rs index 5dbc81e26..07099fa4c 100644 --- a/xet_core_structures/src/xorb_object/byte_grouping/compression_stats/collect_compression_stats.rs +++ b/xet_core_structures/src/xorb_object/byte_grouping/compression_stats/collect_compression_stats.rs @@ -187,8 +187,8 @@ async fn main() -> Result<(), Box> { let writer = Arc::new(Mutex::new(wtr)); - eprintln!("Output File: {}", &args.output); - eprintln!("Input Files: {:?}", &args.files); + eprintln!("Output File: {}", args.output); + eprintln!("Input Files: {:?}", args.files); let max_limiter = Arc::new(Semaphore::new(32 * 1024)); diff --git a/xet_core_structures/src/xorb_object/constants.rs b/xet_core_structures/src/xorb_object/constants.rs index b322326e8..2105f781e 100644 --- a/xet_core_structures/src/xorb_object/constants.rs +++ b/xet_core_structures/src/xorb_object/constants.rs @@ -27,3 +27,50 @@ lazy_static::lazy_static! { /// The maximum chunk size, calculated from the configurable constants above pub static ref MAX_CHUNK_SIZE: usize = (*TARGET_CHUNK_SIZE) * (*MAXIMUM_CHUNK_MULTIPLIER); } + +/// Given a list of chunk boundaries in a file and an arbitrary reference position, +/// returns the next stable chunk boundary at or after that position. +/// +/// `starting_position` may be any byte offset in the file; it does not need to +/// be an existing chunk boundary. The search starts at the first chunk boundary +/// `>= starting_position`. +/// +/// A stable chunk boundary is defined such that any possible changes in the data +/// before `starting_position` would produce the same chunk boundaries at the +/// stable boundary and later. The fixed data between `starting_position` and +/// the returned stable boundary is always sufficient to restore the chunker to +/// its original chunk boundaries. +/// +/// The stability condition requires two consecutive chunks after `starting_position`, +/// both with sizes in `[2 * min_chunk, max_chunk - min_chunk)`. The boundary +/// at the end of the second such chunk is the stable chunk boundary. +/// +/// The lower bound is `2 * min_chunk` rather than `min_chunk` (as used in +/// `find_partitions` in the chunking module) because this function operates on +/// existing chunk boundaries without data access, and cannot verify the absence +/// of hidden hash triggers in the `[c_k, c_k + min_chunk)` skip zone. A +/// shadow-zone trigger can advance a modified chunker by up to `min_chunk`, so +/// the next chunk must be at least `2 * min_chunk` to remain reachable. +/// +/// See `parallel chunking.lyx` for the full proof and `find_stable_start` in +/// `merkle_hash_subtree.rs` for the analogous construction in merkle hashing. +pub fn next_stable_chunk_boundary(starting_position: usize, chunk_boundaries: &[usize]) -> Option { + let minimum_chunk = *TARGET_CHUNK_SIZE / *MINIMUM_CHUNK_DIVISOR; + let maximum_chunk = *TARGET_CHUNK_SIZE * *MAXIMUM_CHUNK_MULTIPLIER; + + let start_idx = chunk_boundaries.partition_point(|&x| x < starting_position); + + for i in start_idx..chunk_boundaries.len().saturating_sub(2) { + let size_a = chunk_boundaries[i + 1] - chunk_boundaries[i]; + let size_b = chunk_boundaries[i + 2] - chunk_boundaries[i + 1]; + + if size_a >= 2 * minimum_chunk + && size_a < maximum_chunk - minimum_chunk + && size_b >= 2 * minimum_chunk + && size_b < maximum_chunk - minimum_chunk + { + return Some(chunk_boundaries[i + 2]); + } + } + None +} diff --git a/xet_core_structures/src/xorb_object/xorb_object_format.rs b/xet_core_structures/src/xorb_object/xorb_object_format.rs index 0284dea50..db22b2a43 100644 --- a/xet_core_structures/src/xorb_object/xorb_object_format.rs +++ b/xet_core_structures/src/xorb_object/xorb_object_format.rs @@ -13,15 +13,15 @@ use super::constants::{TARGET_CHUNK_SIZE, XORB_BLOCK_SIZE}; use super::xorb_chunk_format::{deserialize_chunk, deserialize_chunk_header, serialize_chunk, write_chunk_header}; use super::{CompressionScheme, RawXorbData, XorbChunkHeader}; use crate::error::{CoreError, Validate}; -use crate::merklehash::{DataHash, MerkleHash}; +use crate::merklehash::{ChunkHashList, DataHash, MerkleHash}; use crate::metadata_shard::chunk_verification::range_hash_from_chunks; use crate::serialization_utils::*; pub type XorbObjectIdent = [u8; 7]; -pub(crate) const XORB_OBJECT_FORMAT_IDENT: XorbObjectIdent = [b'X', b'E', b'T', b'B', b'L', b'O', b'B']; +pub(crate) const XORB_OBJECT_FORMAT_IDENT: XorbObjectIdent = *b"XETBLOB"; pub(crate) const XORB_OBJECT_FORMAT_VERSION_V0: u8 = 0; -pub(crate) const XORB_OBJECT_FORMAT_IDENT_HASHES: XorbObjectIdent = [b'X', b'B', b'L', b'B', b'H', b'S', b'H']; -pub(crate) const XORB_OBJECT_FORMAT_IDENT_BOUNDARIES: XorbObjectIdent = [b'X', b'B', b'L', b'B', b'B', b'N', b'D']; +pub(crate) const XORB_OBJECT_FORMAT_IDENT_HASHES: XorbObjectIdent = *b"XBLBHSH"; +pub(crate) const XORB_OBJECT_FORMAT_IDENT_BOUNDARIES: XorbObjectIdent = *b"XBLBBND"; pub(crate) const XORB_OBJECT_FORMAT_VERSION: u8 = 1; pub(crate) const XORB_OBJECT_FORMAT_HASHES_VERSION: u8 = 0; @@ -1240,6 +1240,21 @@ impl XorbObject { Ok(incl_end - before_start) } + /// Returns (chunk_hash, uncompressed_size) pairs for chunks in [start, end). + pub fn chunk_hash_sizes(&self, start: u32, end: u32) -> Result { + self.validate_xorb_object_info()?; + if end > self.info.num_chunks || start > end { + return Err(CoreError::InvalidArguments); + } + (start..end) + .map(|i| { + let hash = self.info.chunk_hashes[i as usize]; + let size = self.uncompressed_chunk_length(i)? as u64; + Ok((hash, size)) + }) + .collect() + } + /// Helper method to verify that info object is complete fn validate_xorb_object_info(&self) -> Result<(), CoreError> { if self.info.num_chunks == 0 { diff --git a/xet_data/src/deduplication/chunking.rs b/xet_data/src/deduplication/chunking.rs index 82976e7d7..97f70026c 100644 --- a/xet_data/src/deduplication/chunking.rs +++ b/xet_data/src/deduplication/chunking.rs @@ -359,52 +359,10 @@ pub fn find_partitions( Ok(partitions) } -/// Given a list of chunk boundaries in a file and an arbitrary reference position, -/// returns the next stable chunk boundary at or after that position. -/// -/// `starting_position` may be any byte offset in the file; it does not need to -/// be an existing chunk boundary. The search starts at the first chunk boundary -/// `>= starting_position`. -/// -/// A stable chunk boundary is defined such that any possible changes in the data -/// before `starting_position` would produce the same chunk boundaries at the -/// stable boundary and later. The fixed data between `starting_position` and -/// the returned stable boundary is always sufficient to restore the chunker to -/// its original chunk boundaries. -/// -/// The stability condition requires two consecutive chunks after `starting_position`, -/// both with sizes in `[2 * min_chunk, max_chunk - min_chunk)`. The boundary -/// at the end of the second such chunk is the stable chunk boundary. -/// -/// The lower bound is `2 * min_chunk` rather than `min_chunk` (as used in -/// [`find_partitions`]) because this function operates on existing chunk -/// boundaries without data access, and cannot verify the absence of hidden -/// hash triggers in the `[c_k, c_k + min_chunk)` skip zone. A shadow-zone -/// trigger can advance a modified chunker by up to `min_chunk`, so the next -/// chunk must be at least `2 * min_chunk` to remain reachable. -/// -/// See `parallel chunking.lyx` for the full proof and `find_stable_start` in -/// `merkle_hash_subtree.rs` for the analogous construction in merkle hashing. -pub fn next_stable_chunk_boundary(starting_position: usize, chunk_boundaries: &[usize]) -> Option { - let minimum_chunk = *TARGET_CHUNK_SIZE / *MINIMUM_CHUNK_DIVISOR; - let maximum_chunk = *TARGET_CHUNK_SIZE * *MAXIMUM_CHUNK_MULTIPLIER; - - let start_idx = chunk_boundaries.partition_point(|&x| x < starting_position); - - for i in start_idx..chunk_boundaries.len().saturating_sub(2) { - let size_a = chunk_boundaries[i + 1] - chunk_boundaries[i]; - let size_b = chunk_boundaries[i + 2] - chunk_boundaries[i + 1]; - - if size_a >= 2 * minimum_chunk - && size_a < maximum_chunk - minimum_chunk - && size_b >= 2 * minimum_chunk - && size_b < maximum_chunk - minimum_chunk - { - return Some(chunk_boundaries[i + 2]); - } - } - None -} +// Re-exported from xet_core_structures where the canonical implementation lives, +// so that downstream users of xet_data::deduplication::next_stable_chunk_boundary +// continue to work without a source change. +pub use xet_core_structures::xorb_object::constants::next_stable_chunk_boundary; #[cfg(test)] mod tests { diff --git a/xet_data/src/deduplication/file_deduplication.rs b/xet_data/src/deduplication/file_deduplication.rs index 1e9ede240..3235fa7c0 100644 --- a/xet_data/src/deduplication/file_deduplication.rs +++ b/xet_data/src/deduplication/file_deduplication.rs @@ -2,7 +2,7 @@ use std::result::Result; use more_asserts::{debug_assert_le, debug_assert_lt}; use xet_core_structures::MerkleHashMap; -use xet_core_structures::merklehash::{MerkleHash, file_hash}; +use xet_core_structures::merklehash::{ChunkHashList, MerkleHash, file_hash}; use xet_core_structures::metadata_shard::file_structs::{ FileDataSequenceEntry, FileDataSequenceHeader, FileMetadataExt, FileVerificationEntry, MDBFileInfo, }; @@ -36,7 +36,7 @@ pub struct FileDeduper { new_data_hash_lookup: MerkleHashMap, /// The current chunk hashes for this file. - chunk_hashes: Vec<(MerkleHash, u64)>, + chunk_hashes: ChunkHashList, /// The current file data entries. file_info: Vec, @@ -399,8 +399,11 @@ impl FileDeduper) -> (MerkleHash, DataAggregator, DeduplicationMetrics) { + /// Returns (file hash, chunk_hashes, data aggregation, deduplication metrics) + pub fn finalize( + self, + metadata_ext: Option, + ) -> (MerkleHash, ChunkHashList, DataAggregator, DeduplicationMetrics) { let file_hash = file_hash(&self.chunk_hashes); let metadata = FileDataSequenceHeader::new(file_hash, self.file_info.len(), true, metadata_ext.is_some()); @@ -434,6 +437,6 @@ impl FileDeduper Resu debug_assert_eq!(size_read, size); - let (file_info, _) = handle.finish().await?; + let (file_info, _metrics) = handle.finish().await?; translator.finalize().await?; diff --git a/xet_data/src/processing/data_client.rs b/xet_data/src/processing/data_client.rs index c268537af..34813fa12 100644 --- a/xet_data/src/processing/data_client.rs +++ b/xet_data/src/processing/data_client.rs @@ -47,7 +47,8 @@ pub async fn clean_bytes( ) -> Result<(XetFileInfo, DeduplicationMetrics)> { let (_id, mut handle) = processor.start_clean(None, Some(bytes.len() as u64), sha256_policy)?; handle.add_data(&bytes).await?; - handle.finish().await + let (info, metrics) = handle.finish().await?; + Ok((info, metrics)) } #[instrument(skip_all, name = "clean_file", fields(file.name = tracing::field::Empty, file.len = tracing::field::Empty))] @@ -76,7 +77,8 @@ pub async fn clean_file( handle.add_data(&buffer[0..bytes]).await?; } - handle.finish().await + let (info, metrics) = handle.finish().await?; + Ok((info, metrics)) } /// Computes the xet hash for a single file without uploading. diff --git a/xet_data/src/processing/file_cleaner.rs b/xet_data/src/processing/file_cleaner.rs index 8ae975b27..d292ea3e9 100644 --- a/xet_data/src/processing/file_cleaner.rs +++ b/xet_data/src/processing/file_cleaner.rs @@ -5,8 +5,9 @@ use std::sync::Arc; use bytes::Bytes; use chrono::{DateTime, Utc}; use tracing::{Instrument, debug_span, info, instrument}; +use xet_core_structures::merklehash::ChunkHashList; use xet_core_structures::metadata_shard::Sha256; -use xet_core_structures::metadata_shard::file_structs::FileMetadataExt; +use xet_core_structures::metadata_shard::file_structs::{FileMetadataExt, MDBFileInfo}; use xet_runtime::core::XetContext; use super::XetFileInfo; @@ -201,15 +202,39 @@ impl SingleFileCleaner { } /// Return the representation of the file after clean as a pointer file instance. - #[instrument(skip_all, name = "FileCleaner::finish", fields(file_name=self.file_name.as_ref().map(|s|s.to_string())))] - pub async fn finish(mut self) -> Result<(XetFileInfo, DeduplicationMetrics)> { - // Chunk the rest of the data. + pub async fn finish(self) -> Result<(XetFileInfo, DeduplicationMetrics)> { + let (info, _chunks, metrics) = self.finish_with_chunks().await?; + Ok((info, metrics)) + } + + /// Same as [`finish`], but also returns the per-chunk hash list produced during CDC. + /// Only needed by composition flows (e.g. `upload_ranges`) that build partial + /// `MerkleHashSubtree` nodes for newly-uploaded windows; regular uploads should call + /// [`finish`] instead. + #[instrument(skip_all, name = "FileCleaner::finish_with_chunks", fields(file_name=self.file_name.as_ref().map(|s|s.to_string())))] + pub async fn finish_with_chunks(self) -> Result<(XetFileInfo, ChunkHashList, DeduplicationMetrics)> { + let (file_info, chunk_hashes, _, deduplication_metrics) = Self::finish_inner(self, true).await?; + Ok((file_info, chunk_hashes, deduplication_metrics)) + } + + /// Like `finish_with_chunks`, but does NOT register the file's MDBFileInfo in the + /// session shard. Returns the MDBFileInfo directly so the caller can compose it into a + /// larger file without creating orphan shard entries. + pub async fn finish_with_chunks_detached( + self, + ) -> Result<(XetFileInfo, ChunkHashList, MDBFileInfo, DeduplicationMetrics)> { + Self::finish_inner(self, false).await + } + + async fn finish_inner( + mut self, + register: bool, + ) -> Result<(XetFileInfo, ChunkHashList, MDBFileInfo, DeduplicationMetrics)> { if let Some(chunk) = self.chunker.finish() { let data = Arc::new([chunk]); self.deduper_process_chunks(data).await?; } - // Resolve the SHA-256: computed, provided, or skipped. let sha256 = if let Some(generator) = self.sha_generator.take() { Some(generator.finalize().await?) } else { @@ -217,7 +242,7 @@ impl SingleFileCleaner { }; let metadata_ext = sha256.map(FileMetadataExt::new); - let (file_hash, remaining_file_data, deduplication_metrics) = + let (file_hash, chunk_hashes, remaining_file_data, deduplication_metrics) = self.dedup_manager_fut.await?.finalize(metadata_ext); let file_info = XetFileInfo { @@ -226,34 +251,33 @@ impl SingleFileCleaner { sha256: sha256.map(|s| s.hex()), }; - // Let's check some things that should be invariants #[cfg(debug_assertions)] { - // There should be exactly one file referenced in the remaining file data. debug_assert_eq!(remaining_file_data.pending_file_info.len(), 1); - - // The size should be total bytes debug_assert_eq!(remaining_file_data.pending_file_info[0].0.file_size(), deduplication_metrics.total_bytes) } - // Now, return all this information to the - self.session - .register_single_file_clean_completion(remaining_file_data, &deduplication_metrics) - .await?; - - // NB: xorb upload is happening in the background, this number is optimistic since it does - // not count transfer time of the uploaded xorbs, which is why `end_processing_ts` + let mdb_file_info = if register { + self.session + .register_single_file_clean_completion(remaining_file_data, &deduplication_metrics) + .await?; + MDBFileInfo::default() + } else { + self.session + .register_single_file_clean_completion_detached(remaining_file_data, &deduplication_metrics) + .await? + }; info!( target: "client_telemetry", action = "clean", - file_name = self.file_name.unwrap_or_default().to_string(), + file_name = self.file_name.as_deref().unwrap_or_default().to_string(), file_size_count = deduplication_metrics.total_bytes, new_bytes_count = deduplication_metrics.new_bytes, start_ts = self.start_time.to_rfc3339(), end_processing_ts = Utc::now().to_rfc3339(), ); - Ok((file_info, deduplication_metrics)) + Ok((file_info, chunk_hashes, mdb_file_info, deduplication_metrics)) } } diff --git a/xet_data/src/processing/file_upload_session.rs b/xet_data/src/processing/file_upload_session.rs index 3b00104cb..b80f93f10 100644 --- a/xet_data/src/processing/file_upload_session.rs +++ b/xet_data/src/processing/file_upload_session.rs @@ -467,36 +467,74 @@ impl FileUploadSession { Ok(()) } + /// Like `register_single_file_clean_completion`, but does NOT register the MDBFileInfo + /// in the session shard. Returns the finalized MDBFileInfo instead. + /// Used by composition flows where only the final composed file should appear in the shard. + pub(crate) async fn register_single_file_clean_completion_detached( + self: &Arc, + file_data: DataAggregator, + dedup_metrics: &DeduplicationMetrics, + ) -> Result { + // Always cut a dedicated xorb for detached files. This avoids mixing with other + // files in current_session_data whose MDBFileInfo must still be registered. + let file_infos = self.process_aggregated_data_as_xorb_detached(file_data).await?; + + self.deduplication_metrics.lock().await.merge_in(dedup_metrics); + + debug_assert_eq!(file_infos.len(), 1); + file_infos + .into_iter() + .next() + .ok_or_else(|| DataError::InternalError("detached completion produced no file info".into())) + } + /// Process the aggregated data, uploading the data as a xorb and registering the files async fn process_aggregated_data_as_xorb(self: &Arc, data_agg: DataAggregator) -> Result<()> { + self.process_aggregated_data_as_xorb_impl(data_agg, true).await.map(|_| ()) + } + + /// Upload the xorb data but do NOT register file reconstruction info in the shard. + /// Returns the finalized MDBFileInfo for each file in the aggregator. + async fn process_aggregated_data_as_xorb_detached( + self: &Arc, + data_agg: DataAggregator, + ) -> Result> { + self.process_aggregated_data_as_xorb_impl(data_agg, false).await + } + + async fn process_aggregated_data_as_xorb_impl( + self: &Arc, + data_agg: DataAggregator, + register_files: bool, + ) -> Result> { let (xorb, new_files) = data_agg.finalize(); let xorb_hash = xorb.hash(); debug_assert_le!(xorb.num_bytes(), *MAX_XORB_BYTES); debug_assert_le!(xorb.data.len(), *MAX_XORB_CHUNKS); - // Now, we need to scan all the file dependencies for dependencies on this xorb, as - // these would not have been registered yet as we just got the xorb hash. let mut new_dependencies = Vec::with_capacity(new_files.len()); + let mut file_infos = Vec::with_capacity(new_files.len()); - { - for (file_id, fi, bytes_in_xorb) in new_files { - new_dependencies.push(FileXorbDependency { - file_id, - xorb_hash, - n_bytes: bytes_in_xorb, - is_external: false, - }); - - // Record the reconstruction. + for (file_id, fi, bytes_in_xorb) in new_files { + new_dependencies.push(FileXorbDependency { + file_id, + xorb_hash, + n_bytes: bytes_in_xorb, + is_external: false, + }); + + if register_files { self.shard_interface.add_file_reconstruction_info(fi).await?; + } else { + file_infos.push(fi); } } // Register the xorb and start the upload process. self.register_new_xorb(xorb, &new_dependencies).await?; - Ok(()) + Ok(file_infos) } /// Register a xorb dependencies that is given as part of the dedup process. @@ -570,6 +608,13 @@ impl FileUploadSession { Ok(()) } + /// Register a pre-composed file reconstruction plan (MDBFileInfo) with this session. + /// Used for append-aware writes where the caller builds the reconstruction plan + /// from existing segments + newly uploaded segments. + pub(crate) async fn register_composed_file(self: &Arc, file_info: MDBFileInfo) -> Result<()> { + self.shard_interface.add_file_reconstruction_info(file_info).await + } + fn check_not_finalized(&self) -> Result<()> { if self.finalized.load(Ordering::Acquire) { return Err(DataError::InvalidOperation("FileUploadSession already finalized".to_string())); @@ -730,7 +775,7 @@ mod tests { .start_clean(Some("test".into()), Some(data.len() as u64), Sha256Policy::Skip) .unwrap(); cleaner.add_data(data).await.unwrap(); - cleaner.finish().await.unwrap(); + let _ = cleaner.finish().await.unwrap(); // Verify that the shard has no metadata_ext (no SHA-256). let (_metrics, file_infos) = upload_session.finalize_with_file_info().await.unwrap(); diff --git a/xet_data/src/processing/mod.rs b/xet_data/src/processing/mod.rs index adcf7ee0e..a0d9bee4e 100644 --- a/xet_data/src/processing/mod.rs +++ b/xet_data/src/processing/mod.rs @@ -5,6 +5,7 @@ mod file_cleaner; mod file_download_session; mod file_upload_session; pub mod migration_tool; +pub mod range_upload; mod remote_client_interface; mod sha256; mod shard_interface; @@ -14,9 +15,11 @@ mod xet_file; pub use file_cleaner::{Sha256Policy, SingleFileCleaner}; pub use file_download_session::FileDownloadSession; pub use file_upload_session::FileUploadSession; +pub use range_upload::{DirtyInput, upload_ranges}; pub use remote_client_interface::create_remote_client; pub use xet_client::cas_client::Client as CasClient; pub use xet_client::chunk_cache::{CacheConfig, ChunkCache, get_cache}; +pub use xet_core_structures::merklehash::ChunkHashList; pub use xet_file::XetFileInfo; pub use crate::deduplication::RawXorbData; diff --git a/xet_data/src/processing/range_upload.rs b/xet_data/src/processing/range_upload.rs new file mode 100644 index 000000000..6c00faac7 --- /dev/null +++ b/xet_data/src/processing/range_upload.rs @@ -0,0 +1,2105 @@ +use std::ops::Range; +use std::pin::Pin; +use std::sync::Arc; + +use tokio::io::{AsyncRead, AsyncReadExt}; +use tracing::{debug, info}; +use xet_client::cas_client::Client; +use xet_client::cas_types::{FileChunkHashesResponse, FileRange, HexMerkleHash}; +use xet_core_structures::merklehash::{ChunkHashList, MerkleHash, MerkleHashSubtree}; +use xet_core_structures::metadata_shard::file_structs::{ + FileDataSequenceEntry, FileDataSequenceHeader, FileVerificationEntry, MDBFileInfo, +}; +use xet_runtime::core::XetContext; + +use super::XetFileInfo; +use super::configurations::TranslatorConfig; +use super::file_cleaner::Sha256Policy; +use super::file_upload_session::FileUploadSession; +use crate::error::{DataError, Result}; +use crate::file_reconstruction::FileReconstructor; + +/// A single edit applied to the original file: replace `original_range` with `new_length` +/// bytes from `reader`. +/// +/// All three combinations are supported: +/// - `original_range.len() == new_length` → in-place edit (no file size change). +/// - `original_range.len() != new_length` → resize edit (file grows or shrinks). +/// - `original_range.start == original_range.end` → pure insert at that position. +/// - `new_length == 0` → pure delete of `original_range`. +/// +/// Pure append at end of file is `{ original_range: original_size..original_size, new_length: N, reader: ... }`. +/// Pure truncate-to-N is `{ original_range: N..original_size, new_length: 0, reader: empty }`. +pub struct DirtyInput { + pub original_range: Range, + pub reader: Pin>, + pub new_length: u64, +} + +/// Size of blocks read from the dirty source and fed to the cleaner. +const STREAM_BLOCK_SIZE: usize = 4 * 1024 * 1024; // 4 MB + +/// A dirty window the server told us to re-upload, augmented with the bytes the cleaner +/// produced and the resulting per-window file info. `start`/`end` are file byte offsets +/// extended for append/truncation past `original_size` if needed. +struct UploadedWindow { + start: u64, + end: u64, + chunks: ChunkHashList, + mdb: MDBFileInfo, +} + +/// Upload an edited version of an existing file, reusing the unchanged regions from the +/// original file's CAS segments and only re-uploading the parts the caller actually +/// rewrites. +/// +/// `dirty_inputs` is a list of edits expressed in the **original file's coordinates**. +/// Each edit replaces `original_range` with `new_length` bytes from its reader; resize +/// edits (including pure inserts and pure deletes) are supported. The output file size is +/// derived from the inputs. +/// +/// # When to use +/// +/// - **In-place edit**: `original_range.len() == new_length`, file size unchanged. +/// - **Resize edit**: any `new_length`. Replaces `original_range.len()` original bytes with `new_length` new bytes. +/// - **Pure insert**: `original_range.start == original_range.end`, `new_length > 0`. +/// - **Pure delete**: `original_range.start < original_range.end`, `new_length == 0`. +/// - **Append**: `original_range == original_size..original_size`, `new_length > 0`. +/// - **Truncate to N**: `original_range == N..original_size`, `new_length == 0`. +/// - **No change**: empty `dirty_inputs`. Returns the original hash without any CAS call. +/// +/// # Arguments +/// +/// * `config` - Translator configuration for creating upload sessions. +/// * `cas_client` - CAS client for fetching original file metadata and downloading boundary bytes. +/// * `original_hash` - Merkle hash of the original file in CAS. +/// * `original_size` - Size of the original file in bytes. +/// * `dirty_inputs` - Edits to apply. Must be sorted by `original_range.start` and non-overlapping +/// (`prev.original_range.end <= next.original_range.start`). Each reader must yield exactly `new_length` bytes. Each +/// reader is consumed exactly once. +/// +/// # Limitations +/// +/// The composed file has no SHA-256 metadata (`metadata_ext = None`), since recomputing +/// it would require reading the full file. This means `upload_ranges` is only suitable +/// for contexts that don't require SHA-256 verification (e.g. HF buckets, xet-native +/// repos), not for Git LFS-backed repos that verify SHA-256 on download. +pub async fn upload_ranges( + config: Arc, + cas_client: Arc, + original_hash: MerkleHash, + original_size: u64, + mut dirty_inputs: Vec, +) -> Result { + validate_dirty_ranges(&dirty_inputs, original_size)?; + let total_size = compute_total_size(original_size, &dirty_inputs)?; + + if dirty_inputs.is_empty() { + debug_assert_eq!(total_size, original_size); + return Ok(XetFileInfo::new(original_hash.hex(), original_size)); + } + + // Empty original: nothing to compose against — upload as a fresh file (concatenation of + // the edits' new bytes, since every `original_range` must be `0..0`). + if original_size == 0 { + return upload_fresh_file(config, dirty_inputs, total_size).await; + } + + let recon_result = cas_client.get_file_reconstruction_info(&original_hash).await?; + let original_mdb = recon_result + .map(|(mdb, _)| mdb) + .ok_or_else(|| DataError::ParameterError(format!("file {} not found in CAS", original_hash.hex())))?; + if original_mdb.file_size() != original_size { + return Err(DataError::ParameterError(format!( + "caller said original_size={original_size} but reconstruction info reports {}", + original_mdb.file_size() + ))); + } + + // `seg_byte_starts[i]` is the first byte of segment `i`; the trailing entry equals `original_size`. + let mut seg_byte_starts: Vec = Vec::with_capacity(original_mdb.segments.len() + 1); + seg_byte_starts.push(0); + let mut acc = 0u64; + for s in &original_mdb.segments { + acc += s.unpacked_segment_bytes as u64; + seg_byte_starts.push(acc); + } + + // Snapping to *segment* boundaries (rather than chunk boundaries) lets us swap whole + // segments during composition, so the client never has to truncate a segment mid-chunk. + // Safe to send to the server because segment edges are chunk edges, so the server's + // chunk-aligned windows come back identical to our snapped ranges. + // + // Pure inserts (`original_range.start == original_range.end`) snap to the segment that + // owns the insert position so the cleaner has enough surrounding bytes to re-chunk + // around the insertion. An insert at `original_size` snaps to the last segment. + let n_segs = original_mdb.segments.len(); + let mut snapped: Vec<(u64, u64)> = Vec::with_capacity(dirty_inputs.len()); + for input in &dirty_inputs { + let r = &input.original_range; + let (s, e) = if r.start == r.end { + // Pure insert: pick the segment containing `r.start`. At end-of-file, fall back to + // the last segment. + if r.start == original_size { + (seg_byte_starts[n_segs - 1], seg_byte_starts[n_segs]) + } else { + (snap_to_segment_start(&seg_byte_starts, r.start), snap_to_segment_end(&seg_byte_starts, r.start + 1)) + } + } else { + (snap_to_segment_start(&seg_byte_starts, r.start), snap_to_segment_end(&seg_byte_starts, r.end)) + }; + snapped.push((s, e)); + } + + snapped.sort_by_key(|&(s, _)| s); + let mut coalesced: Vec<(u64, u64)> = Vec::with_capacity(snapped.len()); + for r in snapped { + if let Some(last) = coalesced.last_mut() + && r.0 <= last.1 + { + last.1 = last.1.max(r.1); + continue; + } + coalesced.push(r); + } + + let server_query: Vec = coalesced.iter().map(|&(s, e)| FileRange::new(s, e)).collect(); + if server_query.is_empty() { + return Err(DataError::InternalError("internal: non-empty dirty_inputs produced no server query".into())); + } + + let response: FileChunkHashesResponse = cas_client.get_file_chunk_hashes(&original_hash, server_query).await?; + // The server may coalesce adjacent/overlapping dirty ranges after extending them to + // stable boundaries, so only enforce shape invariants on the returned payload itself. + if response.windows.is_empty() { + return Err(DataError::InternalError("server returned no windows".into())); + } + if response.hash_ranges.len() != response.windows.len() + 1 { + return Err(DataError::InternalError(format!( + "server returned {} hash_ranges, expected {} (n_windows + 1)", + response.hash_ranges.len(), + response.windows.len() + 1 + ))); + } + let gap_verification = response.gap_verification; + + let ctx = config.ctx.clone(); + let session = FileUploadSession::new(config.clone()).await?; + let mut input_idx = 0usize; + let mut uploaded: Vec = Vec::with_capacity(response.windows.len()); + + let mut buf = vec![0u8; STREAM_BLOCK_SIZE]; + for window in response.windows.iter() { + let w_start = window.dirty_byte_range[0]; + let w_end = window.dirty_byte_range[1]; + + // Find the slice of edits that land in this window. A pure insert at exactly + // `w_end` belongs here only when `w_end == original_size` (no later window can + // take it); anywhere else it belongs to the next window starting at that byte. + // Exercised by `test_resize_insert_at_segment_boundary` and `test_mid_edit_plus_append`. + let edits_end = dirty_inputs[input_idx..] + .iter() + .take_while(|d| { + let r = &d.original_range; + if r.start == r.end { + r.start < w_end || (r.start == w_end && w_end == original_size) + } else { + r.end <= w_end + } + }) + .count() + + input_idx; + let window_edits = &mut dirty_inputs[input_idx..edits_end]; + + let (removed, added): (u64, u64) = window_edits + .iter() + .map(|d| (d.original_range.end - d.original_range.start, d.new_length)) + .fold((0, 0), |(rm, ad), (r, a)| (rm + r, ad + a)); + let middle_size = (w_end - w_start) + added - removed; + + let (_id, mut cleaner) = session.start_clean(None, Some(middle_size), Sha256Policy::Skip)?; + + let mut cursor = w_start; + for input in window_edits.iter_mut() { + let edit_start = input.original_range.start; + let edit_end = input.original_range.end; + debug_assert!(edit_start >= w_start && edit_end <= w_end, "edit straddles window (validation bug)"); + + if cursor < edit_start { + stream_cas_range(&ctx, &cas_client, original_hash, cursor, edit_start, &mut cleaner).await?; + } + + let mut remaining = input.new_length as usize; + while remaining > 0 { + let to_read = buf.len().min(remaining); + input.reader.read_exact(&mut buf[..to_read]).await.map_err(|err| { + DataError::InternalError(format!( + "failed to read dirty input [{}, {}): {err}", + input.original_range.start, input.original_range.end + )) + })?; + cleaner.add_data(&buf[..to_read]).await?; + remaining -= to_read; + } + + cursor = edit_end; + } + input_idx = edits_end; + + if cursor < w_end { + stream_cas_range(&ctx, &cas_client, original_hash, cursor, w_end, &mut cleaner).await?; + } + + let (_info, chunks, mdb, _metrics) = cleaner.finish_with_chunks_detached().await?; + uploaded.push(UploadedWindow { + start: w_start, + end: w_end, + chunks, + mdb, + }); + } + + // Every edit must have been assigned to exactly one window. If the server + // returned narrower `dirty_byte_range`s than requested, leftover edits would + // silently drop and produce a corrupt file with no error. + if input_idx != dirty_inputs.len() { + return Err(DataError::InternalError(format!( + "{} dirty edits not assigned to any window (input_idx={input_idx}, total={})", + dirty_inputs.len() - input_idx, + dirty_inputs.len() + ))); + } + + // Merge sequence: [gap0, w0, gap1, w1, ..., gapN]. Empty gaps (`None`) are skipped. + let mut hash_ranges = response.hash_ranges; + let trailing_gap = hash_ranges.pop().flatten(); + // Leading gap is empty (None) => first window starts at byte 0. + let first_window_at_start = matches!(hash_ranges.first(), Some(None)); + let last_window_at_end = trailing_gap.is_none(); + let last_idx = uploaded.len() - 1; + + let mut merge_seq: Vec = Vec::with_capacity(2 * uploaded.len() + 1); + for (i, (w, gap)) in uploaded.iter().zip(hash_ranges).enumerate() { + if let Some(g) = gap { + merge_seq.push(g); + } + let at_start = i == 0 && first_window_at_start; + let at_end = i == last_idx && last_window_at_end; + merge_seq.push(MerkleHashSubtree::from_chunks(at_start, &w.chunks, at_end)); + } + if let Some(g) = trailing_gap { + merge_seq.push(g); + } + + let merged = MerkleHashSubtree::merge(&merge_seq) + .map_err(|err| DataError::InternalError(format!("MerkleHashSubtree::merge failed: {err}")))?; + // `final_hash()` is the aggregated chunk hash; the file hash is its HMAC with the zero + // salt (matching `file_hash` / the cleaner's output for files without SHA-256 metadata, + // the only flavor `upload_ranges` produces). Empty content is the one exception: + // `file_hash([])` short-circuits to `MerkleHash::default()` *without* HMAC, so we mirror + // that here when the result is zero-length (e.g. truncating to empty). + let aggregated_hash = merged.final_hash().ok_or_else(|| { + DataError::InternalError("merged subtree is not fully closed; cannot derive final hash".into()) + })?; + let combined_hash = if total_size == 0 { + MerkleHash::default() + } else { + aggregated_hash.hmac(MerkleHash::default()) + }; + + let composed_mdb = + compose_mdb(&original_mdb, &seg_byte_starts, &uploaded, gap_verification, combined_hash, original_size)?; + + debug!( + "upload_ranges: composed hash={}, {} segments, {} windows", + combined_hash.hex(), + composed_mdb.segments.len(), + uploaded.len() + ); + + session.register_composed_file(composed_mdb).await?; + session.finalize().await?; + + let total_dirty: u64 = dirty_inputs.iter().map(|d| d.new_length).sum(); + info!( + "upload_ranges: hash={} size={} (original={}, {} windows, {} dirty bytes)", + combined_hash.hex(), + total_size, + original_size, + uploaded.len(), + total_dirty + ); + + Ok(XetFileInfo::new(combined_hash.hex(), total_size)) +} + +/// Assemble the composed MDB by splicing uploaded window segments into the original +/// file's segment list, pulling verification entries from the server (for stable gaps) +/// and from the cleaner (for re-uploaded windows). +fn compose_mdb( + original_mdb: &MDBFileInfo, + seg_byte_starts: &[u64], + uploaded: &[UploadedWindow], + gap_verification: Vec, + combined_hash: MerkleHash, + original_size: u64, +) -> Result { + let mut all_segments: Vec = Vec::new(); + let mut all_verification: Vec = Vec::new(); + let mut seg_idx = 0usize; + let n_segs = original_mdb.segments.len(); + let mut gap_idx = 0usize; + + for w in uploaded { + while seg_idx < n_segs && seg_byte_starts[seg_idx] < w.start { + if seg_byte_starts[seg_idx + 1] > w.start { + return Err(DataError::InternalError(format!( + "server returned a window starting at {} that straddles segment {} \ + ({}..{}); composition requires segment-aligned windows", + w.start, + seg_idx, + seg_byte_starts[seg_idx], + seg_byte_starts[seg_idx + 1] + ))); + } + all_segments.push(original_mdb.segments[seg_idx].clone()); + let entry = gap_verification.get(gap_idx).ok_or_else(|| { + DataError::InternalError(format!( + "ran out of gap_verification entries at stable segment {seg_idx}; \ + server response is inconsistent with the segment layout" + )) + })?; + all_verification.push(FileVerificationEntry::new(entry.into())); + gap_idx += 1; + seg_idx += 1; + } + // Server guarantees window ends are clamped to file_size (see xetcas + // `core::get_file_chunk_hashes`), so this invariant should always hold. + debug_assert!(w.end <= original_size, "window end {} exceeds original_size {}", w.end, original_size); + while seg_idx < n_segs && seg_byte_starts[seg_idx] < w.end { + seg_idx += 1; + } + if w.mdb.verification.len() != w.mdb.segments.len() { + return Err(DataError::InternalError(format!( + "window MDB has {} segments but {} verification entries", + w.mdb.segments.len(), + w.mdb.verification.len() + ))); + } + all_segments.extend_from_slice(&w.mdb.segments); + all_verification.extend_from_slice(&w.mdb.verification); + } + while seg_idx < n_segs { + all_segments.push(original_mdb.segments[seg_idx].clone()); + let entry = gap_verification.get(gap_idx).ok_or_else(|| { + DataError::InternalError(format!( + "ran out of gap_verification entries at stable segment {seg_idx}; \ + server response is inconsistent with the segment layout" + )) + })?; + all_verification.push(FileVerificationEntry::new(entry.into())); + gap_idx += 1; + seg_idx += 1; + } + if gap_idx < gap_verification.len() { + return Err(DataError::InternalError(format!( + "server returned {} gap_verification entries but only {} stable segments were emitted", + gap_verification.len(), + gap_idx + ))); + } + + debug_assert_eq!(all_segments.len(), all_verification.len()); + + Ok(MDBFileInfo { + metadata: FileDataSequenceHeader::new(combined_hash, all_segments.len(), true, false), + segments: all_segments, + verification: all_verification, + metadata_ext: None, + }) +} + +/// Validate the caller-provided dirty ranges. +/// +/// `dirty_inputs` must be sorted by `original_range.start`, non-overlapping +/// (`prev.original_range.end <= next.original_range.start`), and every +/// `original_range.end <= original_size`. Empty ranges (pure inserts) are allowed at any +/// position including `original_size` (append). +fn validate_dirty_ranges(dirty_inputs: &[DirtyInput], original_size: u64) -> Result<()> { + let mut prev_end = 0u64; + for (i, input) in dirty_inputs.iter().enumerate() { + let r = &input.original_range; + if r.start > r.end { + return Err(DataError::ParameterError(format!( + "dirty_inputs[{i}].original_range is reversed: {}..{}", + r.start, r.end + ))); + } + if r.end > original_size { + return Err(DataError::ParameterError(format!( + "dirty_inputs[{i}].original_range end ({}) exceeds original_size ({original_size})", + r.end + ))); + } + if i > 0 && r.start < prev_end { + return Err(DataError::ParameterError(format!( + "dirty_inputs[{i}].original_range overlaps the previous edit (starts at {} < {prev_end})", + r.start + ))); + } + prev_end = r.end; + } + Ok(()) +} + +/// Compute the resulting file size: `original_size` minus bytes removed plus bytes added. +fn compute_total_size(original_size: u64, dirty_inputs: &[DirtyInput]) -> Result { + let (removed, added) = dirty_inputs + .iter() + .fold((0u64, 0u64), |(r, a), d| (r + (d.original_range.end - d.original_range.start), a + d.new_length)); + original_size + .checked_add(added) + .and_then(|s| s.checked_sub(removed)) + .ok_or_else(|| { + DataError::ParameterError(format!( + "total size overflows: original_size={original_size}, added={added}, removed={removed}" + )) + }) +} + +/// Upload a brand-new file from `dirty_inputs` (no original to compose against). +/// Used when the original file is empty: every edit's `original_range` is `0..0`, so we +/// just stream their new bytes through the cleaner. +async fn upload_fresh_file( + config: Arc, + mut dirty_inputs: Vec, + total_size: u64, +) -> Result { + let session = FileUploadSession::new(config).await?; + let (_id, mut cleaner) = session.start_clean(None, Some(total_size), Sha256Policy::Skip)?; + for input in &mut dirty_inputs { + let mut remaining = input.new_length as usize; + let mut buf = vec![0u8; STREAM_BLOCK_SIZE.min(remaining.max(1))]; + while remaining > 0 { + let to_read = buf.len().min(remaining); + input.reader.read_exact(&mut buf[..to_read]).await.map_err(|err| { + DataError::InternalError(format!("failed to read dirty input at {}: {err}", input.original_range.start)) + })?; + cleaner.add_data(&buf[..to_read]).await?; + remaining -= to_read; + } + } + let (info, _metrics) = cleaner.finish().await?; + session.finalize().await?; + Ok(info) +} + +/// Stream a byte range from CAS into the cleaner. +async fn stream_cas_range( + ctx: &XetContext, + cas_client: &Arc, + file_hash: MerkleHash, + start: u64, + end: u64, + cleaner: &mut super::SingleFileCleaner, +) -> Result<()> { + let reconstructor = FileReconstructor::new(ctx, cas_client, file_hash).with_byte_range(FileRange::new(start, end)); + let mut stream = reconstructor.reconstruct_to_stream(); + while let Some(chunk) = stream.next().await? { + cleaner.add_data(&chunk).await?; + } + Ok(()) +} + +/// Largest segment-start byte that is `<= byte`. Used to snap a dirty-range start back +/// to its enclosing segment boundary. +fn snap_to_segment_start(seg_byte_starts: &[u64], byte: u64) -> u64 { + let idx = seg_byte_starts.partition_point(|&s| s <= byte); + seg_byte_starts[idx.saturating_sub(1)] +} + +/// Smallest segment-start byte that is `>= byte`. Used to snap a dirty-range end forward +/// to the segment boundary that fully contains it. +fn snap_to_segment_end(seg_byte_starts: &[u64], byte: u64) -> u64 { + let idx = seg_byte_starts.partition_point(|&s| s < byte); + seg_byte_starts[idx] +} + +#[cfg(test)] +mod tests { + use std::io::Cursor; + use std::ops::Range; + use std::path::Path; + use std::sync::Arc; + + use tempfile::TempDir; + use xet_client::cas_client::{Client, LocalTestServerBuilder}; + use xet_core_structures::merklehash::MerkleHash; + + use super::*; + use crate::processing::configurations::TranslatorConfig; + use crate::processing::file_cleaner::Sha256Policy; + use crate::processing::file_download_session::FileDownloadSession; + use crate::processing::file_upload_session::FileUploadSession; + + fn test_config(endpoint: impl AsRef, base_dir: impl AsRef) -> Arc { + let ctx = XetContext::default().unwrap(); + Arc::new(TranslatorConfig::test_server_config(&ctx, endpoint, base_dir).unwrap()) + } + + /// Test helper: fetch the original file's per-segment byte sizes. Tests use these to + /// build dirty ranges that align with segment (== chunk-group) boundaries — handy for + /// scenarios that want to overwrite or truncate exactly on a boundary. + async fn fetch_segment_sizes(cas_client: &Arc, hash: &MerkleHash) -> Vec { + let (mdb, _) = cas_client.get_file_reconstruction_info(hash).await.unwrap().unwrap(); + mdb.segments.iter().map(|s| s.unpacked_segment_bytes as u64).collect() + } + + /// Build in-place `DirtyInput`s (each edit's `new_length` matches its `original_range` + /// length, so the file size doesn't change) from a source buffer and range list. + fn make_dirty_inputs(ranges: &[(u64, u64)], data: &[u8]) -> Vec { + ranges + .iter() + .map(|&(start, end)| { + let slice = data[start as usize..end as usize].to_vec(); + DirtyInput { + original_range: start..end, + new_length: end - start, + reader: Box::pin(Cursor::new(slice)), + } + }) + .collect() + } + + /// Build `DirtyInput`s with dummy readers (for validation error tests where + /// the reader is never consumed). + fn make_dummy_inputs(ranges: &[(u64, u64)]) -> Vec { + ranges + .iter() + .map(|&(start, end)| DirtyInput { + original_range: start..end, + new_length: end - start, + reader: Box::pin(Cursor::new(Vec::new())), + }) + .collect() + } + + /// Bridge the old `(start, end)` + `total_size`-shaped test cases into the new edit + /// model. `(start, end)` indexes into `data` AND describes the output byte range. + /// Handles in-place edits, pure appends, mid-edit-plus-append spans (split into two + /// edits), and trailing truncations. + fn make_legacy_inputs(specs: &[(u64, u64)], data: &[u8], original_size: u64, total_size: u64) -> Vec { + let mut out: Vec = Vec::new(); + for &(start, end) in specs { + let bytes = data[start as usize..end as usize].to_vec(); + if end <= original_size { + out.push(DirtyInput { + original_range: start..end, + new_length: bytes.len() as u64, + reader: Box::pin(Cursor::new(bytes)), + }); + } else if start >= original_size { + out.push(DirtyInput { + original_range: original_size..original_size, + new_length: bytes.len() as u64, + reader: Box::pin(Cursor::new(bytes)), + }); + } else { + let split = (original_size - start) as usize; + let (head, tail) = bytes.split_at(split); + out.push(DirtyInput { + original_range: start..original_size, + new_length: head.len() as u64, + reader: Box::pin(Cursor::new(head.to_vec())), + }); + out.push(DirtyInput { + original_range: original_size..original_size, + new_length: tail.len() as u64, + reader: Box::pin(Cursor::new(tail.to_vec())), + }); + } + } + if total_size < original_size { + out.push(DirtyInput { + original_range: total_size..original_size, + new_length: 0, + reader: Box::pin(Cursor::new(Vec::new())), + }); + } + out + } + + // original: [=========================== 256 KB ===========================] + // dirty: [== 1 KB ===] + // result: [===stable===][re-uploaded][============stable=================] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_upload_ranges_mid_file_edit() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let endpoint = server.http_endpoint().to_string(); + let config = test_config(&endpoint, base_dir.path()); + + // Use the server directly as the CAS client (bypasses HTTP for get_file_chunk_hashes). + let cas_client: Arc = Arc::new(server); + + // 1. Upload an original file: 256 KB of pseudo-random bytes. + let original_data = random_data(42, 256 * 1024); + let original_hash = { + let upload_session = FileUploadSession::new(config.clone()).await.unwrap(); + let (_id, mut cleaner) = upload_session + .start_clean(Some("original".into()), Some(original_data.len() as u64), Sha256Policy::Skip) + .unwrap(); + cleaner.add_data(&original_data).await.unwrap(); + let (xfi, _metrics) = cleaner.finish().await.unwrap(); + upload_session.finalize().await.unwrap(); + MerkleHash::from_hex(xfi.hash()).unwrap() + }; + let original_size = original_data.len() as u64; + + // 2. Build modified content: overwrite [100_000, 101_000) with 0xBB. + let mut modified_data = original_data.clone(); + let dirty_start = 100_000usize; + let dirty_end = 101_000usize; + modified_data[dirty_start..dirty_end].fill(0xBB); + let total_size = modified_data.len() as u64; + let result = upload_ranges( + config.clone(), + cas_client.clone(), + original_hash, + original_size, + make_legacy_inputs(&[(dirty_start as u64, dirty_end as u64)], &modified_data, original_size, total_size), + ) + .await + .unwrap(); + + assert_eq!(result.file_size, Some(total_size)); + + // 3. Download and verify the composed file. + let composed_hash = MerkleHash::from_hex(result.hash()).unwrap(); + let session = FileDownloadSession::new(config.clone(), None).await.unwrap(); + let file_info = crate::processing::XetFileInfo::new(composed_hash.hex(), total_size); + let out_path = base_dir.path().join("output"); + session.download_file(&file_info, &out_path).await.unwrap(); + let downloaded = std::fs::read(&out_path).unwrap(); + + assert_eq!(downloaded.len(), modified_data.len()); + assert_eq!(downloaded, modified_data); + + // Hash must match a clean upload of the same content. + let clean_hash = upload_file(&config, &modified_data).await; + assert_eq!(result.hash(), clean_hash.hex(), "hash mismatch with clean upload"); + } + + // original: [=========================== 256 KB ===========================] + // result: [========= 100 KB =========] + // ^ cut here (mid-chunk) + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_upload_ranges_truncation() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + // Upload 256 KB file. + let original_data = random_data(43, 256 * 1024); + let original_hash = upload_file(&config, &original_data).await; + let original_size = original_data.len() as u64; + + // Truncate to 100 KB (no dirty ranges, pure truncation). + let truncated_size = 100_000u64; + let result = upload_ranges( + config.clone(), + cas_client.clone(), + original_hash, + original_size, + make_legacy_inputs(&[], &[], original_size, truncated_size), + ) + .await + .unwrap(); + + assert_eq!(result.file_size(), Some(truncated_size)); + + // Download and verify: first truncated_size bytes should match original. + let downloaded = download_file(&config, MerkleHash::from_hex(result.hash()).unwrap(), truncated_size).await; + assert_eq!(downloaded.len(), truncated_size as usize); + assert_eq!(downloaded, &original_data[..truncated_size as usize]); + + let clean_hash = upload_file(&config, &original_data[..truncated_size as usize]).await; + assert_eq!(result.hash(), clean_hash.hex(), "hash mismatch with clean upload"); + } + + // original: [======== 100 KB ========] + // result: [======== 100 KB ========][== 50 KB appended ==] + // ^ last chunk re-chunked with append + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_upload_ranges_append() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + // Upload 100 KB file. + let original_data = random_data(44, 100 * 1024); + let original_hash = upload_file(&config, &original_data).await; + let original_size = original_data.len() as u64; + + // Append 50 KB of pseudo-random data. + let mut full_data = original_data.clone(); + full_data.extend(random_data(99, 50 * 1024)); + let total_size = full_data.len() as u64; + let result = upload_ranges( + config.clone(), + cas_client.clone(), + original_hash, + original_size, + make_legacy_inputs(&[(original_size, total_size)], &full_data, original_size, total_size), + ) + .await + .unwrap(); + + assert_eq!(result.file_size(), Some(total_size)); + + let downloaded = download_file(&config, MerkleHash::from_hex(result.hash()).unwrap(), total_size).await; + assert_eq!(downloaded, full_data); + + let clean_hash = upload_file(&config, &full_data).await; + assert_eq!(result.hash(), clean_hash.hex(), "hash mismatch with clean upload"); + } + + // original: [=========================== 256 KB ==============================] + // dirty: [4K] + // result: [re-uploaded][=================stable=============================] + // ^ no stable prefix + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_upload_ranges_at_file_start() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + // Upload 256 KB file. + let original_data = random_data(45, 256 * 1024); + let original_hash = upload_file(&config, &original_data).await; + let original_size = original_data.len() as u64; + + // Overwrite [0, 4096) with 0xBB (dirty range at offset 0, no stable prefix). + let mut modified_data = original_data.clone(); + modified_data[..4096].fill(0xBB); + let total_size = modified_data.len() as u64; + let result = upload_ranges( + config.clone(), + cas_client.clone(), + original_hash, + original_size, + make_legacy_inputs(&[(0, 4096)], &modified_data, original_size, total_size), + ) + .await + .unwrap(); + + assert_eq!(result.file_size(), Some(total_size)); + + let downloaded = download_file(&config, MerkleHash::from_hex(result.hash()).unwrap(), total_size).await; + assert_eq!(downloaded.len(), modified_data.len()); + assert_eq!(downloaded, modified_data); + + let clean_hash = upload_file(&config, &modified_data).await; + assert_eq!(result.hash(), clean_hash.hex(), "hash mismatch with clean upload"); + } + + // original: [=========================== 256 KB ===========================] + // dirty: [2K] [2K] + // result: [s][re-up][=======stable========][re-up][========stable========] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_upload_ranges_multiple_regions() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + // Upload 256 KB file. + let original_data = random_data(46, 256 * 1024); + let original_hash = upload_file(&config, &original_data).await; + let original_size = original_data.len() as u64; + + // Two non-adjacent dirty ranges with a stable gap between them. + let mut modified_data = original_data.clone(); + modified_data[10_000..12_000].fill(0xBB); // first dirty range + modified_data[200_000..202_000].fill(0xCC); // second dirty range + let total_size = modified_data.len() as u64; + let result = upload_ranges( + config.clone(), + cas_client.clone(), + original_hash, + original_size, + make_legacy_inputs(&[(10_000, 12_000), (200_000, 202_000)], &modified_data, original_size, total_size), + ) + .await + .unwrap(); + + assert_eq!(result.file_size(), Some(total_size)); + + let downloaded = download_file(&config, MerkleHash::from_hex(result.hash()).unwrap(), total_size).await; + assert_eq!(downloaded.len(), modified_data.len()); + assert_eq!(downloaded, modified_data); + + let clean_hash = upload_file(&config, &modified_data).await; + assert_eq!(result.hash(), clean_hash.hex(), "hash mismatch with clean upload"); + } + + // original: [======== 100 KB ========] + // inputs: [======== 100 KB ========][000][=4K written=] + // ^ gap (zeros from seek past EOF) + // + // With the new API, the caller provides the full append region [original_size, total_size) + // as a single DirtyInput, including the sparse gap. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_append_with_gap_before_dirty_range() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let original_data = random_data(50, 100 * 1024); + let original_hash = upload_file(&config, &original_data).await; + let original_size = original_data.len() as u64; + + let gap = 500u64; + let write_data = random_data(101, 4096); + let total_size = original_size + gap + write_data.len() as u64; + + let mut full_data = original_data.clone(); + full_data.extend(vec![0x00u8; gap as usize]); + full_data.extend(&write_data); + + // The caller provides the entire append region (including the sparse gap). + let result = upload_ranges( + config.clone(), + cas_client.clone(), + original_hash, + original_size, + make_legacy_inputs(&[(original_size, total_size)], &full_data, original_size, total_size), + ) + .await + .unwrap(); + + assert_eq!(result.file_size(), Some(total_size)); + + let downloaded = download_file(&config, MerkleHash::from_hex(result.hash()).unwrap(), total_size).await; + assert_eq!(downloaded.len(), full_data.len(), "size mismatch"); + assert_eq!(&downloaded[..], &full_data[..], "content mismatch: gap bytes were lost"); + + let clean_hash = upload_file(&config, &full_data).await; + assert_eq!(result.hash(), clean_hash.hex(), "hash mismatch with clean upload"); + } + + // staging: [000000 (zeros) 000000][=== appended ===] + // CAS: [=== original data ==] + // result: [=== original data ==][=== appended ===] + // ^ boundary prefix must come from CAS, not from staging zeros + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_append_sparse_staging_file() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let original_data = vec![0xDDu8; 100 * 1024]; + let original_hash = upload_file(&config, &original_data).await; + let original_size = original_data.len() as u64; + + let append_data = vec![0xEEu8; 50 * 1024]; + let total_size = original_size + append_data.len() as u64; + + // Build a sparse staging file: zeros for [0, original_size), real data after. + // This is what hf-mount produces (sparse hole + appended bytes). + let mut sparse_staging = vec![0u8; total_size as usize]; + sparse_staging[original_size as usize..].copy_from_slice(&append_data); + let result = upload_ranges( + config.clone(), + cas_client.clone(), + original_hash, + original_size, + make_legacy_inputs(&[(original_size, total_size)], &sparse_staging, original_size, total_size), + ) + .await + .unwrap(); + + // Expected: original data + appended data (not zeros + appended data). + let mut expected = original_data.clone(); + expected.extend(&append_data); + + let downloaded = download_file(&config, MerkleHash::from_hex(result.hash()).unwrap(), total_size).await; + assert_eq!(downloaded.len(), expected.len(), "size mismatch"); + assert_eq!(&downloaded[..], &expected[..], "content mismatch: CAS data replaced by zeros from sparse file"); + + let clean_hash = upload_file(&config, &expected).await; + assert_eq!(result.hash(), clean_hash.hex(), "hash mismatch with clean upload"); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_data_integrity_scenarios() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + // original: [========================= 256 KB =========================] + // dirty: [10K] + // result: [====== 90 KB ======][10K] + // 0 90K 100K ^ truncate here, dirty touches cut + { + let original = vec![0xAAu8; 256 * 1024]; + let mut expected = original[..100_000].to_vec(); + expected[90_000..100_000].fill(0xBB); + assert_range_edit(&config, &cas_client, &original, &expected, &[(90_000, 100_000)], 100_000).await; + } + + // original: [========= 128 KB =========] + // dirty: [========= 128 KB =========] + // result: [======= re-uploaded ======] (no stable regions) + { + let original = vec![0xAAu8; 128 * 1024]; + let expected = vec![0xBBu8; 128 * 1024]; + let size = original.len() as u64; + assert_range_edit(&config, &cas_client, &original, &expected, &[(0, size)], size).await; + } + + // original: [===================== 256 KB =====================] + // dirty: [1K][1K][1K] + // ^-- coalesced into one region + { + let original = vec![0xAAu8; 256 * 1024]; + let mut expected = original.clone(); + expected[50_000..51_000].fill(0xBB); + expected[51_000..52_000].fill(0xCC); + expected[52_000..53_000].fill(0xDD); + let size = original.len() as u64; + assert_range_edit( + &config, + &cas_client, + &original, + &expected, + &[(50_000, 51_000), (51_000, 52_000), (52_000, 53_000)], + size, + ) + .await; + } + + // original: [======== 100 KB ========] + // result: [======== 100 KB ========][== 50 KB ==] + // no dirty_ranges, only total_size > original_size + { + let original = vec![0xAAu8; 100 * 1024]; + let mut expected = original.clone(); + expected.extend(vec![0xEEu8; 50 * 1024]); + let total = expected.len() as u64; + assert_range_edit(&config, &cas_client, &original, &expected, &[], total).await; + } + + // original: [chunk0][chunk1][chunk2][...] + // dirty: [chunk2] + // ^ ^-- starts/ends on chunk boundary + { + let original: Vec = (0..256 * 1024) + .map(|i: usize| { + let x = i.wrapping_mul(2654435761); + (x >> 16) as u8 + }) + .collect(); + let original_hash = upload_file(&config, &original).await; + let seg_sizes = fetch_segment_sizes(&cas_client, &original_hash).await; + if seg_sizes.len() >= 3 { + let boundary: u64 = seg_sizes[0] + seg_sizes[1]; + let dirty_end = boundary + seg_sizes[2]; + let mut expected = original.clone(); + expected[boundary as usize..dirty_end as usize].fill(0xFF); + let size = original.len() as u64; + let result = upload_ranges( + config.clone(), + cas_client.clone(), + original_hash, + size, + make_legacy_inputs(&[(boundary, dirty_end)], &expected, size, size), + ) + .await + .unwrap(); + let downloaded = download_file(&config, MerkleHash::from_hex(result.hash()).unwrap(), size).await; + assert_eq!(downloaded, expected, "chunk-boundary edit mismatch"); + + let clean_hash = upload_file(&config, &expected).await; + assert_eq!(result.hash(), clean_hash.hex(), "hash mismatch with clean upload"); + } + } + } + + // No changes: dirty_ranges=[], total_size == original_size -> early return. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_noop_returns_original_hash() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let data = random_data(70, 256 * 1024); + let hash = upload_file(&config, &data).await; + let size = data.len() as u64; + let result = upload_ranges(config, cas_client, hash, size, make_legacy_inputs(&[], &[], size, size)) + .await + .unwrap(); + + assert_eq!(result.hash(), hash.hex()); + assert_eq!(result.file_size(), Some(size)); + } + + // dirty_range end > total_size -> rejected. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_rejects_dirty_range_past_total_size() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let data = random_data(71, 256 * 1024); + let hash = upload_file(&config, &data).await; + let size = data.len() as u64; + let err = upload_ranges(config, cas_client, hash, size, make_dummy_inputs(&[(100, size + 1)])).await; + assert!(err.is_err(), "dirty range past total_size should be rejected"); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_rejects_overlapping_dirty_ranges() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let data = random_data(60, 256 * 1024); + let hash = upload_file(&config, &data).await; + let size = data.len() as u64; + let err = upload_ranges(config, cas_client, hash, size, make_dummy_inputs(&[(100, 300), (200, 400)])).await; + assert!(err.is_err(), "overlapping ranges should be rejected"); + } + + // Regression: validation must run *before* the `original_size == 0` short-circuit, or + // the empty-original path silently accepts ranges that exceed `original_size`. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_empty_original_validates_ranges() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let original_hash = upload_file(&config, &[]).await; + + // `0..10` and `5..15` both have `end > original_size == 0`; either would corrupt + // the upload if validation didn't fire. + let inputs = vec![ + DirtyInput { + original_range: 0..10, + new_length: 10, + reader: Box::pin(Cursor::new(vec![0xAA; 10])), + }, + DirtyInput { + original_range: 5..15, + new_length: 10, + reader: Box::pin(Cursor::new(vec![0xBB; 10])), + }, + ]; + let err = upload_ranges(config, cas_client, original_hash, 0, inputs).await; + assert!(err.is_err(), "ranges with end > original_size must be rejected for empty originals too"); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_rejects_unsorted_dirty_ranges() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let data = random_data(62, 256 * 1024); + let hash = upload_file(&config, &data).await; + let size = data.len() as u64; + let err = upload_ranges(config, cas_client, hash, size, make_dummy_inputs(&[(300, 400), (100, 200)])).await; + assert!(err.is_err(), "unsorted ranges should be rejected"); + } + + // original: [chunk0][chunk1][chunk2][chunk3][...more chunks...] + // input: [========= single large write ==========] + // + // A single DirtyInput that spans many segments. Verifies the reader is consumed + // correctly across the multi-segment window. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_single_input_spanning_many_chunks() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let original_data = random_data(99, 256 * 1024); + let original_hash = upload_file(&config, &original_data).await; + let original_size = original_data.len() as u64; + + // Overwrite a large middle section (likely spans many CDC chunks). + let mut modified = original_data.clone(); + let dirty_start = 10_000u64; + let dirty_end = 200_000u64; + modified[dirty_start as usize..dirty_end as usize].fill(0xFF); + + let result = upload_ranges( + config.clone(), + cas_client.clone(), + original_hash, + original_size, + make_legacy_inputs(&[(dirty_start, dirty_end)], &modified, original_size, original_size), + ) + .await + .unwrap(); + + let downloaded = download_file(&config, MerkleHash::from_hex(result.hash()).unwrap(), original_size).await; + assert_eq!(downloaded, modified, "large spanning input produced wrong content"); + + let clean_hash = upload_file(&config, &modified).await; + assert_eq!(result.hash(), clean_hash.hex(), "hash mismatch with clean upload"); + } + + // original: b"AAAA_HEADER_AAAA|" (17 bytes, single CAS chunk) + // dirty: [SPARSE] (bytes [5, 11)) + // expected: b"AAAA_SPARSE_AAAA|" (17 bytes) + // + // Tests mid-file edit on a very small file (single chunk, smaller than + // typical CDC minimum). See test_truncate_then_mid_edit for the regression + // test that reproduces the real production bug. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_upload_ranges_small_file_mid_edit() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let original_data = b"AAAA_HEADER_AAAA|"; + let original_hash = upload_file(&config, original_data).await; + let original_size = original_data.len() as u64; + + let dirty_data = b"SPARSE"; + let dirty_inputs = vec![DirtyInput { + original_range: 5..11, + new_length: dirty_data.len() as u64, + reader: Box::pin(Cursor::new(dirty_data.to_vec())), + }]; + + let result = upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, dirty_inputs) + .await + .unwrap(); + + assert_eq!(result.file_size(), Some(original_size)); + + let downloaded = download_file(&config, MerkleHash::from_hex(result.hash()).unwrap(), original_size).await; + assert_eq!(downloaded.len(), original_size as usize, "reconstructed size mismatch"); + assert_eq!(&downloaded[..5], b"AAAA_", "prefix from CAS"); + assert_eq!(&downloaded[5..11], b"SPARSE", "dirty range"); + assert_eq!(&downloaded[11..], b"_AAAA|", "suffix from CAS"); + + let expected = b"AAAA_SPARSE_AAAA|"; + let clean_hash = upload_file(&config, expected).await; + assert_eq!(result.hash(), clean_hash.hex(), "hash mismatch with clean upload"); + } + + // original: [=========================== 256 KB ===========================] + // staging: [0000000000000000000000000000] (all zeros, file never opened for write) + // result: [====== 100 KB from CAS =====] + // ^ cut here (mid-chunk) + // + // The boundary chunk bytes must come from CAS, not from the zero-filled staging. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_upload_ranges_truncation_empty_staging() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let original_data = random_data(77, 256 * 1024); + let original_hash = upload_file(&config, &original_data).await; + let original_size = original_data.len() as u64; + + let truncated_size = 100_000u64; + + let result = upload_ranges( + config.clone(), + cas_client.clone(), + original_hash, + original_size, + make_legacy_inputs(&[], &[], original_size, truncated_size), + ) + .await + .unwrap(); + + assert_eq!(result.file_size(), Some(truncated_size)); + + let downloaded = download_file(&config, MerkleHash::from_hex(result.hash()).unwrap(), truncated_size).await; + assert_eq!(downloaded.len(), truncated_size as usize); + assert_eq!( + &downloaded[..], + &original_data[..truncated_size as usize], + "truncated content should match original CAS data, not staging zeros" + ); + + let clean_hash = upload_file(&config, &original_data[..truncated_size as usize]).await; + assert_eq!(result.hash(), clean_hash.hex(), "hash mismatch with clean upload"); + } + + // original: [=========================== 256 KB ===========================] + // staging: [000000000000][0xBB][0000000] (zeros except the dirty range) + // ^ ^ + // 90K 95K (dirty from caller) + // result: [==CAS==][stg][===CAS===] + // ^ cut at 100K (mid-chunk) + // + // Dirty bytes [90K,95K) come from staging; boundary bytes from CAS. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_upload_ranges_truncation_with_overlapping_dirty() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let original_data = random_data(88, 256 * 1024); + let original_hash = upload_file(&config, &original_data).await; + let original_size = original_data.len() as u64; + + let truncated_size = 100_000u64; + + let dirty_start = 90_000u64; + let dirty_end = 95_000u64; + + let mut expected = original_data[..truncated_size as usize].to_vec(); + expected[dirty_start as usize..dirty_end as usize].fill(0xBB); + + let mut staging = vec![0u8; truncated_size as usize]; + staging[dirty_start as usize..dirty_end as usize].fill(0xBB); + let result = upload_ranges( + config.clone(), + cas_client.clone(), + original_hash, + original_size, + make_legacy_inputs(&[(dirty_start, dirty_end)], &staging, original_size, truncated_size), + ) + .await + .unwrap(); + + assert_eq!(result.file_size(), Some(truncated_size)); + + let downloaded = download_file(&config, MerkleHash::from_hex(result.hash()).unwrap(), truncated_size).await; + assert_eq!(downloaded.len(), expected.len()); + assert_eq!(&downloaded[..], &expected[..], "dirty bytes should come from staging, boundary bytes from CAS"); + + let clean_hash = upload_file(&config, &expected).await; + assert_eq!(result.hash(), clean_hash.hex(), "hash mismatch with clean upload"); + } + + // ── Helpers ────────────────────────────────────────────────────── + + fn random_data(seed: u64, len: usize) -> Vec { + (0..len) + .map(|i| { + let x = (i as u64).wrapping_add(seed).wrapping_mul(2654435761); + (x >> 16) as u8 + }) + .collect() + } + + #[derive(Clone, Debug)] + struct DeterministicRng { + state: u64, + } + + impl DeterministicRng { + fn new(seed: u64) -> Self { + Self { state: seed } + } + + fn next_u64(&mut self) -> u64 { + self.state = self.state.wrapping_mul(6364136223846793005).wrapping_add(1); + self.state + } + + fn gen_range(&mut self, start: usize, end: usize) -> usize { + if end <= start { + return start; + } + start + (self.next_u64() as usize % (end - start)) + } + + fn gen_bytes(&mut self, len: usize) -> Vec { + (0..len).map(|_| (self.next_u64() >> 56) as u8).collect() + } + } + + #[derive(Clone, Debug)] + struct PlannedEdit { + original_range: Range, + replacement: Vec, + } + + fn build_random_non_overlapping_edits( + rng: &mut DeterministicRng, + original_len: usize, + max_edits: usize, + ) -> Vec { + if original_len == 0 { + let replacement_len = 1 + rng.gen_range(0, 8 * 1024); + return vec![PlannedEdit { + original_range: 0..0, + replacement: rng.gen_bytes(replacement_len), + }]; + } + + let target_edits = 1 + rng.gen_range(0, max_edits.max(1)); + let mut edits: Vec = Vec::with_capacity(target_edits); + let mut cursor = 0usize; + + while edits.len() < target_edits && cursor <= original_len { + let remaining = original_len - cursor; + let max_gap = remaining.min(64 * 1024); + let start = cursor + rng.gen_range(0, max_gap + 1); + + let (end, replacement_len) = if start == original_len { + (start, 1 + rng.gen_range(0, 32 * 1024)) + } else { + let op = rng.gen_range(0, 5); + let max_span = (original_len - start).clamp(1, 64 * 1024); + let span = 1 + rng.gen_range(0, max_span); + let end = start + span; + match op { + 0 => (start, 1 + rng.gen_range(0, 32 * 1024)), + 1 => (end, span), + 2 => (end, span + 1 + rng.gen_range(0, 16 * 1024)), + 3 => (end, rng.gen_range(0, span + 1)), + _ => (end, 0), + } + }; + + edits.push(PlannedEdit { + original_range: start..end, + replacement: rng.gen_bytes(replacement_len), + }); + cursor = if end > start { end } else { start.saturating_add(1) }; + } + + if edits.is_empty() { + let replacement_len = 1 + rng.gen_range(0, 32 * 1024); + edits.push(PlannedEdit { + original_range: original_len..original_len, + replacement: rng.gen_bytes(replacement_len), + }); + } + + for w in edits.windows(2) { + assert!(w[0].original_range.end <= w[1].original_range.start); + } + + edits + } + + fn apply_planned_edits(original: &[u8], edits: &[PlannedEdit]) -> Vec { + let removed: usize = edits.iter().map(|e| e.original_range.end - e.original_range.start).sum(); + let added: usize = edits.iter().map(|e| e.replacement.len()).sum(); + let mut out: Vec = Vec::with_capacity(original.len() + added.saturating_sub(removed)); + let mut cursor = 0usize; + + for edit in edits { + assert!(edit.original_range.start >= cursor); + out.extend_from_slice(&original[cursor..edit.original_range.start]); + out.extend_from_slice(&edit.replacement); + cursor = edit.original_range.end; + } + + out.extend_from_slice(&original[cursor..]); + out + } + + fn edits_to_dirty_inputs(edits: &[PlannedEdit]) -> Vec { + edits + .iter() + .map(|e| DirtyInput { + original_range: e.original_range.start as u64..e.original_range.end as u64, + new_length: e.replacement.len() as u64, + reader: Box::pin(Cursor::new(e.replacement.clone())), + }) + .collect() + } + + fn summarize_edits(edits: &[PlannedEdit]) -> String { + edits + .iter() + .map(|e| format!("[{}..{}, new_len={}]", e.original_range.start, e.original_range.end, e.replacement.len())) + .collect::>() + .join(", ") + } + + async fn upload_file(config: &Arc, data: &[u8]) -> MerkleHash { + let session = FileUploadSession::new(config.clone()).await.unwrap(); + let (_id, mut cleaner) = session + .start_clean(Some("test".into()), Some(data.len() as u64), Sha256Policy::Skip) + .unwrap(); + cleaner.add_data(data).await.unwrap(); + let (xfi, _metrics) = cleaner.finish().await.unwrap(); + session.finalize().await.unwrap(); + MerkleHash::from_hex(xfi.hash()).unwrap() + } + + async fn download_file(config: &Arc, hash: MerkleHash, size: u64) -> Vec { + let session = FileDownloadSession::new(config.clone(), None).await.unwrap(); + let xfi = crate::processing::XetFileInfo::new(hash.hex(), size); + let dir = TempDir::new().unwrap(); + let out = dir.path().join("out"); + session.download_file(&xfi, &out).await.unwrap(); + std::fs::read(&out).unwrap() + } + + /// Empty `Box::pin(Cursor::new(Vec::new()))` for delete / dummy edits. + fn empty_reader() -> Pin> { + Box::pin(Cursor::new(Vec::::new())) + } + + /// End-to-end check: upload `original`, apply `inputs` via `upload_ranges`, download, + /// compare against `expected`, and verify the hash matches a clean upload of `expected`. + async fn assert_edits( + config: &Arc, + cas_client: &Arc, + original: &[u8], + inputs: Vec, + expected: &[u8], + ) { + let original_hash = upload_file(config, original).await; + let original_size = original.len() as u64; + let result = upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, inputs) + .await + .unwrap(); + assert_eq!(result.file_size(), Some(expected.len() as u64), "file size mismatch"); + let downloaded = + download_file(config, MerkleHash::from_hex(result.hash()).unwrap(), expected.len() as u64).await; + assert_eq!(downloaded, expected, "content mismatch"); + let clean = upload_file(config, expected).await; + assert_eq!(result.hash(), clean.hex(), "hash diverges from clean upload"); + } + + async fn assert_range_edit( + config: &Arc, + cas_client: &Arc, + original_data: &[u8], + expected: &[u8], + dirty_ranges: &[(u64, u64)], + total_size: u64, + ) { + let original_hash = upload_file(config, original_data).await; + let original_size = original_data.len() as u64; + + // Build dirty inputs from the caller's ranges. + let mut inputs = make_dirty_inputs(dirty_ranges, expected); + + // For appends, ensure the appended region is included as a pure-insert edit at the + // end of the original. + if total_size > original_size { + let append_start = original_size; + let already_covered = dirty_ranges.iter().any(|&(s, e)| s <= append_start && e >= total_size); + if !already_covered { + inputs.push(DirtyInput { + original_range: original_size..original_size, + new_length: total_size - original_size, + reader: Box::pin(Cursor::new(expected[append_start as usize..total_size as usize].to_vec())), + }); + inputs.sort_by_key(|d| d.original_range.start); + } + } + + // Truncation: drop bytes past `total_size` from the original via a pure-delete edit. + if total_size < original_size { + inputs.push(DirtyInput { + original_range: total_size..original_size, + new_length: 0, + reader: empty_reader(), + }); + inputs.sort_by_key(|d| d.original_range.start); + } + + let result = upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, inputs) + .await + .unwrap(); + + assert_eq!(result.file_size(), Some(total_size), "file size mismatch"); + let downloaded = download_file(config, MerkleHash::from_hex(result.hash()).unwrap(), total_size).await; + assert_eq!(downloaded.len(), expected.len(), "downloaded length mismatch"); + assert_eq!(&downloaded[..], expected, "content mismatch"); + + let clean_hash = upload_file(config, expected).await; + assert_eq!(result.hash(), clean_hash.hex(), "hash mismatch with clean upload"); + } + + // Regression: mid-file edit + tail append in a single call. Codex caught that the last + // existing window (covering the mid-file edit) was being stretched to `total_size`, + // dropping the stable bytes between the edit and the appended region. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_mid_edit_plus_append() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let original_data = random_data(7, 256 * 1024); + let original_hash = upload_file(&config, &original_data).await; + let original_size = original_data.len() as u64; + + let dirty_start = 50_000usize; + let dirty_end = 51_000usize; + let append_extra: Vec = (0..16 * 1024).map(|i| (i % 251) as u8).collect(); + let mut expected = original_data.clone(); + expected[dirty_start..dirty_end].fill(0xAA); + expected.extend_from_slice(&append_extra); + let total_size = expected.len() as u64; + + let inputs = vec![ + DirtyInput { + original_range: dirty_start as u64..dirty_end as u64, + new_length: (dirty_end - dirty_start) as u64, + reader: Box::pin(Cursor::new(expected[dirty_start..dirty_end].to_vec())), + }, + DirtyInput { + original_range: original_size..original_size, + new_length: append_extra.len() as u64, + reader: Box::pin(Cursor::new(append_extra)), + }, + ]; + let result = upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, inputs) + .await + .unwrap(); + + let downloaded = download_file(&config, MerkleHash::from_hex(result.hash()).unwrap(), total_size).await; + assert_eq!(downloaded, expected, "content mismatch (mid-edit + append regression)"); + let clean_hash = upload_file(&config, &expected).await; + assert_eq!(result.hash(), clean_hash.hex(), "hash mismatch with clean upload"); + } + + // Regression: empty original + append. Codex caught that we'd return an internal error + // instead of treating it as a fresh upload. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_empty_original_append() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let original_data: &[u8] = &[]; + let original_hash = upload_file(&config, original_data).await; + let new_data: Vec = (0..32 * 1024).map(|i| (i % 251) as u8).collect(); + let total_size = new_data.len() as u64; + + let inputs = vec![DirtyInput { + original_range: 0..0, + new_length: total_size, + reader: Box::pin(Cursor::new(new_data.clone())), + }]; + let result = upload_ranges(config.clone(), cas_client.clone(), original_hash, 0, inputs) + .await + .unwrap(); + + let downloaded = download_file(&config, MerkleHash::from_hex(result.hash()).unwrap(), total_size).await; + assert_eq!(downloaded, new_data); + let clean_hash = upload_file(&config, &new_data).await; + assert_eq!(result.hash(), clean_hash.hex(), "hash mismatch with clean upload"); + } + + // Regression: truncating to empty must produce the canonical empty-file hash + // (`MerkleHash::default()` without HMAC), not `default.hmac(default)`. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_truncate_to_empty_matches_clean_empty() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let original_data = random_data(11, 64 * 1024); + let original_hash = upload_file(&config, &original_data).await; + let original_size = original_data.len() as u64; + + let result = upload_ranges( + config.clone(), + cas_client.clone(), + original_hash, + original_size, + make_legacy_inputs(&[], &[], original_size, 0), + ) + .await + .unwrap(); + + let clean_empty = upload_file(&config, &[]).await; + assert_eq!(result.hash(), clean_empty.hex(), "truncate-to-empty must match clean empty upload hash"); + } + + // The three small examples that motivated the resize-edit API: replace, insert, delete + // on a tiny 3-byte file. Each produces an output of a different length than the input. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_resize_edits_abc() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + // abc + replace [0, 1) with "foo" => "foobc" + assert_edits( + &config, + &cas_client, + b"abc", + vec![DirtyInput { + original_range: 0..1, + new_length: 3, + reader: Box::pin(Cursor::new(b"foo".to_vec())), + }], + b"foobc", + ) + .await; + + // abc + insert "foo" at 0 => "fooabc" + assert_edits( + &config, + &cas_client, + b"abc", + vec![DirtyInput { + original_range: 0..0, + new_length: 3, + reader: Box::pin(Cursor::new(b"foo".to_vec())), + }], + b"fooabc", + ) + .await; + + // abc + delete [0, 1) => "bc" + assert_edits( + &config, + &cas_client, + b"abc", + vec![DirtyInput { + original_range: 0..1, + new_length: 0, + reader: empty_reader(), + }], + b"bc", + ) + .await; + } + + // original: [============= 256 KB =============] + // edit: [== 4K stale ==] + // ^ replaced with 32 K (resize +28 K) + // result: [== prefix ==][== 32K new ==][== suffix ==] + // + // Big in-place replace where new_length >> original_range.len(); exercises the + // resize branch with a multi-segment-touching window. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_resize_large_replace_grows_file() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let original = random_data(101, 256 * 1024); + let drop_start = 100_000usize; + let drop_end = 104_000usize; + let new_bytes: Vec = (0..32 * 1024).map(|i| (i % 251) as u8).collect(); + let mut expected = Vec::with_capacity(original.len() - (drop_end - drop_start) + new_bytes.len()); + expected.extend_from_slice(&original[..drop_start]); + expected.extend_from_slice(&new_bytes); + expected.extend_from_slice(&original[drop_end..]); + + assert_edits( + &config, + &cas_client, + &original, + vec![DirtyInput { + original_range: drop_start as u64..drop_end as u64, + new_length: new_bytes.len() as u64, + reader: Box::pin(Cursor::new(new_bytes)), + }], + &expected, + ) + .await; + } + + // original: [============= 256 KB =============] + // edit: [============= 80 KB stale =============] + // ^ replaced with 4 K (resize -76 K) + // result: [== prefix ==][4 K new][== suffix ==] + // + // Big in-place replace where new_length << original_range.len(); shrink case. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_resize_large_replace_shrinks_file() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let original = random_data(102, 256 * 1024); + let drop_start = 80_000usize; + let drop_end = 160_000usize; + let new_bytes: Vec = (0..4 * 1024).map(|i| (0xCC ^ i) as u8).collect(); + let mut expected = Vec::with_capacity(original.len() - (drop_end - drop_start) + new_bytes.len()); + expected.extend_from_slice(&original[..drop_start]); + expected.extend_from_slice(&new_bytes); + expected.extend_from_slice(&original[drop_end..]); + + assert_edits( + &config, + &cas_client, + &original, + vec![DirtyInput { + original_range: drop_start as u64..drop_end as u64, + new_length: new_bytes.len() as u64, + reader: Box::pin(Cursor::new(new_bytes)), + }], + &expected, + ) + .await; + } + + // original: [============= 100 KB =============] + // edit: [insert 8 KB here] + // result: [== prefix ==][== 8 KB new ==][======= suffix =======] + // + // Pure mid-file insert (range start == end), large payload that forces re-chunking + // around the insertion point. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_resize_mid_file_insert() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let original = random_data(103, 100 * 1024); + let at = 40_000usize; + let new_bytes: Vec = (0..8u32 * 1024).map(|i| (i.wrapping_mul(13) % 251) as u8).collect(); + let mut expected = Vec::with_capacity(original.len() + new_bytes.len()); + expected.extend_from_slice(&original[..at]); + expected.extend_from_slice(&new_bytes); + expected.extend_from_slice(&original[at..]); + + assert_edits( + &config, + &cas_client, + &original, + vec![DirtyInput { + original_range: at as u64..at as u64, + new_length: new_bytes.len() as u64, + reader: Box::pin(Cursor::new(new_bytes)), + }], + &expected, + ) + .await; + } + + // original: [============= 256 KB =============] + // edit: [== 64 KB hole ==] + // ^ delete this slice (no replacement) + // result: [== prefix ==][== suffix ==] + // + // Pure mid-file delete: shrinks the file without touching the tail boundary. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_resize_mid_file_delete() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let original = random_data(104, 256 * 1024); + let drop_start = 80_000usize; + let drop_end = 144_000usize; + let mut expected = Vec::with_capacity(original.len() - (drop_end - drop_start)); + expected.extend_from_slice(&original[..drop_start]); + expected.extend_from_slice(&original[drop_end..]); + + assert_edits( + &config, + &cas_client, + &original, + vec![DirtyInput { + original_range: drop_start as u64..drop_end as u64, + new_length: 0, + reader: empty_reader(), + }], + &expected, + ) + .await; + } + + // original: [== seg0 ==][== seg1 ==][== seg2 ==] + // edits: [shrink][grow][ delete ] + // + // Three independent edits in one call: an in-place shrink, a pure insert in seg1, and + // a pure delete in seg2. Mix of resize directions, far enough apart that they don't + // coalesce into one window. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_resize_multi_edit_mix() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let original = random_data(105, 384 * 1024); + let (a_start, a_end) = (10 * 1024usize, 20 * 1024usize); + let a_new: Vec = vec![0xAA; 2 * 1024]; + let b_at = 150 * 1024usize; + let b_new: Vec = vec![0xBB; 4 * 1024]; + let (c_start, c_end) = (300 * 1024usize, 320 * 1024usize); + + let mut expected = Vec::with_capacity(original.len() + b_new.len()); + expected.extend_from_slice(&original[..a_start]); + expected.extend_from_slice(&a_new); + expected.extend_from_slice(&original[a_end..b_at]); + expected.extend_from_slice(&b_new); + expected.extend_from_slice(&original[b_at..c_start]); + expected.extend_from_slice(&original[c_end..]); + + assert_edits( + &config, + &cas_client, + &original, + vec![ + DirtyInput { + original_range: a_start as u64..a_end as u64, + new_length: a_new.len() as u64, + reader: Box::pin(Cursor::new(a_new)), + }, + DirtyInput { + original_range: b_at as u64..b_at as u64, + new_length: b_new.len() as u64, + reader: Box::pin(Cursor::new(b_new)), + }, + DirtyInput { + original_range: c_start as u64..c_end as u64, + new_length: 0, + reader: empty_reader(), + }, + ], + &expected, + ) + .await; + } + + // original: [== seg0 ==][== seg1 ==] + // ^ insert here, exactly on segment boundary + // result: [== seg0 ==][== inserted ==][== seg1 ==] + // + // Pure insert on a segment boundary mid-file. Snap picks the segment starting at the + // boundary; insert lands at `w_start` of that window. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_resize_insert_at_segment_boundary() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let original = random_data(106, 200 * 1024); + let original_hash = upload_file(&config, &original).await; + let original_size = original.len() as u64; + + // Pick the first interior segment boundary; bail if the file lands in one segment. + let seg_sizes = fetch_segment_sizes(&cas_client, &original_hash).await; + let Some(boundary) = seg_sizes + .iter() + .scan(0u64, |acc, s| { + *acc += s; + Some(*acc) + }) + .find(|&b| b > 0 && b < original_size) + else { + return; + }; + + let new_bytes: Vec = vec![0x42; 4 * 1024]; + let mut expected = Vec::with_capacity(original.len() + new_bytes.len()); + expected.extend_from_slice(&original[..boundary as usize]); + expected.extend_from_slice(&new_bytes); + expected.extend_from_slice(&original[boundary as usize..]); + + assert_edits( + &config, + &cas_client, + &original, + vec![DirtyInput { + original_range: boundary..boundary, + new_length: new_bytes.len() as u64, + reader: Box::pin(Cursor::new(new_bytes)), + }], + &expected, + ) + .await; + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + #[ignore = "stress test"] + async fn test_stress_random_resize_sequences() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + for seed in 0..6u64 { + let mut rng = DeterministicRng::new(0x9E37_79B9_7F4A_7C15 ^ seed.wrapping_mul(0xD1B5_4A32_D192_ED03)); + let mut expected = random_data(10_000 + seed, 1_048_576 + (seed as usize * 91_117 % 262_144)); + let mut original_hash = upload_file(&config, &expected).await; + let mut original_size = expected.len() as u64; + + for round in 0..25usize { + let edits = build_random_non_overlapping_edits(&mut rng, expected.len(), 8); + let expected_next = apply_planned_edits(&expected, &edits); + let inputs = edits_to_dirty_inputs(&edits); + let result = upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, inputs) + .await + .unwrap(); + let result_hash = MerkleHash::from_hex(result.hash()).unwrap(); + + assert_eq!( + result.file_size(), + Some(expected_next.len() as u64), + "seed={seed}, round={round}: size mismatch" + ); + let clean_hash = upload_file(&config, &expected_next).await; + assert_eq!(result.hash(), clean_hash.hex(), "seed={seed}, round={round}: hash mismatch"); + + let downloaded = download_file(&config, result_hash, expected_next.len() as u64).await; + assert_eq!(downloaded, expected_next, "seed={seed}, round={round}: content mismatch"); + + expected = expected_next; + original_hash = result_hash; + original_size = expected.len() as u64; + } + } + } + + #[cfg(not(feature = "smoke-test"))] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_regression_hash_matches_clean_upload_seed1_round17() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let seed = 1u64; + let mut rng = DeterministicRng::new(0x9E37_79B9_7F4A_7C15 ^ seed.wrapping_mul(0xD1B5_4A32_D192_ED03)); + let mut expected = random_data(10_000 + seed, 1_048_576 + (seed as usize * 91_117 % 262_144)); + let mut original_hash = upload_file(&config, &expected).await; + let mut original_size = expected.len() as u64; + + for round in 0..=17usize { + let edits = build_random_non_overlapping_edits(&mut rng, expected.len(), 8); + let edits_summary = summarize_edits(&edits); + let expected_next = apply_planned_edits(&expected, &edits); + let inputs = edits_to_dirty_inputs(&edits); + let result = upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, inputs) + .await + .unwrap(); + let result_hash = MerkleHash::from_hex(result.hash()).unwrap(); + + let clean_hash = upload_file(&config, &expected_next).await; + assert_eq!( + result.hash(), + clean_hash.hex(), + "seed={seed}, round={round}: hash mismatch; original_size={original_size}, expected_size={}, edits={edits_summary}", + expected_next.len() + ); + + let downloaded = download_file(&config, result_hash, expected_next.len() as u64).await; + assert_eq!(downloaded, expected_next, "seed={seed}, round={round}: content mismatch"); + + expected = expected_next; + original_hash = result_hash; + original_size = expected.len() as u64; + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + #[ignore = "stress test"] + async fn test_stress_many_sparse_windows_single_call() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let original = random_data(13_337, 16 * 1024 * 1024); + let mut rng = DeterministicRng::new(0xA5A5_5A5A_0123_4567); + let mut edits: Vec = Vec::new(); + let stride = original.len() / 200; + let mut cursor = stride / 2; + + while edits.len() < 128 && cursor < original.len() { + let start = cursor; + let max_span = (original.len() - start).clamp(1, 1536); + let span = 128 + rng.gen_range(0, max_span); + let end = (start + span).min(original.len()); + let replacement_len = match rng.gen_range(0, 4) { + 0 => end - start, + 1 => (end - start) + 64 + rng.gen_range(0, 512), + 2 => rng.gen_range(0, end - start + 1), + _ => 0, + }; + edits.push(PlannedEdit { + original_range: start..end, + replacement: rng.gen_bytes(replacement_len), + }); + cursor = cursor.saturating_add(stride.max(1)); + } + + let expected = apply_planned_edits(&original, &edits); + let inputs = edits_to_dirty_inputs(&edits); + assert_edits(&config, &cas_client, &original, inputs, &expected).await; + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 8)] + #[ignore = "stress test"] + async fn test_stress_parallel_random_resize_sequences() { + let server = LocalTestServerBuilder::new().start().await; + let base_dir = TempDir::new().unwrap(); + let config = test_config(server.http_endpoint(), base_dir.path()); + let cas_client: Arc = Arc::new(server); + + let mut handles = Vec::new(); + for worker in 0..8u64 { + let config = config.clone(); + let cas_client = cas_client.clone(); + handles.push(tokio::spawn(async move { + let mut rng = DeterministicRng::new(0xC0FF_EE00_1234_5678 ^ worker.wrapping_mul(0x94D0_49BB_1331_11EB)); + let mut expected = random_data(20_000 + worker, 786_432 + worker as usize * 17_321); + let mut original_hash = upload_file(&config, &expected).await; + let mut original_size = expected.len() as u64; + + for round in 0..18usize { + let edits = build_random_non_overlapping_edits(&mut rng, expected.len(), 6); + let expected_next = apply_planned_edits(&expected, &edits); + let inputs = edits_to_dirty_inputs(&edits); + let result = + upload_ranges(config.clone(), cas_client.clone(), original_hash, original_size, inputs) + .await + .unwrap(); + let result_hash = MerkleHash::from_hex(result.hash()).unwrap(); + + assert_eq!( + result.file_size(), + Some(expected_next.len() as u64), + "worker={worker}, round={round}: size mismatch" + ); + let clean_hash = upload_file(&config, &expected_next).await; + assert_eq!(result.hash(), clean_hash.hex(), "worker={worker}, round={round}: hash mismatch"); + + let downloaded = download_file(&config, result_hash, expected_next.len() as u64).await; + assert_eq!(downloaded, expected_next, "worker={worker}, round={round}: content mismatch"); + + expected = expected_next; + original_hash = result_hash; + original_size = expected.len() as u64; + } + })); + } + + for handle in handles { + handle.await.unwrap(); + } + } +} diff --git a/xet_data/tests/test_session_resume.rs b/xet_data/tests/test_session_resume.rs index 4e8a5271b..d8f688c8f 100644 --- a/xet_data/tests/test_session_resume.rs +++ b/xet_data/tests/test_session_resume.rs @@ -145,7 +145,7 @@ mod tests { // Add all the data. Roughly the first half should dedup. cleaner.add_data(&data).await.unwrap(); - cleaner.finish().await.unwrap(); + let _ = cleaner.finish().await.unwrap(); let report = file_upload_session.report(); assert!(report.total_bytes > 0); diff --git a/xet_pkg/src/xet_session/file_download_group.rs b/xet_pkg/src/xet_session/file_download_group.rs index d78310197..c8f70d4d7 100644 --- a/xet_pkg/src/xet_session/file_download_group.rs +++ b/xet_pkg/src/xet_session/file_download_group.rs @@ -178,7 +178,7 @@ impl XetFileDownloadGroup { pub fn abort(&self) -> Result<(), XetError> { info!(group_id = %self.id(), "Download group abort"); self.task_runtime.cancel_subtree()?; - for (_tracking_id, handle) in self.inner.active_tasks.read()?.iter() { + for handle in self.inner.active_tasks.read()?.values() { handle.cancel(); } Ok(()) diff --git a/xet_pkg/src/xet_session/upload_commit.rs b/xet_pkg/src/xet_session/upload_commit.rs index b1e71fe6e..e6e5921d8 100644 --- a/xet_pkg/src/xet_session/upload_commit.rs +++ b/xet_pkg/src/xet_session/upload_commit.rs @@ -430,8 +430,6 @@ impl XetUploadCommit { .await } - /// Queue raw bytes for upload, starting the transfer immediately. - /// /// Returns a [`XetFileUpload`] whose /// [`finalize_ingestion`](XetFileUpload::finalize_ingestion) method yields /// per-file [`XetFileMetadata`] once ingestion completes. diff --git a/xet_runtime/src/file_utils/safe_file_creator.rs b/xet_runtime/src/file_utils/safe_file_creator.rs index 85cb81d90..3bb7b15c7 100644 --- a/xet_runtime/src/file_utils/safe_file_creator.rs +++ b/xet_runtime/src/file_utils/safe_file_creator.rs @@ -138,7 +138,7 @@ impl SafeFileCreator { Some(wr) => Ok(wr), None => Err(io::Error::new( io::ErrorKind::BrokenPipe, - format!("Writing to {:?} already completed.", &self.dest_path), + format!("Writing to {:?} already completed.", self.dest_path), )), } } @@ -163,7 +163,7 @@ impl Seek for SafeFileCreator { impl Drop for SafeFileCreator { fn drop(&mut self) { if let Err(e) = self.close() { - eprintln!("Error: Failed to close writer for {:?}: {}", &self.dest_path, e); + eprintln!("Error: Failed to close writer for {:?}: {}", self.dest_path, e); } } }