diff --git a/src/app/materialize.rs b/src/app/materialize.rs index 33b7969..aca9664 100644 --- a/src/app/materialize.rs +++ b/src/app/materialize.rs @@ -242,6 +242,10 @@ struct RemoteFile { size_bytes: Option, } +const SMALL_BLOB_MAX_BYTES: u64 = 256 * 1024; +const MAX_SMALL_BLOBS_PER_REQUEST: usize = 128; +const MAX_SMALL_BLOB_REQUEST_BYTES: u64 = 4 * 1024 * 1024; + struct PendingAction { action: ReconciliationAction, missing_blob_count: usize, @@ -397,22 +401,15 @@ async fn receive_and_materialize_files( plan: &mut MaterializationPlan, downloads: Vec, ) -> anyhow::Result<()> { - let mut downloads = stream::iter(downloads) - .map(|file| async move { - let blob_hash = iroh_blobs::Hash::from_str(&file.blob_hash)?; - tracing::info!(path = %file.path, size_bytes = ?file.size_bytes, "Receiving file"); - context - .node - .download_blob(blob_hash, peer.clone()) - .await - .with_context(|| format!("could not receive {}", file.path))?; - Ok::<_, anyhow::Error>(file.blob_hash) - }) + let mut downloads = stream::iter(plan_download_batches(downloads)) + .map(|files| download_files_with_isolation(context, peer, files)) .buffer_unordered(MAX_CONCURRENT_BLOB_TRANSFERS); while let Some(result) = downloads.next().await { match result { - Ok(blob_hash) => { - plan.note_completed_blob(blob_hash); + Ok(blob_hashes) => { + for blob_hash in blob_hashes { + plan.note_completed_blob(blob_hash); + } materialize_ready_actions(context, merged, synchronized_count, plan).await?; } Err(error) => { @@ -423,6 +420,99 @@ async fn receive_and_materialize_files( Ok(()) } +fn plan_download_batches(downloads: Vec) -> Vec> { + let mut batches = Vec::new(); + let mut small_files = Vec::new(); + let mut small_file_bytes = 0; + for file in downloads { + if !is_small_file(&file) { + push_small_file_batch(&mut batches, &mut small_files, &mut small_file_bytes); + batches.push(vec![file]); + continue; + } + if small_files.len() == MAX_SMALL_BLOBS_PER_REQUEST + || small_file_bytes + file.size_bytes.unwrap_or_default() > MAX_SMALL_BLOB_REQUEST_BYTES + { + push_small_file_batch(&mut batches, &mut small_files, &mut small_file_bytes); + } + small_file_bytes += file.size_bytes.unwrap_or_default(); + small_files.push(file); + } + push_small_file_batch(&mut batches, &mut small_files, &mut small_file_bytes); + batches +} + +fn is_small_file(file: &RemoteFile) -> bool { + file.size_bytes + .is_some_and(|size| size <= SMALL_BLOB_MAX_BYTES) +} + +fn push_small_file_batch( + batches: &mut Vec>, + small_files: &mut Vec, + small_file_bytes: &mut u64, +) { + if small_files.is_empty() { + return; + } + batches.push(std::mem::take(small_files)); + *small_file_bytes = 0; +} + +async fn download_files_with_isolation( + context: &SyncContext<'_>, + peer: &iroh::EndpointAddr, + files: Vec, +) -> anyhow::Result> { + let mut completed_hashes = Vec::new(); + let mut pending_batches = VecDeque::from([files]); + while let Some(batch) = pending_batches.pop_front() { + match download_file_batch(context, peer, &batch).await { + Ok(hashes) => completed_hashes.extend(hashes), + Err(error) if batch.len() == 1 => { + tracing::warn!(path = %batch[0].path, error = ?error, "Could not receive file; deferring dependent changes"); + } + Err(error) => { + tracing::debug!(file_count = batch.len(), error = ?error, "Splitting failed blob request"); + let midpoint = batch.len() / 2; + let mut first_half = batch; + let second_half = first_half.split_off(midpoint); + pending_batches.push_back(first_half); + pending_batches.push_back(second_half); + } + } + } + Ok(completed_hashes) +} + +async fn download_file_batch( + context: &SyncContext<'_>, + peer: &iroh::EndpointAddr, + files: &[RemoteFile], +) -> anyhow::Result> { + for file in files { + tracing::debug!(path = %file.path, size_bytes = ?file.size_bytes, "Receiving file"); + } + let hashes = files + .iter() + .map(|file| iroh_blobs::Hash::from_str(&file.blob_hash).map_err(anyhow::Error::from)) + .collect::>>()?; + if let [blob_hash] = hashes.as_slice() { + context + .node + .download_blob(*blob_hash, peer.clone()) + .await + .with_context(|| format!("could not receive {}", files[0].path))?; + } else { + context + .node + .download_blobs(hashes, peer.clone()) + .await + .with_context(|| format!("could not receive {} files", files.len()))?; + } + Ok(files.iter().map(|file| file.blob_hash.clone()).collect()) +} + fn collect_remote_file( files: &mut BTreeMap, required_blobs: &mut BTreeSet, @@ -473,7 +563,10 @@ mod tests { storage::FolderConfig, }; - use super::{plan_remote_materialization, remove_existing_destination, replace_non_directory}; + use super::{ + MAX_SMALL_BLOBS_PER_REQUEST, RemoteFile, plan_download_batches, + plan_remote_materialization, remove_existing_destination, replace_non_directory, + }; #[test] fn replaces_a_file_with_a_directory() -> anyhow::Result<()> { @@ -547,6 +640,37 @@ mod tests { Ok(()) } + #[test] + fn batches_tiny_files_without_mixing_in_large_files() { + let tiny_files = (0..=MAX_SMALL_BLOBS_PER_REQUEST) + .map(|number| RemoteFile { + path: format!("tiny-{number}"), + blob_hash: format!("hash-{number}"), + size_bytes: Some(1), + }) + .collect(); + let downloads = plan_download_batches(tiny_files); + + assert_eq!(downloads.len(), 2); + assert_eq!(downloads[0].len(), MAX_SMALL_BLOBS_PER_REQUEST); + assert_eq!(downloads[1].len(), 1); + + let downloads = plan_download_batches(vec![ + RemoteFile { + path: "tiny".to_owned(), + blob_hash: "tiny-hash".to_owned(), + size_bytes: Some(1), + }, + RemoteFile { + path: "large".to_owned(), + blob_hash: "large-hash".to_owned(), + size_bytes: Some(257 * 1024), + }, + ]); + + assert_eq!(downloads.iter().map(Vec::len).collect::>(), [1, 1]); + } + fn file_entry(path: &str, blob_hash: &str) -> Entry { Entry { path: path.to_owned(), diff --git a/src/iroh.rs b/src/iroh.rs index 8e026eb..9945c97 100644 --- a/src/iroh.rs +++ b/src/iroh.rs @@ -12,7 +12,7 @@ use anyhow::Context; use futures_util::{StreamExt, stream}; use iroh_blobs::{ BlobsProtocol, Hash, - api::downloader::Downloader, + protocol::GetManyRequest, provider::events::{EventMask, EventSender, ProviderMessage, RequestMode}, store::fs::FsStore, }; @@ -84,7 +84,6 @@ pub(super) struct HostedFolder { pub struct NodeHost { endpoint: Endpoint, store: FsStore, - downloader: Downloader, router: Router, folders: Arc>>, discovered_peers: Arc>>, @@ -115,7 +114,6 @@ impl NodeHost { .bind() .await?; let store = FsStore::load(&blob_directory).await?; - let downloader = store.downloader(&endpoint); let folders = Arc::new(RwLock::new(BTreeMap::new())); let discovered_peers = Arc::new(Mutex::new(BTreeMap::new())); let announced_folders = Arc::new(Mutex::new(BTreeMap::new())); @@ -146,7 +144,6 @@ impl NodeHost { Ok(Self { endpoint, store, - downloader, router, folders, discovered_peers, @@ -199,9 +196,22 @@ impl NodeHost { blob_hash: Hash, provider: EndpointAddr, ) -> anyhow::Result<()> { + self.download_blobs(vec![blob_hash], provider).await + } + + pub async fn download_blobs( + &self, + blob_hashes: Vec, + provider: EndpointAddr, + ) -> anyhow::Result<()> { + if blob_hashes.is_empty() { + return Ok(()); + } + let connection = self.connect_to_blob_provider(provider).await?; + let request = blob_hashes.into_iter().collect::(); wait_for_blob_operation( - "download", - self.downloader.download(blob_hash, Some(provider.id)), + "batch download", + self.store.remote().execute_get_many(connection, request), ) .await?; Ok(()) @@ -415,6 +425,21 @@ impl NodeHost { Ok(connection) } + async fn connect_to_blob_provider( + &self, + provider_address: EndpointAddr, + ) -> anyhow::Result { + let connection = tokio::time::timeout( + CONTROL_REQUEST_TIMEOUT, + self.endpoint + .connect(provider_address.clone(), iroh_blobs::ALPN), + ) + .await + .map_err(|_| anyhow::anyhow!("Appa blob provider connection timed out"))??; + verify_connected_peer(&connection, &provider_address)?; + Ok(connection) + } + pub fn take_announced_folders(&self) -> Vec { std::mem::take( &mut *self diff --git a/src/iroh/tests.rs b/src/iroh/tests.rs index 2eef407..bc32405 100644 --- a/src/iroh/tests.rs +++ b/src/iroh/tests.rs @@ -58,6 +58,48 @@ async fn transfers_a_file_between_shared_blob_stores() -> anyhow::Result<()> { Ok(()) } +#[tokio::test] +async fn transfers_tiny_files_in_one_blob_request() -> anyhow::Result<()> { + let source_directory = TempDir::new()?; + let target_directory = TempDir::new()?; + let source = + NodeHost::load_with_lan_discovery(source_directory.path(), SecretKey::generate(), false) + .await?; + let target = + NodeHost::load_with_lan_discovery(target_directory.path(), SecretKey::generate(), false) + .await?; + let first_source = source_directory.path().join("first.txt"); + let second_source = source_directory.path().join("second.txt"); + fs::write(&first_source, "first")?; + fs::write(&second_source, "second")?; + let first_hash = source.import_file(&first_source).await?; + let second_hash = source.import_file(&second_source).await?; + source.wait_until_online().await?; + target.wait_until_online().await?; + + target + .download_blobs(vec![first_hash, second_hash], source.endpoint_address()) + .await?; + target + .export_blob(first_hash, &target_directory.path().join("first.txt")) + .await?; + target + .export_blob(second_hash, &target_directory.path().join("second.txt")) + .await?; + + assert_eq!( + fs::read_to_string(target_directory.path().join("first.txt"))?, + "first" + ); + assert_eq!( + fs::read_to_string(target_directory.path().join("second.txt"))?, + "second" + ); + source.shutdown().await?; + target.shutdown().await?; + Ok(()) +} + #[tokio::test] async fn announces_folder_changes_to_an_authorized_peer() -> anyhow::Result<()> { let source_directory = TempDir::new()?;