Phase 5h: streaming chunk-level prewarm

Bounded-memory cross-peer prewarm: instead of buffering the whole blob in RAM (previous 5f path), iterate the upstream manifest chunk by chunk, ask downstream `HasChunk`, stream missing chunks one at a time. Memory ceiling is 1 chunk (4 MiB) regardless of blob size — a 5 GiB target dir no longer needs 5 GiB of mediator RAM. * rpc/client.rs: new `prewarm_missing_chunks_between(upstream, downstream, blob_id)` helper. Returns `(uploaded, total)` — the difference is the dedup save. Retries once if the downstream `PutManifest` reports missing chunks after our push (guards a narrow eviction race); a second failure surfaces as `Err`. * claw_cargo.rs: `prewarm` now streams by default; new `--buffered` flag for the old whole-blob path (kept for diagnostic comparability during rollout). Human-readable output shows the mode + dedup count. +2 tests: - cold downstream: 3-chunk payload (with a partial tail chunk) is copied exactly, reassembly is byte-equal to source - partial dedup: pre-seed 1 of 3 chunks on downstream → uploaded=2; rerun is a full no-op (uploaded=0), proving idempotence 251 tests pass (+2 from Phase 5j). Pre-existing macOS failure unchanged. Follow-ons: (1) parallel chunk transfer (uploaded chunks in fan-out) would speed multi-GB prewarms further; (2) exposing an accurate transferred-bytes counter needs router-side accounting instead of the current chunk-count × CHUNK_SIZE approximation.
2026-07-12 04:48:55 -07:00
parent 523b22f148
commit 6fb16286dd
3 changed files with 362 additions and 48 deletions
@@ -903,3 +903,205 @@ async fn stream_put_deduplicates_with_prior_put_bytes() {
    tokio::time::sleep(Duration::from_millis(50)).await;
    accept_task.abort();
 }
+
+#[tokio::test]
+async fn end_to_end_streaming_prewarm_copies_chunks_bounded_memory() {
+    // Phase 5h: `prewarm_missing_chunks_between` streams chunks one at
+    // a time from upstream to downstream without holding the whole
+    // blob in RAM. Two full RPC servers; seed a multi-chunk blob on A;
+    // stream it to C; verify downstream ended up with every chunk +
+    // manifest and can serve the assembled bytes back.
+    use crate::cluster::blob::CHUNK_SIZE;
+
+    let (id_a, id_b_for_a) = NodeIdentity::generate_test_pair("a", "b_a").unwrap();
+    let (id_c, id_b_for_c) = NodeIdentity::generate_test_pair("c", "b_c").unwrap();
+
+    let (_tmp_a, router_a) = router_with_full_stack("a", next_port()).await;
+    let (_tmp_c, router_c) = router_with_full_stack("c", next_port()).await;
+
+    // 3 chunks worth so we're not testing a single-chunk edge case.
+    // The last chunk is intentionally short (not a full CHUNK_SIZE) so
+    // we also cover the tail-chunk path.
+    let payload: Vec<u8> = (0..(2 * CHUNK_SIZE + CHUNK_SIZE / 4))
+        .map(|i| ((i * 31) % 251) as u8)
+        .collect();
+    let blob_id = router_a
+        .blob_store()
+        .unwrap()
+        .put_bytes(&payload)
+        .await
+        .unwrap();
+    let up_manifest = router_a
+        .blob_store()
+        .unwrap()
+        .load_manifest(&blob_id)
+        .await
+        .unwrap()
+        .expect("manifest on A");
+    assert_eq!(up_manifest.chunks.len(), 3, "expected 3 chunks in test payload");
+
+    let server_a = QuicServer::bind(loopback(0), id_a).unwrap();
+    let server_a_addr = server_a.local_addr().unwrap();
+    let router_a_srv = router_a.clone();
+    let accept_a = tokio::spawn(async move {
+        while let Some(Ok(conn)) = server_a.accept().await {
+            let r = router_a_srv.clone();
+            tokio::spawn(async move {
+                let _ = serve_connection(conn, r).await;
+            });
+        }
+    });
+
+    let server_c = QuicServer::bind(loopback(0), id_c).unwrap();
+    let server_c_addr = server_c.local_addr().unwrap();
+    let router_c_srv = router_c.clone();
+    let accept_c = tokio::spawn(async move {
+        while let Some(Ok(conn)) = server_c.accept().await {
+            let r = router_c_srv.clone();
+            tokio::spawn(async move {
+                let _ = serve_connection(conn, r).await;
+            });
+        }
+    });
+
+    // Mediator client — one QUIC client per side (distinct identities
+    // are needed since each test pair generates its own CA).
+    let up_client = QuicClient::new(loopback(0), id_b_for_a).unwrap();
+    let up_conn = up_client.connect(server_a_addr, "a").await.unwrap();
+    let down_client = QuicClient::new(loopback(0), id_b_for_c).unwrap();
+    let down_conn = down_client.connect(server_c_addr, "c").await.unwrap();
+
+    let (uploaded, total) =
+        prewarm_missing_chunks_between(&up_conn, &down_conn, &blob_id)
+            .await
+            .unwrap();
+    assert_eq!(total, 3);
+    assert_eq!(uploaded, 3, "cold downstream must receive every chunk");
+
+    // Downstream reassembly proves the manifest committed AND all
+    // chunks landed correctly.
+    let round = router_c
+        .blob_store()
+        .unwrap()
+        .get_bytes(&blob_id)
+        .await
+        .unwrap();
+    assert_eq!(round.as_deref(), Some(payload.as_slice()));
+
+    up_conn.close(quinn::VarInt::from_u32(0), b"done");
+    down_conn.close(quinn::VarInt::from_u32(0), b"done");
+    up_client.shutdown().await;
+    down_client.shutdown().await;
+    tokio::time::sleep(Duration::from_millis(50)).await;
+    accept_a.abort();
+    accept_c.abort();
+}
+
+#[tokio::test]
+async fn streaming_prewarm_skips_chunks_already_present_downstream() {
+    // Phase 5h: verify the dedup path. Pre-seed a subset of the
+    // upstream blob's chunks on downstream; run the streaming prewarm;
+    // `uploaded` must be less than `total` by exactly the pre-seeded
+    // count.
+    use crate::cluster::blob::CHUNK_SIZE;
+
+    let (id_a, id_b_for_a) = NodeIdentity::generate_test_pair("a", "b_a").unwrap();
+    let (id_c, id_b_for_c) = NodeIdentity::generate_test_pair("c", "b_c").unwrap();
+
+    let (_tmp_a, router_a) = router_with_full_stack("a", next_port()).await;
+    let (_tmp_c, router_c) = router_with_full_stack("c", next_port()).await;
+
+    let payload: Vec<u8> = (0..(3 * CHUNK_SIZE)).map(|i| ((i * 17) % 251) as u8).collect();
+    let blob_id = router_a
+        .blob_store()
+        .unwrap()
+        .put_bytes(&payload)
+        .await
+        .unwrap();
+    let up_manifest = router_a
+        .blob_store()
+        .unwrap()
+        .load_manifest(&blob_id)
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(up_manifest.chunks.len(), 3);
+
+    // Pre-seed the first chunk on downstream so it's genuinely already
+    // present. Reads it from A's store to guarantee identical bytes.
+    let first_hash = up_manifest.chunks[0];
+    let first_bytes = router_a
+        .blob_store()
+        .unwrap()
+        .read_chunk(&first_hash)
+        .await
+        .unwrap()
+        .unwrap();
+    router_c
+        .blob_store()
+        .unwrap()
+        .put_chunk(&first_hash, &first_bytes)
+        .await
+        .unwrap();
+
+    let server_a = QuicServer::bind(loopback(0), id_a).unwrap();
+    let addr_a = server_a.local_addr().unwrap();
+    let ra = router_a.clone();
+    let acc_a = tokio::spawn(async move {
+        while let Some(Ok(conn)) = server_a.accept().await {
+            let r = ra.clone();
+            tokio::spawn(async move {
+                let _ = serve_connection(conn, r).await;
+            });
+        }
+    });
+    let server_c = QuicServer::bind(loopback(0), id_c).unwrap();
+    let addr_c = server_c.local_addr().unwrap();
+    let rc = router_c.clone();
+    let acc_c = tokio::spawn(async move {
+        while let Some(Ok(conn)) = server_c.accept().await {
+            let r = rc.clone();
+            tokio::spawn(async move {
+                let _ = serve_connection(conn, r).await;
+            });
+        }
+    });
+
+    let up_client = QuicClient::new(loopback(0), id_b_for_a).unwrap();
+    let up_conn = up_client.connect(addr_a, "a").await.unwrap();
+    let down_client = QuicClient::new(loopback(0), id_b_for_c).unwrap();
+    let down_conn = down_client.connect(addr_c, "c").await.unwrap();
+
+    let (uploaded, total) =
+        prewarm_missing_chunks_between(&up_conn, &down_conn, &blob_id)
+            .await
+            .unwrap();
+    assert_eq!(total, 3);
+    assert_eq!(
+        uploaded, 2,
+        "dedup should skip the pre-seeded first chunk (uploaded={uploaded})"
+    );
+
+    // Idempotent: rerun should upload zero chunks (all present now).
+    let (uploaded2, _) =
+        prewarm_missing_chunks_between(&up_conn, &down_conn, &blob_id)
+            .await
+            .unwrap();
+    assert_eq!(uploaded2, 0, "second prewarm should be a full-dedup no-op");
+
+    let round = router_c
+        .blob_store()
+        .unwrap()
+        .get_bytes(&blob_id)
+        .await
+        .unwrap();
+    assert_eq!(round.as_deref(), Some(payload.as_slice()));
+
+    up_conn.close(quinn::VarInt::from_u32(0), b"done");
+    down_conn.close(quinn::VarInt::from_u32(0), b"done");
+    up_client.shutdown().await;
+    down_client.shutdown().await;
+    tokio::time::sleep(Duration::from_millis(50)).await;
+    acc_a.abort();
+    acc_c.abort();
+}