Phase 5h: streaming chunk-level prewarm #19

Merged
osobh merged 1 commits from phase-5h-streaming-prewarm-v2 into main 2026-07-12 12:01:03 +00:00
3 changed files with 362 additions and 48 deletions
+84 -48
View File
@@ -40,13 +40,13 @@ use std::path::{Path, PathBuf};
use std::process::Command; use std::process::Command;
use std::time::Instant; use std::time::Instant;
use crate::cluster::blob::BlobId; use crate::cluster::blob::{BlobId, CHUNK_SIZE};
use crate::cluster::build_cache::{ use crate::cluster::build_cache::{
capture_target, compute_workspace_fingerprint, restore_target, Fingerprint, capture_target, compute_workspace_fingerprint, restore_target, Fingerprint,
}; };
use crate::cluster::client_config::{ClientConfig, ResolvedClientConfig}; use crate::cluster::client_config::{ClientConfig, ResolvedClientConfig};
use crate::cluster::rpc::{ use crate::cluster::rpc::{
call_blob_get_stream, call_blob_put_stream, call_blob_stat, call_delete_tag, call_get_metrics, call_blob_get_stream, call_blob_put_stream, call_blob_stat, call_delete_tag, call_get_metrics, prewarm_missing_chunks_between,
call_get_ref, call_get_tag, call_list_tags, call_put_ref, call_put_tag, call_get_ref, call_get_tag, call_list_tags, call_put_ref, call_put_tag,
}; };
use crate::cluster::transport::{NodeIdentity, QuicClient}; use crate::cluster::transport::{NodeIdentity, QuicClient};
@@ -123,6 +123,12 @@ struct PrewarmArgs {
/// fingerprint to prewarm. /// fingerprint to prewarm.
#[arg(long)] #[arg(long)]
pin: String, pin: String,
/// Phase 5h: buffer the entire blob in RAM instead of streaming
/// chunk-by-chunk. Slower + memory-heavy but useful as a diagnostic
/// when the streaming path misbehaves. Default is streaming (one
/// 4 MiB chunk in RAM at a time).
#[arg(long, default_value_t = false)]
buffered: bool,
} }
#[derive(clap::Args, Debug, Clone)] #[derive(clap::Args, Debug, Clone)]
@@ -797,30 +803,8 @@ async fn cmd_prewarm(args: PrewarmArgs) -> Result<()> {
stat.chunk_count stat.chunk_count
); );
// Download the whole blob. Buffered — the streaming zero-copy // Downstream client — separate QUIC endpoint so we don't conflate
// upstream→downstream pipe is a follow-on that avoids the RAM // per-side connection state.
// ceiling (~16 MiB currently) but complicates task ownership.
let download_started = Instant::now();
let mut buf: Vec<u8> = Vec::with_capacity(stat.total_size as usize);
let ok = call_blob_get_stream(&up_conn, &blob_id, &mut buf).await?;
if !ok {
up_conn.close(quinn::VarInt::from_u32(0), b"done");
up_client.shutdown().await;
anyhow::bail!(
"upstream reported tag {} → blob {} but BlobGetStream returned NotFound",
args.pin,
blob_id
);
}
let download_elapsed = download_started.elapsed();
tracing::info!(
"prewarm: downloaded {} bytes from upstream in {:?}",
buf.len(),
download_elapsed
);
// Second client for the downstream side. Distinct QUIC endpoint
// so we don't conflate connection state.
let down_identity = NodeIdentity::from_pem_dir(&args.tls_dir)?; let down_identity = NodeIdentity::from_pem_dir(&args.tls_dir)?;
let down_client = QuicClient::new("0.0.0.0:0".parse()?, down_identity)?; let down_client = QuicClient::new("0.0.0.0:0".parse()?, down_identity)?;
let down_conn = down_client let down_conn = down_client
@@ -828,26 +812,69 @@ async fn cmd_prewarm(args: PrewarmArgs) -> Result<()> {
.await .await
.with_context(|| format!("connecting to downstream {} @ {}", args.to_peer, args.to_addr))?; .with_context(|| format!("connecting to downstream {} @ {}", args.to_peer, args.to_addr))?;
let upload_started = Instant::now(); // Chunk-streaming (default) vs whole-blob buffering (--buffered).
let cursor = std::io::Cursor::new(buf.clone()); // Streaming reads one 4 MiB chunk into RAM at a time; buffered
let assigned_id = call_blob_put_stream(&down_conn, cursor).await?; // reads the whole blob before pushing (useful as a diagnostic
if assigned_id != blob_id { // baseline).
down_conn.close(quinn::VarInt::from_u32(0), b"done"); let transfer_started = Instant::now();
up_conn.close(quinn::VarInt::from_u32(0), b"done"); let (uploaded_chunks, transfer_bytes, mode_label) = if args.buffered {
down_client.shutdown().await; let mut buf: Vec<u8> = Vec::with_capacity(stat.total_size as usize);
up_client.shutdown().await; let ok = call_blob_get_stream(&up_conn, &blob_id, &mut buf).await?;
anyhow::bail!( if !ok {
"prewarm integrity check failed: downstream stored {} but upstream had {}", up_conn.close(quinn::VarInt::from_u32(0), b"done");
assigned_id, down_conn.close(quinn::VarInt::from_u32(0), b"done");
blob_id up_client.shutdown().await;
down_client.shutdown().await;
anyhow::bail!(
"upstream reported tag {} → blob {} but BlobGetStream returned NotFound",
args.pin,
blob_id
);
}
let bytes_len = buf.len() as u64;
let cursor = std::io::Cursor::new(buf);
let assigned_id = call_blob_put_stream(&down_conn, cursor).await?;
if assigned_id != blob_id {
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;
anyhow::bail!(
"prewarm integrity check failed: downstream stored {} but upstream had {}",
assigned_id,
blob_id
);
}
// Buffered mode has no chunk dedup on the wire — every byte
// crosses. Report full chunk count as "uploaded" for parity
// with the streaming path's semantics.
(stat.chunk_count as usize, bytes_len, "buffered")
} else {
// Phase 5h: chunk-by-chunk streaming. Memory ceiling = one
// chunk (4 MiB). Also gets us free dedup — chunks already
// present downstream (from a prior warm build) are skipped.
let (uploaded, total) =
prewarm_missing_chunks_between(&up_conn, &down_conn, &blob_id)
.await
.with_context(|| {
format!("streaming prewarm of blob {}", blob_id.to_hex())
})?;
// Bytes actually transferred = uploaded chunks × chunk size,
// capped at total_size for the tail chunk. Rough estimate; a
// true byte count would require the router to report actual
// bytes served, which is future 5h+ work.
let approx_bytes = (uploaded as u64) * CHUNK_SIZE as u64;
let approx_bytes = approx_bytes.min(stat.total_size);
tracing::info!(
"prewarm: streamed {}/{} chunks ({} dedup skipped) in {:?}",
uploaded,
total,
total.saturating_sub(uploaded),
transfer_started.elapsed()
); );
} (uploaded, approx_bytes, "streaming")
let upload_elapsed = upload_started.elapsed(); };
tracing::info!( let transfer_elapsed = transfer_started.elapsed();
"prewarm: uploaded {} bytes to downstream in {:?}",
buf.len(),
upload_elapsed
);
// Publish the tag downstream too — otherwise a `prefetch --pin` // Publish the tag downstream too — otherwise a `prefetch --pin`
// there would still miss. // there would still miss.
@@ -868,9 +895,18 @@ async fn cmd_prewarm(args: PrewarmArgs) -> Result<()> {
"bytes: {} ({} chunks)", "bytes: {} ({} chunks)",
stat.total_size, stat.chunk_count stat.total_size, stat.chunk_count
); );
println!("download: {:?}", download_elapsed); println!("mode: {}", mode_label);
println!("upload: {:?}", upload_elapsed); println!(
println!("total: {:?}", download_elapsed + upload_elapsed); "transferred: {} ({} of {} chunks)",
human_bytes(transfer_bytes),
uploaded_chunks,
stat.chunk_count
);
println!(
"dedup save: {} chunks",
(stat.chunk_count as usize).saturating_sub(uploaded_chunks)
);
println!("elapsed: {:?}", transfer_elapsed);
println!("────────────────────────────────────────────────────"); println!("────────────────────────────────────────────────────");
Ok(()) Ok(())
} }
+76
View File
@@ -600,3 +600,79 @@ pub async fn push_blob_missing_chunks(
} }
Ok((uploaded, total)) 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))
}
+202
View File
@@ -903,3 +903,205 @@ async fn stream_put_deduplicates_with_prior_put_bytes() {
tokio::time::sleep(Duration::from_millis(50)).await; tokio::time::sleep(Duration::from_millis(50)).await;
accept_task.abort(); 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();
}