Make Storage::get_slice zero copy for single-segment bodies (#6336)

Co-authored-by: PSeitz <PSeitz@users.noreply.github.com>

Author: guilload

quickwit/Cargo.lock

@@ -240,6 +240,7 @@ serde_yaml = "0.9"
 serial_test = { version = "3.2", features = ["file_locks"] }
 sha2 = "0.11"
 siphasher = "1.0"
+stable_deref_trait = "1.2"
 stateright = "0.31"
 storekey = { version = "0.11", default-features = false }
 smallvec = "1"

quickwit/Cargo.toml

@@ -29,6 +29,7 @@ quick_cache = { workspace = true }
 regex = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
+stable_deref_trait = { workspace = true }
 tantivy = { workspace = true }
 tempfile = { workspace = true }
 thiserror = { workspace = true }

quickwit/quickwit-storage/Cargo.toml

@@ -50,6 +50,7 @@ mod prefix_storage;
 mod ram_storage;
 mod split;
 mod split_cache;
+mod stable_deref_bytes;
 mod storage_factory;
 mod storage_resolver;
 mod versioned_component;

quickwit/quickwit-storage/src/lib.rs

@@ -26,7 +26,7 @@ use azure_storage::Error as AzureError;
 use azure_storage::prelude::*;
 use azure_storage_blobs::blob::operations::GetBlobResponse;
 use azure_storage_blobs::prelude::*;
-use bytes::Bytes;
+use bytes::{Bytes, BytesMut};
 use futures::io::Error as FutureError;
 use futures::stream::{StreamExt, TryStreamExt};
 use md5::Digest;
@@ -44,6 +44,7 @@ use tracing::{instrument, warn};
 
 use crate::debouncer::DebouncedStorage;
 use crate::metrics::object_storage_get_slice_in_flight_guards;
+use crate::stable_deref_bytes::into_owned_bytes;
 use crate::storage::SendableAsync;
 use crate::{
     BulkDeleteError, DeleteFailure, MultiPartPolicy, PutPayload, STORAGE_METRICS, Storage,
@@ -203,12 +204,12 @@ impl AzureBlobStorage {
         key_path.to_string_lossy().to_string()
     }
 
-    /// Downloads a blob as vector of bytes.
-    async fn get_to_vec(
+    /// Downloads a blob as `Bytes` — zero-copy when the blob arrives as a single chunk.
+    async fn get_to_bytes(
         &self,
         path: &Path,
         range_opt: Option<Range<usize>>,
-    ) -> StorageResult<Vec<u8>> {
+    ) -> StorageResult<Bytes> {
         let name = self.blob_name(path);
         let capacity = range_opt.as_ref().map(Range::len).unwrap_or(0);
         retry(&self.retry_params, || async {
@@ -226,10 +227,8 @@ impl AzureBlobStorage {
                 let stream = self.container_client.blob_client(&name).get().into_stream();
                 (stream, None)
             };
-            let mut buf: Vec<u8> = Vec::with_capacity(capacity);
-            download_all(&mut response_stream, &mut buf).await?;
-
-            Result::<_, AzureErrorWrapper>::Ok(buf)
+            let bytes = download_all(&mut response_stream).await?;
+            Result::<_, AzureErrorWrapper>::Ok(bytes)
         })
         .await
         .map_err(StorageError::from)
@@ -442,9 +441,9 @@ impl Storage for AzureBlobStorage {
 
     #[instrument(level = "debug", skip(self, range), fields(range.start = range.start, range.end = range.end))]
     async fn get_slice(&self, path: &Path, range: Range<usize>) -> StorageResult<OwnedBytes> {
-        self.get_to_vec(path, Some(range.clone()))
+        self.get_to_bytes(path, Some(range.clone()))
             .await
-            .map(OwnedBytes::new)
+            .map(into_owned_bytes)
             .map_err(|err| {
                 err.add_context(format!(
                     "failed to fetch slice {:?} for object: {}/{}",
@@ -493,9 +492,9 @@ impl Storage for AzureBlobStorage {
     #[instrument(level = "debug", skip(self), fields(fetched_bytes_len))]
     async fn get_all(&self, path: &Path) -> StorageResult<OwnedBytes> {
         let data = self
-            .get_to_vec(path, None)
+            .get_to_bytes(path, None)
             .await
-            .map(OwnedBytes::new)
+            .map(into_owned_bytes)
             .map_err(|err| {
                 err.add_context(format!(
                     "failed to fetch object: {}/{}",
@@ -559,28 +558,45 @@ pub fn parse_azure_uri(uri: &Uri) -> Option<(String, PathBuf)> {
     Some((container, prefix))
 }
 
-/// Collect a download stream into an output buffer.
+/// Collect a download stream into a single [`Bytes`].
+///
+/// `Bytes` segments yielded by the SDK are preserved so that the single-segment case avoids the
+/// extra copy into a contiguous buffer. When more than one segment is received, they are
+/// concatenated exactly once into a freshly allocated `Bytes`.
 async fn download_all(
     chunk_stream: &mut Pageable<GetBlobResponse, AzureError>,
-    output: &mut Vec<u8>,
-) -> Result<(), AzureErrorWrapper> {
-    output.clear();
+) -> Result<Bytes, AzureErrorWrapper> {
+    let mut segments: Vec<Bytes> = Vec::new();
+    let mut total_num_bytes: usize = 0;
     while let Some(chunk_result) = chunk_stream.next().await {
         let chunk_response = chunk_result?;
-        let chunk_response_body_stream = chunk_response
-            .data
-            .map_err(FutureError::other)
-            .into_async_read()
-            .compat();
-        let mut body_stream_reader = BufReader::new(chunk_response_body_stream);
-        let num_bytes_copied = tokio::io::copy_buf(&mut body_stream_reader, output).await?;
-        crate::STORAGE_METRICS
-            .object_storage_download_num_bytes
-            .inc_by(num_bytes_copied);
+        let mut data_stream = chunk_response.data;
+        while let Some(bytes_res) = data_stream.next().await {
+            let bytes = bytes_res?;
+            total_num_bytes += bytes.len();
+            segments.push(bytes);
+        }
+    }
+    crate::STORAGE_METRICS
+        .object_storage_download_num_bytes
+        .inc_by(total_num_bytes as u64);
+    Ok(coalesce_segments(segments, total_num_bytes))
+}
+
+/// Returns a single [`Bytes`] covering `segments`. Zero-copy when there is at most one segment;
+/// otherwise a single allocation concatenates them.
+fn coalesce_segments(mut segments: Vec<Bytes>, total_num_bytes: usize) -> Bytes {
+    match segments.len() {
+        0 => Bytes::new(),
+        1 => segments.remove(0),
+        _ => {
+            let mut out = BytesMut::with_capacity(total_num_bytes);
+            for segment in segments {
+                out.extend_from_slice(&segment);
+            }
+            out.freeze()
+        }
     }
-    // When calling `get_all`, the Vec capacity is not properly set.
-    output.shrink_to_fit();
-    Ok(())
 }
 
 #[derive(Error, Debug)]

quickwit/quickwit-storage/src/object_storage/azure_blob_storage.rs

@@ -30,10 +30,11 @@ use aws_sdk_s3::config::{
 use aws_sdk_s3::error::{ProvideErrorMetadata, SdkError};
 use aws_sdk_s3::operation::delete_objects::DeleteObjectsOutput;
 use aws_sdk_s3::operation::get_object::{GetObjectError, GetObjectOutput};
-use aws_sdk_s3::primitives::ByteStream;
+use aws_sdk_s3::primitives::{AggregatedBytes, ByteStream};
 use aws_sdk_s3::types::builders::ObjectIdentifierBuilder;
 use aws_sdk_s3::types::{CompletedMultipartUpload, CompletedPart, Delete, ObjectIdentifier};
 use base64::prelude::{BASE64_STANDARD, Engine};
+use bytes::Bytes;
 use futures::{StreamExt, stream};
 use quickwit_aws::retry::{AwsRetryable, aws_retry};
 use quickwit_aws::{aws_behavior_version, get_aws_config};
@@ -48,6 +49,7 @@ use tracing::{info, instrument, warn};
 
 use crate::metrics::object_storage_get_slice_in_flight_guards;
 use crate::object_storage::MultiPartPolicy;
+use crate::stable_deref_bytes::into_owned_bytes;
 use crate::storage::SendableAsync;
 use crate::{
     BulkDeleteError, DeleteFailure, OwnedBytes, STORAGE_METRICS, Storage, StorageError,
@@ -445,11 +447,11 @@ impl S3CompatibleObjectStorage {
             .upload_id(upload_id.0)
             .send()
             .await
-            .map_err(|s3_err| {
-                if s3_err.is_retryable() {
-                    Retry::Transient(StorageError::from(s3_err))
+            .map_err(|sdk_error| {
+                if sdk_error.is_retryable() {
+                    Retry::Transient(StorageError::from(sdk_error))
                 } else {
-                    Retry::Permanent(StorageError::from(s3_err))
+                    Retry::Permanent(StorageError::from(sdk_error))
                 }
             })?;
 
@@ -484,7 +486,7 @@ impl S3CompatibleObjectStorage {
             .collect::<Vec<_>>()
             .await
             .into_iter()
-            .map(|res| res.map_err(|e| e.into_inner()))
+            .map(|res| res.map_err(|error| error.into_inner()))
             .collect();
         match completed_parts_res {
             Ok(completed_parts) => {
@@ -564,22 +566,19 @@ impl S3CompatibleObjectStorage {
         Ok(get_object_output)
     }
 
-    async fn get_to_vec(
+    async fn get_to_bytes(
         &self,
         path: &Path,
         range_opt: Option<Range<usize>>,
-    ) -> StorageResult<Vec<u8>> {
-        let cap = range_opt.as_ref().map(Range::len).unwrap_or(0);
+    ) -> StorageResult<Bytes> {
         let get_object_output = aws_retry(&self.retry_params, || {
             self.get_object(path, range_opt.clone())
         })
         .await?;
         // only record ranged get request as being in flight
         let _in_flight_guards =
             range_opt.map(|range| object_storage_get_slice_in_flight_guards(range.len()));
-        let mut buf: Vec<u8> = Vec::with_capacity(cap);
-        download_all(get_object_output.body, &mut buf).await?;
-        Ok(buf)
+        download_all(get_object_output.body).await
     }
 
     /// Bulk delete implementation based on the DeleteObject API:
@@ -720,16 +719,35 @@ impl S3CompatibleObjectStorage {
     }
 }
 
-async fn download_all(byte_stream: ByteStream, output: &mut Vec<u8>) -> io::Result<()> {
-    output.clear();
-    let mut body_stream_reader = BufReader::new(byte_stream.into_async_read());
-    let num_bytes_copied = tokio::io::copy_buf(&mut body_stream_reader, output).await?;
+async fn download_all(byte_stream: ByteStream) -> StorageResult<Bytes> {
+    let aggregated: AggregatedBytes = byte_stream
+        .collect()
+        .await
+        .map_err(byte_stream_to_storage_error)?;
+    // `AggregatedBytes::into_bytes` returns the underlying `Bytes` without copying when the body
+    // was received as a single segment, and concatenates into a fresh `Bytes` otherwise.
+    let bytes = aggregated.into_bytes();
     STORAGE_METRICS
         .object_storage_download_num_bytes
-        .inc_by(num_bytes_copied);
-    // When calling `get_all`, the Vec capacity is not properly set.
-    output.shrink_to_fit();
-    Ok(())
+        .inc_by(bytes.len() as u64);
+    Ok(bytes)
+}
+
+/// Classifies a `ByteStream::collect` failure. The response headers were already received
+/// successfully at this point; what fails here is draining the body. The underlying cause can be a
+/// local `io::Error` (which we classify as `Io`) or a transport-level streaming failure — stalled
+/// stream, malformed body, checksum mismatch — which we classify as `Internal` and report with the
+/// original error preserved as source.
+fn byte_stream_to_storage_error(error: aws_sdk_s3::primitives::ByteStreamError) -> StorageError {
+    let is_io_error = std::error::Error::source(&error)
+        .and_then(|source| source.downcast_ref::<io::Error>())
+        .is_some();
+    let kind = if is_io_error {
+        StorageErrorKind::Io
+    } else {
+        StorageErrorKind::Internal
+    };
+    kind.with_error(error)
 }
 
 #[async_trait]
@@ -816,9 +834,9 @@ impl Storage for S3CompatibleObjectStorage {
     #[instrument(level = "debug", skip(self, range), fields(range.start = range.start, range.end = range.end))]
     async fn get_slice(&self, path: &Path, range: Range<usize>) -> StorageResult<OwnedBytes> {
         let _permit = REQUEST_SEMAPHORE.acquire().await;
-        self.get_to_vec(path, Some(range.clone()))
+        self.get_to_bytes(path, Some(range.clone()))
             .await
-            .map(OwnedBytes::new)
+            .map(into_owned_bytes)
             .map_err(|err| {
                 err.add_context(format!(
                     "failed to fetch slice {:?} for object: {}/{}",
@@ -850,9 +868,9 @@ impl Storage for S3CompatibleObjectStorage {
     async fn get_all(&self, path: &Path) -> StorageResult<OwnedBytes> {
         let _permit = REQUEST_SEMAPHORE.acquire().await;
         let bytes = self
-            .get_to_vec(path, None)
+            .get_to_bytes(path, None)
             .await
-            .map(OwnedBytes::new)
+            .map(into_owned_bytes)
             .map_err(|err| {
                 err.add_context(format!(
                     "failed to fetch object: {}/{}",

quickwit/quickwit-storage/src/object_storage/s3_compatible_storage.rs

@@ -24,6 +24,7 @@ use tokio::io::{AsyncRead, AsyncWriteExt as TokioAsyncWriteExt};
 use tokio_util::compat::{FuturesAsyncReadCompatExt, FuturesAsyncWriteCompatExt};
 
 use crate::metrics::object_storage_get_slice_in_flight_guards;
+use crate::stable_deref_bytes::into_owned_bytes;
 use crate::storage::SendableAsync;
 use crate::{
     BulkDeleteError, MultiPartPolicy, OwnedBytes, PutPayload, Storage, StorageError,
@@ -123,8 +124,11 @@ impl Storage for OpendalStorage {
         // recorded before issuing the query to the object store.
         let _inflight_guards = object_storage_get_slice_in_flight_guards(size);
         crate::STORAGE_METRICS.object_storage_get_total.inc();
-        let storage_content = self.op.read_with(&path).range(range).await?.to_vec();
-        Ok(OwnedBytes::new(storage_content))
+        // `Buffer::to_bytes` is zero-copy when the underlying buffer is contiguous, and coalesces
+        // into a single `Bytes` otherwise — avoiding the extra `Vec<u8>` round-trip `to_vec` would
+        // perform.
+        let storage_content = self.op.read_with(&path).range(range).await?.to_bytes();
+        Ok(into_owned_bytes(storage_content))
     }
 
     async fn get_slice_stream(
@@ -146,8 +150,8 @@ impl Storage for OpendalStorage {
 
     async fn get_all(&self, path: &Path) -> StorageResult<OwnedBytes> {
         let path = path.as_os_str().to_string_lossy();
-        let storage_content = self.op.read(&path).await?.to_vec();
-        Ok(OwnedBytes::new(storage_content))
+        let storage_content = self.op.read(&path).await?.to_bytes();
+        Ok(into_owned_bytes(storage_content))
     }
 
     async fn delete(&self, path: &Path) -> StorageResult<()> {

quickwit/quickwit-storage/src/opendal_storage/base.rs

@@ -0,0 +1,45 @@
+// Copyright 2021-Present Datadog, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::ops::Deref;
+
+use bytes::Bytes;
+use stable_deref_trait::StableDeref;
+
+use crate::OwnedBytes;
+
+/// `StableDeref` wrapper around [`Bytes`] so it can back an [`OwnedBytes`] without an extra copy.
+///
+/// `Bytes` dereferences to a heap-allocated slice whose address does not change when the `Bytes`
+/// itself is moved, which is exactly the contract `StableDeref` requires.
+pub(crate) struct StableDerefBytes(pub Bytes);
+
+impl Deref for StableDerefBytes {
+    type Target = [u8];
+
+    #[inline]
+    fn deref(&self) -> &[u8] {
+        &self.0
+    }
+}
+
+// SAFETY: `Bytes` stores its payload behind a stable heap pointer; moving the `Bytes` (and thus
+// the `StableDerefBytes`) does not invalidate the slice returned by `deref`.
+unsafe impl StableDeref for StableDerefBytes {}
+
+/// Wraps a [`Bytes`] into an [`OwnedBytes`] without copying its contents.
+#[inline]
+pub(crate) fn into_owned_bytes(bytes: Bytes) -> OwnedBytes {
+    OwnedBytes::new(StableDerefBytes(bytes))
+}