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.
This commit is contained in:
Omar Sobh
@@ -600,3 +600,79 @@ pub async fn push_blob_missing_chunks(
|
||||
}
|
||||
Ok((uploaded, total))
|
||||
}
|
||||
|
||||
/// Phase 5h: chunk-level cross-peer prewarm with bounded memory.
|
||||
///
|
||||
/// Unlike [`push_blob_missing_chunks`] this doesn't need a local
|
||||
/// `BlobStore` — it forwards chunk data directly from `upstream` to
|
||||
/// `downstream` one chunk at a time. Memory ceiling is one chunk
|
||||
/// (4 MiB by default), independent of blob size, so multi-GB blobs
|
||||
/// don't blow up RAM on the mediator.
|
||||
///
|
||||
/// Flow:
|
||||
/// 1. Load the manifest from `upstream`.
|
||||
/// 2. For each chunk, ask `downstream` if it already has it.
|
||||
/// 3. For missing chunks: `GetChunk` upstream → `PutChunk` downstream.
|
||||
/// Bytes are held only for the duration of the pair of calls.
|
||||
/// 4. Commit the manifest downstream via `PutManifest`. If the peer
|
||||
/// still reports missing chunks (concurrent eviction, storage
|
||||
/// failure), we retry each once; a second failure surfaces as
|
||||
/// `Err`.
|
||||
///
|
||||
/// Returns `(uploaded_chunks, total_chunks)` — the difference is the
|
||||
/// dedup save. When `downstream` is already fully warm the returned
|
||||
/// `uploaded` is 0 and no chunk bytes crossed the wire.
|
||||
///
|
||||
/// Note: this streams chunks sequentially. Fan-out (N concurrent
|
||||
/// chunk transfers) would speed multi-GB prewarms but complicates
|
||||
/// error handling; a future revision can layer parallelism on top
|
||||
/// without changing the memory story per in-flight chunk.
|
||||
pub async fn prewarm_missing_chunks_between(
|
||||
upstream: &Connection,
|
||||
downstream: &Connection,
|
||||
id: &BlobId,
|
||||
) -> Result<(usize, usize)> {
|
||||
let manifest = call_blob_load_manifest(upstream, id)
|
||||
.await?
|
||||
.with_context(|| format!("upstream missing blob {}", id.to_hex()))?;
|
||||
let total = manifest.chunks.len();
|
||||
let mut uploaded = 0usize;
|
||||
for hash in &manifest.chunks {
|
||||
if call_has_chunk(downstream, hash).await? {
|
||||
continue;
|
||||
}
|
||||
let bytes = call_get_chunk(upstream, hash)
|
||||
.await?
|
||||
.with_context(|| {
|
||||
format!(
|
||||
"upstream manifest referenced chunk {} but GetChunk returned NotFound",
|
||||
hash.to_hex()
|
||||
)
|
||||
})?;
|
||||
call_put_chunk(downstream, hash, &bytes).await?;
|
||||
uploaded += 1;
|
||||
// `bytes` is dropped here — memory ceiling is one chunk at a
|
||||
// time. Reassigned on the next iteration.
|
||||
}
|
||||
let still_missing = call_put_manifest(downstream, &manifest).await?;
|
||||
if !still_missing.is_empty() {
|
||||
// One retry pass: fetch + push each chunk the peer still lacks.
|
||||
// Guards against a narrow race where the peer evicted a chunk
|
||||
// between our HasChunk and its PutManifest verification pass.
|
||||
for hash in &still_missing {
|
||||
let bytes = call_get_chunk(upstream, hash).await?.with_context(|| {
|
||||
format!("retry fetch of chunk {} failed", hash.to_hex())
|
||||
})?;
|
||||
call_put_chunk(downstream, hash, &bytes).await?;
|
||||
uploaded += 1;
|
||||
}
|
||||
let final_missing = call_put_manifest(downstream, &manifest).await?;
|
||||
if !final_missing.is_empty() {
|
||||
bail!(
|
||||
"downstream still missing {} chunks after retry; storage may be failing",
|
||||
final_missing.len()
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok((uploaded, total))
|
||||
}
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user