From 371e52cb48bae99232547f6fae6a0487a85c2d04 Mon Sep 17 00:00:00 2001 From: tv Date: Sat, 24 Oct 2026 16:18:58 +0200 Subject: add chunked download and use it during rejoin --- src/swarm/blob.rs | 491 ++++++++++++++++++++++++++++++++++++++++++++++++++++ src/swarm/mod.rs | 1 + src/swarm/peer.rs | 3 + src/swarm/source.rs | 2 +- src/swarm/sync.rs | 6 + 5 files changed, 502 insertions(+), 1 deletion(-) create mode 100644 src/swarm/blob.rs (limited to 'src/swarm') diff --git a/src/swarm/blob.rs b/src/swarm/blob.rs new file mode 100644 index 0000000..6ff1618 --- /dev/null +++ b/src/swarm/blob.rs @@ -0,0 +1,491 @@ +//! Blobs: content-addressed bytes fetched in chunks from several peers at once. +//! +//! Large transfers can't go in one Mycelium message (it must arrive within Mycelium's +//! five-minute deadline), and a blob every node holds identically can come from many +//! of them. Today the only kind of blob is an Automerge change by its hash, e.g. an +//! epoch's snapshot change, which every node on that epoch holds byte for byte. +//! +//! * Serving: a request without chunks gets the manifest (size, chunk size, a BLAKE3 +//! hash per chunk); one with chunks gets each chunk as a message of its own. A node +//! that doesn't hold the blob answers `Missing`. The bytes of the blob served last +//! are kept, so serving a download doesn't extract them for every chunk. Serving runs +//! in tasks of its own, at most [`MAX_SERVING`] chunk sends at a time. +//! * Fetching: the manifest from the first peer that has it (closest first), then the +//! chunks, a few in flight per peer from up to [`MAX_SOURCES`] peers. Every chunk is +//! checked against the manifest; a peer that sends a bad one is dropped, a chunk +//! that doesn't arrive in time is asked of another peer. The caller checks the whole +//! (for a change: its hash). +//! +//! The chunk size is the serving node's (`--blob-chunk-kib`); the manifest carries it. + +use std::collections::{HashMap, HashSet, VecDeque}; +use std::sync::{Arc, Mutex, MutexGuard}; +use std::time::Duration; + +use anyhow::{Result, anyhow, bail}; +use automerge::ChangeHash; +use tokio::sync::{Semaphore, mpsc}; +use tokio::time::Instant; +use tracing::{debug, info, warn}; + +use crate::state::AppState; +use crate::transport::{BlobBody, BlobId, BlobMessage, ParsedMessage}; + +/// Chunk sends in progress at once, over all downloads served. +const MAX_SERVING: usize = 8; +/// Peers a download fetches chunks from at once. +const MAX_SOURCES: usize = 4; +/// Chunks in flight per peer. +const WINDOW: usize = 4; +/// Largest blob a download accepts. +const MAX_SIZE: u64 = 4 << 30; +/// Largest chunk a download accepts. +const MAX_CHUNK: u32 = 4 << 20; +/// How long to wait for a manifest from one peer. +const MANIFEST_TIMEOUT: Duration = Duration::from_secs(15); +/// How long to wait for a chunk before asking another peer (raised by the peer's +/// round-trip time). +const CHUNK_TIMEOUT: Duration = Duration::from_secs(60); +/// A download gives up after this long without any chunk arriving. +const STALL_TIMEOUT: Duration = Duration::from_secs(180); +/// Bad chunks after which a peer is no longer asked. +const MAX_BAD: u32 = 2; + +/// Per node: what is being served and fetched. +pub(crate) struct Blobs { + chunk_size: u32, + serving: Arc, + /// The blob served last. + cache: Mutex)>>, + /// Answers for downloads in progress, by blob. + fetches: Mutex>>, +} + +struct Prepared { + bytes: Vec, + hashes: Vec<[u8; 32]>, +} + +impl Blobs { + pub(crate) fn new(chunk_size: u32) -> Self { + Self { + chunk_size: chunk_size.max(1), + serving: Arc::new(Semaphore::new(MAX_SERVING)), + cache: Mutex::new(None), + fetches: Mutex::new(HashMap::new()), + } + } + + fn fetches( + &self, + ) -> MutexGuard<'_, HashMap>> { + self.fetches + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } + + fn cached(&self, id: BlobId) -> Option> { + self.cache + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .as_ref() + .filter(|(cached, _)| *cached == id) + .map(|(_, prepared)| Arc::clone(prepared)) + } + + fn remember(&self, id: BlobId, prepared: Arc) { + *self + .cache + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some((id, prepared)); + } +} + +fn chunk_hashes(bytes: &[u8], chunk_size: u32) -> Vec<[u8; 32]> { + bytes + .chunks(chunk_size as usize) + .map(|chunk| *blake3::hash(chunk).as_bytes()) + .collect() +} + +fn chunk_count(size: u64, chunk_size: u32) -> u64 { + size.div_ceil(u64::from(chunk_size)) +} + +/// A blob message from `peer`. +pub(crate) async fn handle(state: &AppState, peer: &str, message: BlobMessage) { + match message.body { + BlobBody::Request { chunks } => { + let state = state.clone(); + let peer = peer.to_owned(); + // Off the inbound loop: answering may send many large messages. + tokio::spawn(async move { serve(&state, &peer, message.id, chunks).await }); + } + body => { + let sender = state.blobs.fetches().get(&message.id).cloned(); + match sender { + Some(sender) => { + let _ = sender.send((peer.to_owned(), body)); + } + None => { + debug!(peer = %peer, blob = %message.id, "Blob answer for no download in progress") + } + } + } + } +} + +async fn prepare(state: &AppState, id: BlobId) -> Option> { + if let Some(prepared) = state.blobs.cached(id) { + return Some(prepared); + } + let bytes = match id { + BlobId::Change(hash) => { + let doc = Arc::clone(&state.doc); + tokio::task::spawn_blocking(move || { + // get_change_by_hash needs the document mutably. + let mut replica = doc.blocking_write(); + replica + .doc + .get_change_by_hash(&hash) + .map(|change| change.raw_bytes().to_vec()) + }) + .await + .ok() + .flatten()? + } + }; + let chunk_size = state.blobs.chunk_size; + let prepared = tokio::task::spawn_blocking(move || { + let hashes = chunk_hashes(&bytes, chunk_size); + Prepared { bytes, hashes } + }) + .await + .ok()?; + let prepared = Arc::new(prepared); + state.blobs.remember(id, Arc::clone(&prepared)); + Some(prepared) +} + +async fn reply(state: &AppState, peer: &str, id: BlobId, body: BlobBody) { + let message = ParsedMessage::Blob(BlobMessage { + epoch: state.local_epoch(), + id, + body, + }); + if let Err(error) = state.transport.send_message(peer, &message).await { + debug!(peer = %peer, blob = %id, error = %format!("{error:#}"), "Blob answer failed"); + } +} + +async fn serve(state: &AppState, peer: &str, id: BlobId, chunks: Vec) { + let Some(prepared) = prepare(state, id).await else { + return reply(state, peer, id, BlobBody::Missing).await; + }; + let chunk_size = state.blobs.chunk_size; + if chunks.is_empty() { + let body = BlobBody::Manifest { + size: prepared.bytes.len() as u64, + chunk_size, + hashes: prepared.hashes.clone(), + }; + return reply(state, peer, id, body).await; + } + for index in chunks.into_iter().take(WINDOW * 2) { + let start = index as usize * chunk_size as usize; + if start >= prepared.bytes.len() { + continue; + } + let end = (start + chunk_size as usize).min(prepared.bytes.len()); + let Ok(_permit) = state.blobs.serving.acquire().await else { + return; + }; + let data = bytes::Bytes::copy_from_slice(&prepared.bytes[start..end]); + reply(state, peer, id, BlobBody::Chunk { index, data }).await; + } +} + +// --- fetching ---------------------------------------------------------------------- + +/// Removes the download's answer channel when it ends, however it ends. +struct Registration<'a> { + state: &'a AppState, + id: BlobId, +} + +impl Drop for Registration<'_> { + fn drop(&mut self) { + self.state.blobs.fetches().remove(&self.id); + } +} + +async fn request(state: &AppState, peer: &str, id: BlobId, chunks: Vec) { + let message = ParsedMessage::Blob(BlobMessage { + epoch: state.local_epoch(), + id, + body: BlobBody::Request { chunks }, + }); + if let Err(error) = state.transport.send_message(peer, &message).await { + debug!(peer = %peer, blob = %id, error = %format!("{error:#}"), "Blob request failed"); + } +} + +struct Manifest { + size: u64, + chunk_size: u32, + hashes: Vec<[u8; 32]>, +} + +/// One peer during a download. +struct Source { + peer: String, + in_flight: HashMap, + window: usize, + bad: u32, + timeout: Duration, +} + +/// Downloads blob `id` from `peers` (all thought to hold it). +pub(crate) async fn fetch(state: &AppState, id: BlobId, peers: Vec) -> Result> { + let (sender, mut answers) = mpsc::unbounded_channel(); + if state.blobs.fetches().insert(id, sender).is_some() { + bail!("{id} is being downloaded already"); + } + let _registration = Registration { state, id }; + + // Closest first; peers without a round-trip time last. + let now = Instant::now(); + let mut peers: Vec<(String, Option)> = { + let rtt = state.lock_rtt(); + peers + .into_iter() + .map(|peer| { + let estimate = rtt.effective(&peer, now, state.timing.rtt_fresh_for()); + (peer, estimate) + }) + .collect() + }; + peers.sort_by_key(|(_, rtt)| rtt.unwrap_or(Duration::MAX)); + let rtts: HashMap> = peers.iter().cloned().collect(); + let mut candidates: VecDeque = peers.into_iter().map(|(peer, _)| peer).collect(); + if candidates.is_empty() { + bail!("no peer to download {id} from"); + } + + // The manifest, from the first peer that has the blob. + let mut manifest = None; + let mut missing: HashSet = HashSet::new(); + for peer in &candidates { + request(state, peer, id, Vec::new()).await; + let deadline = Instant::now() + MANIFEST_TIMEOUT; + loop { + let answer = tokio::time::timeout_at(deadline, answers.recv()).await; + match answer { + Ok(Some(( + from, + BlobBody::Manifest { + size, + chunk_size, + hashes, + }, + ))) => { + if size > MAX_SIZE + || chunk_size == 0 + || chunk_size > MAX_CHUNK + || hashes.len() as u64 != chunk_count(size, chunk_size) + { + warn!(peer = %from, blob = %id, size, chunk_size, "Unusable blob manifest; asking another peer"); + missing.insert(from); + break; + } + manifest = Some(Manifest { + size, + chunk_size, + hashes, + }); + break; + } + Ok(Some((from, BlobBody::Missing))) => { + missing.insert(from.clone()); + if &from == peer { + break; + } + } + Ok(Some(_)) => {} // chunks of an earlier attempt + Ok(None) | Err(_) => break, + } + } + if manifest.is_some() { + break; + } + } + let Some(manifest) = manifest else { + bail!("no peer sent a manifest for {id}"); + }; + candidates.retain(|peer| !missing.contains(peer)); + + let total = manifest.hashes.len(); + info!(blob = %id, bytes = manifest.size, chunks = total, peers = candidates.len(), "Downloading blob"); + let started = Instant::now(); + let mut chunks: Vec> = vec![None; total]; + let mut done = 0usize; + let mut sources: Vec = Vec::new(); + let mut last_progress = Instant::now(); + + while done < total { + // Fill up the sources, then their windows. + while sources.len() < MAX_SOURCES { + let Some(peer) = candidates.pop_front() else { + break; + }; + let timeout = + CHUNK_TIMEOUT + rtts.get(&peer).copied().flatten().unwrap_or_default() * 4; + sources.push(Source { + peer, + in_flight: HashMap::new(), + window: WINDOW, + bad: 0, + timeout, + }); + } + if sources.is_empty() { + bail!("no peer left to download {id} from ({done} of {total} chunks)"); + } + let now = Instant::now(); + let assigned: HashSet = sources + .iter() + .flat_map(|source| source.in_flight.keys().copied()) + .collect(); + let mut wanted = (0..total as u32) + .filter(|index| chunks[*index as usize].is_none() && !assigned.contains(index)); + for source in &mut sources { + let free = source.window.saturating_sub(source.in_flight.len()); + let batch: Vec = wanted.by_ref().take(free).collect(); + if batch.is_empty() { + continue; + } + for index in &batch { + source.in_flight.insert(*index, now); + } + request(state, &source.peer, id, batch).await; + } + + match tokio::time::timeout(Duration::from_millis(500), answers.recv()).await { + Ok(Some((from, BlobBody::Chunk { index, data }))) => { + let Some(source) = sources.iter_mut().find(|source| source.peer == from) else { + continue; + }; + let sent_at = source.in_flight.remove(&index); + let slot = index as usize; + if slot >= total || chunks[slot].is_some() { + continue; + } + let expected = expected_len(&manifest, slot); + if data.len() != expected + || *blake3::hash(&data).as_bytes() != manifest.hashes[slot] + { + source.bad += 1; + warn!(peer = %from, blob = %id, index, "Bad blob chunk"); + continue; + } + if let Some(sent_at) = sent_at { + debug!(peer = %from, blob = %id, index, ms = sent_at.elapsed().as_millis() as u64, "Blob chunk"); + } + chunks[slot] = Some(data); + done += 1; + last_progress = Instant::now(); + // A peer that delivers gets its full window back. + source.window = WINDOW; + } + Ok(Some((from, BlobBody::Missing))) => { + sources.retain(|source| source.peer != from); + } + Ok(Some(_)) => {} + Ok(None) => bail!("download of {id} cancelled"), + Err(_) => {} + } + + // Chunks that took too long go to another peer; the slow peer gets a smaller + // window. Peers that sent bad chunks are dropped. + let now = Instant::now(); + for source in &mut sources { + let timeout = source.timeout; + let before = source.in_flight.len(); + source + .in_flight + .retain(|_, sent| now.saturating_duration_since(*sent) < timeout); + if source.in_flight.len() < before { + source.window = 1; + debug!(peer = %source.peer, blob = %id, "Blob chunks timed out; asking others"); + } + } + sources.retain(|source| { + let keep = source.bad < MAX_BAD; + if !keep { + warn!(peer = %source.peer, blob = %id, "Dropping a peer that sent bad blob chunks"); + } + keep + }); + if last_progress.elapsed() > STALL_TIMEOUT { + bail!( + "download of {id} stalled ({done} of {total} chunks, nothing for {}s)", + STALL_TIMEOUT.as_secs() + ); + } + } + + let mut bytes = Vec::with_capacity(manifest.size as usize); + for chunk in chunks.into_iter().flatten() { + bytes.extend_from_slice(&chunk); + } + info!( + blob = %id, + bytes = bytes.len(), + secs = started.elapsed().as_secs_f32(), + "Downloaded blob" + ); + Ok(bytes) +} + +fn expected_len(manifest: &Manifest, slot: usize) -> usize { + let chunk_size = manifest.chunk_size as u64; + let start = slot as u64 * chunk_size; + (manifest.size - start).min(chunk_size) as usize +} + +/// Downloads change `hash` from `peers` and checks it. +pub(crate) async fn fetch_change( + state: &AppState, + hash: ChangeHash, + peers: Vec, +) -> Result { + let bytes = fetch(state, BlobId::Change(hash), peers).await?; + let change = automerge::Change::from_bytes(bytes) + .map_err(|error| anyhow!("downloaded change {hash} doesn't load: {error}"))?; + if change.hash() != hash { + bail!("downloaded change has hash {}, not {hash}", change.hash()); + } + Ok(change) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn chunking() { + let bytes: Vec = (0..1000u32).map(|n| n as u8).collect(); + let hashes = chunk_hashes(&bytes, 300); + assert_eq!(hashes.len(), 4); + assert_eq!(chunk_count(1000, 300), 4); + assert_eq!(chunk_count(900, 300), 3); + let manifest = Manifest { + size: 1000, + chunk_size: 300, + hashes, + }; + assert_eq!(expected_len(&manifest, 0), 300); + assert_eq!(expected_len(&manifest, 3), 100); + assert_eq!(manifest.hashes[3], *blake3::hash(&bytes[900..]).as_bytes()); + } +} diff --git a/src/swarm/mod.rs b/src/swarm/mod.rs index 6f194dc..adb5d1c 100644 --- a/src/swarm/mod.rs +++ b/src/swarm/mod.rs @@ -2,6 +2,7 @@ //! Automerge synchronization with every member (`sync`). pub(crate) mod auth; +pub(crate) mod blob; pub(crate) mod bootstrap; pub(crate) mod members; pub(crate) mod pair; diff --git a/src/swarm/peer.rs b/src/swarm/peer.rs index 3b85539..dd70bbb 100644 --- a/src/swarm/peer.rs +++ b/src/swarm/peer.rs @@ -824,6 +824,9 @@ pub(crate) async fn handle_message( debug!(peer = %peer, error = %format!("{error:#}"), "PONG failed"); } } + transport::ParsedMessage::Blob(blob) => { + crate::swarm::blob::handle(state, &peer, blob).await; + } transport::ParsedMessage::Pong(remote) => { state.lock_rtt().pong_received(&peer, Instant::now()); if reactivated || peer_swarm(state, &peer).await.is_none() { diff --git a/src/swarm/source.rs b/src/swarm/source.rs index 4025f0e..378c40a 100644 --- a/src/swarm/source.rs +++ b/src/swarm/source.rs @@ -199,7 +199,7 @@ pub(crate) fn finished(state: &AppState, from: &str) { } /// After giving up a source: let every peer be considered again. -fn resync_everyone(state: &AppState) { +pub(crate) fn resync_everyone(state: &AppState) { let state = state.clone(); tokio::spawn(async move { let peers: Vec = state.peer_states.read().await.keys().cloned().collect(); diff --git a/src/swarm/sync.rs b/src/swarm/sync.rs index 3284015..58a08db 100644 --- a/src/swarm/sync.rs +++ b/src/swarm/sync.rs @@ -32,6 +32,12 @@ pub(crate) struct ChangeSet { } async fn deliver_sync_to_peer(state: &AppState, peer: &str) -> Result<()> { + if state.remake.fetching_snapshot() { + // Rejoining: the snapshot comes as a download, not in a sync message. Sync + // starts once it is here (`remake::catchup`). + trace!("Waiting for the snapshot download; not syncing yet"); + return Ok(()); + } let Some(peer_state_mutex) = state.get_peer_state(peer).await else { trace!("Peer was forgotten; skipping synchronization"); return Ok(()); -- cgit v1.3.1