introduce bounded blob downloads

4 files changed, 55 insertions, 20 deletions

diff --git a/src/remake/content.rs b/src/remake/content.rs
index f3e549e..d036b87 100644
--- a/src/remake/content.rs
+++ b/src/remake/content.rs
@@ -66,7 +66,7 @@ pub(super) async fn ensure<H: Host>(
 }
 
 async fn fetch<H: Host>(host: &H, hash: [u8; 32], reference: &ContentRef) -> Option<Arc<Vec<u8>>> {
-    match host.fetch_content(hash).await {
+    match host.fetch_content(hash, reference.size).await {
         Ok(bytes) if bytes.len() as u64 == reference.size => {
             info!(content = %reference.blake3, bytes = bytes.len(), "Downloaded remake content");
             let bytes = Arc::new(bytes);
diff --git a/src/remake/mod.rs b/src/remake/mod.rs
index 38c2b9c..ad027a7 100644
--- a/src/remake/mod.rs
+++ b/src/remake/mod.rs
@@ -87,10 +87,13 @@ pub(crate) trait Host: Clone + Send + Sync + 'static {
         hash: automerge::ChangeHash,
     ) -> impl Future<Output = anyhow::Result<automerge::Change>> + Send;
 
-    /// Downloads content by its BLAKE3 hash from peers on the local epoch, checked
-    /// against it.
-    fn fetch_content(&self, hash: [u8; 32])
-    -> impl Future<Output = anyhow::Result<Vec<u8>>> + Send;
+    /// Downloads content of `size` bytes by its BLAKE3 hash from peers on the local
+    /// epoch, checked against it.
+    fn fetch_content(
+        &self,
+        hash: [u8; 32],
+        size: u64,
+    ) -> impl Future<Output = anyhow::Result<Vec<u8>>> + Send;
 
     /// Starts sync with every peer (e.g. after holding it back).
     fn request_sync(&self);
diff --git a/src/state/host.rs b/src/state/host.rs
index 6506274..f6c460d 100644
--- a/src/state/host.rs
+++ b/src/state/host.rs
@@ -100,11 +100,15 @@ impl Host for AppState {
         }
     }
 
-    fn fetch_content(&self, hash: [u8; 32]) -> impl Future<Output = Result<Vec<u8>>> + Send {
+    fn fetch_content(
+        &self,
+        hash: [u8; 32],
+        size: u64,
+    ) -> impl Future<Output = Result<Vec<u8>>> + Send {
         let state = self.clone();
         async move {
             let peers = state.active_peers_on_local_epoch().await;
-            crate::swarm::blob::fetch_content(&state, hash, peers).await
+            crate::swarm::blob::fetch_content(&state, hash, size, peers).await
         }
     }
 
diff --git a/src/swarm/blob.rs b/src/swarm/blob.rs
index 4b69562..a2feb98 100644
--- a/src/swarm/blob.rs
+++ b/src/swarm/blob.rs
@@ -10,8 +10,10 @@
 //!   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
+//!   in tasks of its own, at most [`MAX_ANSWERING`] requests and [`MAX_SERVING`] chunk
+//!   sends at a time.
+//! * Fetching: the manifest from the first peer that has it (closest first, and only
+//!   from a peer asked; of content, only with the size its reference names), 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
@@ -39,10 +41,16 @@ const MAX_SERVING: usize = 8;
 const MAX_SOURCES: usize = 4;
 /// Chunks in flight per peer.
 const WINDOW: usize = 4;
-/// Largest blob a download accepts.
+/// Largest blob a download of unknown size accepts (a change: an epoch's snapshot is as
+/// large as the document). Content has a known size, and only that is accepted.
 const MAX_SIZE: u64 = 4 << 30;
-/// Largest chunk a download accepts.
+/// Smallest and largest chunk a download accepts (`--blob-chunk`'s range). The lower
+/// bound keeps a manifest's chunk list, and the chunk map built from it, small.
+const MIN_CHUNK: u32 = 1 << 10;
 const MAX_CHUNK: u32 = 4 << 20;
+/// Requests being answered at once; more are dropped (the requester asks again, or
+/// another peer). Bounds the tasks anyone can make a node spawn.
+const MAX_ANSWERING: usize = 32;
 /// 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
@@ -57,6 +65,8 @@ const MAX_BAD: u32 = 2;
 pub(crate) struct Blobs {
     chunk_size: u32,
     serving: Arc<Semaphore>,
+    /// Requests being answered (`MAX_ANSWERING`).
+    answering: Arc<Semaphore>,
     /// The blob served last.
     cache: Mutex<Option<(BlobId, Arc<Prepared>)>>,
     /// Answers for downloads in progress, by blob.
@@ -73,6 +83,7 @@ impl Blobs {
         Self {
             chunk_size: chunk_size.max(1),
             serving: Arc::new(Semaphore::new(MAX_SERVING)),
+            answering: Arc::new(Semaphore::new(MAX_ANSWERING)),
             cache: Mutex::new(None),
             fetches: Mutex::new(HashMap::new()),
         }
@@ -112,10 +123,17 @@ fn chunk_count(size: u64, chunk_size: u32) -> u64 {
 pub(crate) async fn handle(state: &AppState, peer: &str, message: BlobMessage) {
     match message.body {
         BlobBody::Request { chunks } => {
+            let Ok(permit) = Arc::clone(&state.blobs.answering).try_acquire_owned() else {
+                debug!(peer = %peer, blob = %message.id, "Answering too many blob requests; dropping one");
+                return;
+            };
             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 });
+            tokio::spawn(async move {
+                serve(&state, &peer, message.id, chunks).await;
+                drop(permit);
+            });
         }
         body => {
             let sender = state.blobs.fetches().get(&message.id).cloned();
@@ -247,8 +265,14 @@ struct Source {
     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>> {
+/// Downloads blob `id` from `peers` (all thought to hold it): of `expected_size` bytes
+/// if known, otherwise at most `MAX_SIZE`.
+pub(crate) async fn fetch(
+    state: &AppState,
+    id: BlobId,
+    peers: Vec<String>,
+    expected_size: Option<u64>,
+) -> 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");
@@ -290,10 +314,13 @@ pub(crate) async fn fetch(state: &AppState, id: BlobId, peers: Vec<String>) -> R
                         chunk_size,
                         hashes,
                     },
-                ))) => {
-                    if size > MAX_SIZE
-                        || chunk_size == 0
-                        || chunk_size > MAX_CHUNK
+                ))) if candidates.contains(&from) => {
+                    let size_fits = match expected_size {
+                        Some(expected) => size == expected,
+                        None => size <= MAX_SIZE,
+                    };
+                    if !size_fits
+                        || !(MIN_CHUNK..=MAX_CHUNK).contains(&chunk_size)
                         || hashes.len() as u64 != chunk_count(size, chunk_size)
                     {
                         warn!(peer = %from, blob = %id, size, chunk_size, "Unusable blob manifest; asking another peer");
@@ -460,10 +487,11 @@ fn expected_len(manifest: &Manifest, slot: usize) -> usize {
 pub(crate) async fn fetch_content(
     state: &AppState,
     hash: [u8; 32],
+    size: u64,
     peers: Vec<String>,
 ) -> Result<Vec<u8>> {
     let id = BlobId::Content(hash);
-    let bytes = fetch(state, id, peers).await?;
+    let bytes = fetch(state, id, peers, Some(size)).await?;
     if *blake3::hash(&bytes).as_bytes() != hash {
         bail!("downloaded {id} doesn't match its hash");
     }
@@ -475,7 +503,7 @@ pub(crate) async fn fetch_change(
     hash: ChangeHash,
     peers: Vec<String>,
 ) -> Result<automerge::Change> {
-    let bytes = fetch(state, BlobId::Change(hash), peers).await?;
+    let bytes = fetch(state, BlobId::Change(hash), peers, None).await?;
     let change = automerge::Change::from_bytes(bytes)
         .map_err(|error| anyhow!("downloaded change {hash} doesn't load: {error}"))?;
     if change.hash() != hash {