diff --git a/crab/docs/design/add.md b/crab/docs/design/add.md index 16fbe9e6e..fd06d91bb 100644 --- a/crab/docs/design/add.md +++ b/crab/docs/design/add.md @@ -298,6 +298,14 @@ closed and asks the user to run `crab add` again. Recipe leases and push snapshots retain shared bodies until the final owner is released; restaging reclaims only globally unleased payloads. +Capsule publication adopts a complete prepared xorb only when every chunk is +selected for the push and still unresolved. Mixed sources are repacked using +only missing chunks, retaining committed placements for reused bytes. The shared +staging read plan groups residual requests across files by source, verifies each +prepared body once, and yields bounded decoded batches. Compression runs off the +async worker. Ordered recipe occurrences still produce complete file terms; +local recipe ownership and flush-before-pointer publication are unchanged. + There is no persisted per-file JSON plan or per-file payload copy. Runtime `FilePushPlan` values are derived from normalized recipe, remote, prepared, and segment rows for the push attempt. diff --git a/crab/src/git/xet_publication.rs b/crab/src/git/xet_publication.rs index 266fa1d6f..42d2b3637 100644 --- a/crab/src/git/xet_publication.rs +++ b/crab/src/git/xet_publication.rs @@ -140,7 +140,6 @@ pub(crate) async fn prepare_delta( let mut xorb_bytes_uploaded = 0_u64; let mut prepared_files = Vec::with_capacity(unresolved.len()); let mut prepared_candidates = Vec::::new(); - let mut candidate_indices = HashMap::::new(); let mut staged_candidates = HashMap::new(); for (file_hash, file_size) in unresolved { @@ -311,7 +310,6 @@ pub(crate) async fn prepare_delta( xorb_entries.insert(hash, entry); origins.insert(hash, origin); } - drop(requested); let mut external_candidates = HashMap::>::new(); for (hash, candidate) in verified_staged.candidates { let xorb_ref = candidate.xorb_ref; @@ -378,13 +376,11 @@ pub(crate) async fn prepare_delta( } let mut candidate_chunks = placements.keys().copied().collect::>(); + let mut seen_candidates = HashMap::new(); for (_, _, _, plan) in &prepared_files { for planned in &plan.prepared_xorbs { let planned_hash = planned.hash()?; - if let Some(index) = candidate_indices.get(&planned_hash).copied() { - let existing = prepared_candidates.get(index).ok_or_else(|| { - CrabError::Internal("prepared xorb candidate index is invalid".to_owned()) - })?; + if let Some(existing) = seen_candidates.insert(planned_hash, planned) { if !planned_xorbs_match(existing, planned) { return Err(CrabError::StagingCorrupt(format!( "prepared xorb {} has conflicting plans", @@ -401,10 +397,13 @@ pub(crate) async fn prepare_delta( .iter() .map(crab_staging::push_plan::PlannedPlacement::to_placement) .collect::>>()?; - if planned_placements - .iter() - .all(|placement| candidate_chunks.contains(&placement.chunk_hash)) - { + // Direct adoption is only for wholly new outgoing payloads. A + // single missing chunk must not upload an otherwise reused xorb; + // residual packing below reads only the still-needed chunks. + if !planned_placements.iter().all(|placement| { + requested.contains(placement.chunk_hash.hex().as_str()) + && !candidate_chunks.contains(&placement.chunk_hash) + }) { continue; } candidate_chunks.extend( @@ -412,11 +411,12 @@ pub(crate) async fn prepare_delta( .iter() .map(|placement| placement.chunk_hash), ); - candidate_indices.insert(planned_hash, prepared_candidates.len()); prepared_candidates.push(planned.clone()); } } + drop(requested); + let uploads = futures_util::stream::iter(prepared_candidates.into_iter().enumerate().map( |(ordinal, planned)| async move { check_cancelled(cancel)?; @@ -484,51 +484,55 @@ pub(crate) async fn prepare_delta( xorb_entries.insert(planned_hash, entry); } - for (file_ordinal, (file_hash, file_size, chunks, _)) in prepared_files.into_iter().enumerate() - { - check_cancelled(cancel)?; - let mut missing = Vec::new(); - let mut seen = HashSet::new(); - for (hash, _) in &chunks { - if !placements.contains_key(hash) && seen.insert(*hash) { - missing.push(*hash); + let mut missing = prepared_files + .iter() + .flat_map(|(_, _, chunks, _)| chunks.iter().map(|(hash, _)| *hash)) + .filter(|hash| !placements.contains_key(hash)) + .collect::>(); + if !missing.is_empty() { + // The same prepared body can supply several files. Gather the whole + // residual schedule before reading, so its verified source opens once + // rather than once per file or fixed-size request page. + let mut read_plan = staging.begin_coalesced_chunk_read_plan()?; + let mut requests = Vec::with_capacity(256); + for (ordinal, (_, _, chunks, _)) in prepared_files.iter().enumerate() { + let context = u64::try_from(ordinal) + .map_err(|_| CrabError::Internal("pointer file ordinal overflowed".to_owned()))?; + for (hash, size) in chunks { + if !missing.remove(hash) { + continue; + } + requests.push((*hash, *size, context)); + if requests.len() == 256 { + check_cancelled(cancel)?; + read_plan.append(&requests)?; + requests.clear(); + } } } - if !missing.is_empty() { - let mut builder = XorbBuilder::new(); - for batch in missing.chunks(256) { - let loaded = staging.get_chunks_batch(batch).await?; - let by_hash = loaded.into_iter().collect::>(); - for hash in batch { - let data = by_hash - .get(hash) - .cloned() - .ok_or_else(|| CrabError::ChunkNotFound { hash: hash.hex() })?; - builder.push( - &Chunk { hash: *hash, data }, - RunId(u64::try_from(file_ordinal).map_err(|_| { - CrabError::Internal("pointer file ordinal overflowed".to_owned()) - })?), - )?; + read_plan.append(&requests)?; + drop(missing); + let mut builder = XorbBuilder::new(); + while let Some(payloads) = read_plan.read_next_bounded(256, 16 * 1024 * 1024).await? { + check_cancelled(cancel)?; + // Compression runs on a bounded decoded batch so ref-lease and + // upload tasks remain schedulable while residual payloads are packed. + let packing_cancel = cancel.clone(); + let (next_builder, completed) = tokio::task::spawn_blocking(move || { + let mut completed = Vec::new(); + for (context, hash, data) in payloads { + check_cancelled(&packing_cancel)?; + builder.push(&Chunk { hash, data }, RunId(context))?; while let Some(result) = builder.take_completed() { - let uploaded_bytes = publish_built_xorb( - layout, - caching_store, - result, - &mut placements, - &mut xorb_entries, - &mut uploaded_xorbs, - &mut origins, - ) - .await?; - xorb_bytes_uploaded = xorb_bytes_uploaded.saturating_add(uploaded_bytes); - if let Some(metrics) = metrics { - metrics.add_bytes_uploaded(uploaded_bytes); - } + completed.push(result); } } - } - for result in builder.finalize()? { + Ok::<_, CrabError>((builder, completed)) + }) + .await + .map_err(std::io::Error::other)??; + builder = next_builder; + for result in completed { let uploaded_bytes = publish_built_xorb( layout, caching_store, @@ -545,7 +549,29 @@ pub(crate) async fn prepare_delta( } } } + let finalized = tokio::task::spawn_blocking(move || builder.finalize()) + .await + .map_err(std::io::Error::other)??; + for result in finalized { + let uploaded_bytes = publish_built_xorb( + layout, + caching_store, + result, + &mut placements, + &mut xorb_entries, + &mut uploaded_xorbs, + &mut origins, + ) + .await?; + xorb_bytes_uploaded = xorb_bytes_uploaded.saturating_add(uploaded_bytes); + if let Some(metrics) = metrics { + metrics.add_bytes_uploaded(uploaded_bytes); + } + } + } + for (file_hash, file_size, chunks, _) in prepared_files { + check_cancelled(cancel)?; let ordered_hashes = chunks.iter().map(|(hash, _)| *hash).collect::>(); let file_info = file_info_from_placements(file_hash, &ordered_hashes, &placements)?; let dependency_hashes = file_info @@ -1486,6 +1512,8 @@ mod tests { } } + mod residual_tests; + #[derive(Default)] struct UploadObserver { active: std::sync::atomic::AtomicUsize, diff --git a/crab/src/git/xet_publication/tests/residual_tests.rs b/crab/src/git/xet_publication/tests/residual_tests.rs new file mode 100644 index 000000000..c5286fe77 --- /dev/null +++ b/crab/src/git/xet_publication/tests/residual_tests.rs @@ -0,0 +1,160 @@ +use super::*; + +#[tokio::test(flavor = "multi_thread")] +async fn mixed_prepared_xorb_uploads_only_unresolved_chunks() { + use crab_staging::stream::{ + StreamStageProgress, StreamStageXorbBuilder, stage_reader_streaming, + }; + use std::sync::atomic::{AtomicU64, Ordering}; + + #[derive(Default)] + struct Reads(AtomicU64); + impl crab_staging::StagingMetrics for Reads { + fn inc_prepared_source_xorb_open(&self) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + let root = tempfile::tempdir().unwrap(); + let staging = crab_staging::StagingArea::open(root.path().to_owned()) + .await + .unwrap(); + let mut bytes = vec![0; 32 * 1024 * 1024]; + blake3::Hasher::new() + .update(b"mixed prepared publication") + .finalize_xof() + .fill(&mut bytes); + let staged = stage_reader_streaming( + &mut std::io::Cursor::new(&bytes), + Path::new("model.bin"), + &staging, + StreamStageProgress { + xorb_builder: Some(StreamStageXorbBuilder::new(1, || { + XorbBuilder::with_policy(Arc::new(FixedCompression::new(CompressionScheme::None)) + as Arc) + })), + ..Default::default() + }, + &CancellationToken::new(), + ) + .await + .unwrap(); + staging + .delegate_publication(&[(&staged.batch_id, Path::new("model.bin"))]) + .unwrap(); + staging.close().await.unwrap(); + let mut staging = StagingAreaReadOnly::open(root.path().to_owned()) + .await + .unwrap(); + let reads = Arc::new(Reads::default()); + staging.set_metrics(reads.clone()); + let staging = Arc::new(staging); + let plans = staging + .prepared_xorbs_for_recipes(&[MerkleHash::from(staged.recipe.hash())]) + .unwrap(); + let plans = &plans[&MerkleHash::from(staged.recipe.hash())]; + assert_eq!(plans.len(), 1); + let original = XorbParser::parse(Bytes::from( + std::fs::read(crab_staging::push_plan::prepared_xorb_path( + staging.root(), + &plans[0].hash().unwrap(), + )) + .unwrap(), + )) + .unwrap(); + let count = plans[0].placements.len(); + assert!( + count > 257, + "fixture must span the old 256-chunk read batches" + ); + // Exercise both a small edit and an almost entirely new recipe. The + // latter must not reopen the same prepared body for every read batch. + for committed_count in [count - 1, 1] { + let mut builder = + XorbBuilder::with_policy(Arc::new(FixedCompression::new(CompressionScheme::None)) + as Arc); + for index in 0..committed_count { + builder + .push(&original.get_chunk(index as u32).unwrap(), RunId(0)) + .unwrap(); + } + let mut results = builder.finalize().unwrap(); + assert_eq!(results.len(), 1); + let committed = results.remove(0); + let entry = catalog_xorb_entry( + committed.bytes.len() as u64, + blake3::hash(&committed.bytes).to_hex().to_string(), + &committed.placements, + ) + .unwrap(); + let mut catalog = PointerCatalog::new(); + catalog.insert_xorb(committed.hash.hex(), entry).unwrap(); + // Only the selected new shard is reconstructed below. This old + // binding supplies the authenticated reverse-chunk lookup fixture. + let shard = "a".repeat(64); + catalog + .insert_shard( + shard.clone(), + ShardCatalogEntry::new(1, vec![committed.hash.hex()]), + ) + .unwrap(); + catalog + .insert_file("b".repeat(64), FileCatalogEntry::new(1, shard)) + .unwrap(); + let (layout, base) = publish_catalog(catalog.encode().unwrap()).await; + layout + .store() + .create_strict(&layout.xorb_path(&committed.hash), committed.bytes) + .await + .unwrap(); + reads.0.store(0, Ordering::Relaxed); + let pointer = Pointer { + file_hash: staged.file_hash, + size: staged.size, + shard_hint: None, + }; + let (delta, stats) = prepare_delta( + &layout, + &base, + std::slice::from_ref(&pointer), + Some(&staging), + None, + None, + &Default::default(), + 1, + false, + &CancellationToken::new(), + ) + .await + .unwrap(); + crab_read::verify_catalog_file_recipe( + &layout, + &pointer, + &delta.files()[&MerkleHash::from(staged.file_hash).hex()], + &CancellationToken::new(), + ) + .await + .unwrap(); + let expected = plans[0] + .placements + .iter() + .skip(committed_count) + .map(|placement| placement.chunk_hash.clone()) + .collect::>(); + let actual = delta + .xorbs() + .values() + .flat_map(|entry| entry.chunks()) + .map(|chunk| chunk.hash().to_owned()) + .collect::>(); + eprintln!( + "committed={committed_count}/{count}, uploaded_bytes={}, prepared_reads={}", + stats.xorb_bytes_uploaded, + reads.0.load(Ordering::Relaxed) + ); + assert_eq!( + (actual, reads.0.load(Ordering::Relaxed)), + (expected, 1), + "publication must upload only missing chunks and open their source once" + ); + } +}