add chunked download and use it during rejoin

5 files changed, 502 insertions, 1 deletions

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<Semaphore>,
+    /// The blob served last.
+    cache: Mutex<Option<(BlobId, Arc<Prepared>)>>,
+    /// Answers for downloads in progress, by blob.
+    fetches: Mutex<HashMap<BlobId, mpsc::UnboundedSender<(String, BlobBody)>>>,
+}
+
+struct Prepared {
+    bytes: Vec<u8>,
+    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<BlobId, mpsc::UnboundedSender<(String, BlobBody)>>> {
+        self.fetches
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner())
+    }
+
+    fn cached(&self, id: BlobId) -> Option<Arc<Prepared>> {
+        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<Prepared>) {
+        *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<Arc<Prepared>> {
+    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<u32>) {
+    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<u32>) {
+    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<u32, Instant>,
+    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<String>) -> Result<Vec<u8>> {
+    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<Duration>)> = {
+        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<String, Option<Duration>> = peers.iter().cloned().collect();
+    let mut candidates: VecDeque<String> = 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<String> = 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<Option<bytes::Bytes>> = vec![None; total];
+    let mut done = 0usize;
+    let mut sources: Vec<Source> = 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<u32> = 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<u32> = 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<String>,
+) -> Result<automerge::Change> {
+    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<u8> = (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<String> = 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(());