diff --git a/Cargo.lock b/Cargo.lock index 1db41cd755..438b68493b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4714,6 +4714,7 @@ dependencies = [ "scopeguard", "serde", "serde_json", + "sync_wrapper", "test-context", "tokio", "tokio-stream", @@ -5956,6 +5957,9 @@ name = "sync_wrapper" version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2047c6ded9c721764247e62cd3b03c09ffc529b2ba5b10ec482ae507a4a70160" +dependencies = [ + "futures-core", +] [[package]] name = "synstructure" @@ -7505,6 +7509,7 @@ dependencies = [ "subtle", "syn 1.0.109", "syn 2.0.52", + "sync_wrapper", "time", "time-macros", "tokio", diff --git a/libs/remote_storage/Cargo.toml b/libs/remote_storage/Cargo.toml index 4a53f485ca..78da01c9a0 100644 --- a/libs/remote_storage/Cargo.toml +++ b/libs/remote_storage/Cargo.toml @@ -38,6 +38,7 @@ azure_storage_blobs.workspace = true futures-util.workspace = true http-types.workspace = true itertools.workspace = true +sync_wrapper = { workspace = true, features = ["futures"] } [dev-dependencies] camino-tempfile.workspace = true diff --git a/libs/remote_storage/src/azure_blob.rs b/libs/remote_storage/src/azure_blob.rs index 5fff3e25c9..24c1248304 100644 --- a/libs/remote_storage/src/azure_blob.rs +++ b/libs/remote_storage/src/azure_blob.rs @@ -3,6 +3,7 @@ use std::borrow::Cow; use std::collections::HashMap; use std::env; +use std::io; use std::num::NonZeroU32; use std::pin::Pin; use std::str::FromStr; @@ -20,6 +21,7 @@ use azure_storage_blobs::blob::CopyStatus; use azure_storage_blobs::prelude::ClientBuilder; use azure_storage_blobs::{blob::operations::GetBlobBuilder, prelude::ContainerClient}; use bytes::Bytes; +use futures::future::Either; use futures::stream::Stream; use futures_util::StreamExt; use futures_util::TryStreamExt; @@ -128,12 +130,12 @@ impl AzureBlobStorage { let kind = RequestKind::Get; let _permit = self.permit(kind, cancel).await?; + let cancel_or_timeout = crate::support::cancel_or_timeout(self.timeout, cancel.clone()); + let cancel_or_timeout_ = crate::support::cancel_or_timeout(self.timeout, cancel.clone()); let mut etag = None; let mut last_modified = None; let mut metadata = HashMap::new(); - // TODO give proper streaming response instead of buffering into RAM - // https://github.com/neondatabase/neon/issues/5563 let download = async { let response = builder @@ -152,39 +154,46 @@ impl AzureBlobStorage { Err(_elapsed) => Err(DownloadError::Timeout), }); - let mut response = std::pin::pin!(response); + let mut response = Box::pin(response); - let mut bufs = Vec::new(); - while let Some(part) = response.next().await { - let part = part?; - if etag.is_none() { - etag = Some(part.blob.properties.etag); - } - if last_modified.is_none() { - last_modified = Some(part.blob.properties.last_modified.into()); - } - if let Some(blob_meta) = part.blob.metadata { - metadata.extend(blob_meta.iter().map(|(k, v)| (k.to_owned(), v.to_owned()))); - } - let data = part - .data - .collect() - .await - .map_err(|e| DownloadError::Other(e.into()))?; - bufs.push(data); - } - - if bufs.is_empty() { + let Some(part) = response.next().await else { return Err(DownloadError::Other(anyhow::anyhow!( - "Azure GET response contained no buffers" + "Azure GET response contained no response body" ))); + }; + let part = part?; + if etag.is_none() { + etag = Some(part.blob.properties.etag); } + if last_modified.is_none() { + last_modified = Some(part.blob.properties.last_modified.into()); + } + if let Some(blob_meta) = part.blob.metadata { + metadata.extend(blob_meta.iter().map(|(k, v)| (k.to_owned(), v.to_owned()))); + } + // unwrap safety: if these were None, bufs would be empty and we would have returned an error already let etag = etag.unwrap(); let last_modified = last_modified.unwrap(); + let tail_stream = response + .map(|part| match part { + Ok(part) => Either::Left(part.data.map(|r| r.map_err(io::Error::other))), + Err(e) => { + Either::Right(futures::stream::once(async { Err(io::Error::other(e)) })) + } + }) + .flatten(); + let stream = part + .data + .map(|r| r.map_err(io::Error::other)) + .chain(sync_wrapper::SyncStream::new(tail_stream)); + //.chain(SyncStream::from_pin(Box::pin(tail_stream))); + + let download_stream = crate::support::DownloadStream::new(cancel_or_timeout_, stream); + Ok(Download { - download_stream: Box::pin(futures::stream::iter(bufs.into_iter().map(Ok))), + download_stream: Box::pin(download_stream), etag, last_modified, metadata: Some(StorageMetadata(metadata)), @@ -193,7 +202,10 @@ impl AzureBlobStorage { tokio::select! { bufs = download => bufs, - _ = cancel.cancelled() => Err(DownloadError::Cancelled), + cancel_or_timeout = cancel_or_timeout => match cancel_or_timeout { + TimeoutOrCancel::Timeout => Err(DownloadError::Timeout), + TimeoutOrCancel::Cancel => Err(DownloadError::Cancelled), + }, } } diff --git a/workspace_hack/Cargo.toml b/workspace_hack/Cargo.toml index a225984688..b2da33e44a 100644 --- a/workspace_hack/Cargo.toml +++ b/workspace_hack/Cargo.toml @@ -66,6 +66,7 @@ serde_json = { version = "1", features = ["raw_value"] } sha2 = { version = "0.10", features = ["asm"] } smallvec = { version = "1", default-features = false, features = ["const_new", "write"] } subtle = { version = "2" } +sync_wrapper = { version = "0.1", default-features = false, features = ["futures"] } time = { version = "0.3", features = ["local-offset", "macros", "serde-well-known"] } tokio = { version = "1", features = ["fs", "io-std", "io-util", "macros", "net", "process", "rt-multi-thread", "signal", "test-util"] } tokio-rustls = { version = "0.24" }