fix(push): upload only unresolved prepared chunks

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.

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::<crab_staging::push_plan::PlannedXorb>::new();
-    let mut candidate_indices = HashMap::<MerkleHash, usize>::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::<MerkleHash, HashMap<MerkleHash, XorbRef>>::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::<HashSet<_>>();
+    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,22 +397,26 @@ pub(crate) async fn prepare_delta(
                 .iter()
                 .map(crab_staging::push_plan::PlannedPlacement::to_placement)
                 .collect::<crab_staging::Result<Vec<_>>>()?;
-            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(
                 planned_placements
                     .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::<HashSet<_>>();
+    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::<HashMap<_, _>>();
-                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::<Vec<_>>();
         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,

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<dyn CompressionPolicy>)
+            })),
+            ..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<dyn CompressionPolicy>);
+        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::<HashSet<_>>();
+        let actual = delta
+            .xorbs()
+            .values()
+            .flat_map(|entry| entry.chunks())
+            .map(|chunk| chunk.hash().to_owned())
+            .collect::<HashSet<_>>();
+        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"
+        );
+    }
+}