From 5f099dc7603d0b41418ad9b5e7267e377f24534c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arpad=20M=C3=BCller?= Date: Thu, 2 May 2024 20:19:00 +0200 Subject: [PATCH] Use streaming downloads for Azure as well (#7579) The main challenge was in the second commit, as `DownloadStream` requires the inner to be Sync but the stream returned by the Azure SDK wasn't Sync. This left us with three options: * Change the Azure SDK to return Sync streams. This was abandoned after we realized that we couldn't just make `TokenCredential`'s returned future Sync: it uses the `async_trait` macro and as the `TokenCredential` trait is used in dyn form, one can't use Rust's new "async fn in Trait" feature. * Change `DownloadStream` to not require `Sync`. This was abandoned after it turned into a safekeeper refactoring project. * Put the stream into a `Mutex` and make it obtain a lock on every poll. This adds some performance overhead but locks that actually don't do anything should be comparatively cheap. We went with the third option in the end as the change still represents an improvement. Follow up of #5446 , fixes #5563 --- Cargo.lock | 5 ++ libs/remote_storage/Cargo.toml | 1 + libs/remote_storage/src/azure_blob.rs | 66 ++++++++++++++++----------- workspace_hack/Cargo.toml | 1 + 4 files changed, 46 insertions(+), 27 deletions(-) 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" }