perf(get): avoid memory body stream wrapper (#6163)

Use MemoryTrackedBytesStream directly as an s3s ByteStream so in-memory GET bodies avoid the generic StreamingBlob::wrap adapter while preserving exact remaining length, request lifecycle tracking, and length-mismatch failure semantics.

Co-authored-by: heihutu <heihutu@gmail.com>
This commit is contained in:
houseme
2026-08-17 10:45:24 +08:00
committed by GitHub
parent 01e0af6312
commit 3f3e3f4f05
+34 -5
View File
@@ -1074,7 +1074,7 @@ where
} }
impl futures::Stream for MemoryTrackedBytesStream { impl futures::Stream for MemoryTrackedBytesStream {
type Item = std::io::Result<Bytes>; type Item = Result<Bytes, S3StdError>;
fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut(); let this = self.get_mut();
@@ -1105,7 +1105,8 @@ impl futures::Stream for MemoryTrackedBytesStream {
return Poll::Ready(Some(Err(std::io::Error::new( return Poll::Ready(Some(Err(std::io::Error::new(
std::io::ErrorKind::InvalidData, std::io::ErrorKind::InvalidData,
format!("materialized GET body length mismatch: expected {}, got {}", this.expected, actual), format!("materialized GET body length mismatch: expected {}, got {}", this.expected, actual),
)))); )
.into())));
} }
let Some(bytes) = this.bytes.take() else { let Some(bytes) = this.bytes.take() else {
@@ -1132,6 +1133,16 @@ impl futures::Stream for MemoryTrackedBytesStream {
} }
} }
impl ByteStream for MemoryTrackedBytesStream {
fn remaining_length(&self) -> RemainingLength {
if self.emitted || self.bytes.is_none() {
RemainingLength::new_exact(0)
} else {
RemainingLength::new_exact(self.expected)
}
}
}
impl Drop for MemoryTrackedBytesStream { impl Drop for MemoryTrackedBytesStream {
fn drop(&mut self) { fn drop(&mut self) {
if self.lifecycle.is_finished() { if self.lifecycle.is_finished() {
@@ -4149,7 +4160,7 @@ impl DefaultObjectUsecase {
let bytes_len = bytes.len(); let bytes_len = bytes.len();
let guard = rustfs_io_metrics::track_get_object_buffered_bytes(bytes_len); let guard = rustfs_io_metrics::track_get_object_buffered_bytes(bytes_len);
let remaining = usize::try_from(response_content_length.max(0)).unwrap_or(usize::MAX); let remaining = usize::try_from(response_content_length.max(0)).unwrap_or(usize::MAX);
let blob = StreamingBlob::wrap(MemoryTrackedBytesStream::new(bytes, remaining, source, guard, lifecycle)); let blob = StreamingBlob::new(MemoryTrackedBytesStream::new(bytes, remaining, source, guard, lifecycle));
if let Some(handoff_start) = handoff_start { if let Some(handoff_start) = handoff_start {
rustfs_io_metrics::record_get_object_response_handoff( rustfs_io_metrics::record_get_object_response_handoff(
"single_chunk", "single_chunk",
@@ -12882,7 +12893,10 @@ mod tests {
.await .await
.expect("mismatched memory body must yield an item") .expect("mismatched memory body must yield an item")
.expect_err("a short memory body must fail the stream instead of serving a truncated body"); .expect_err("a short memory body must fail the stream instead of serving a truncated body");
assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); assert_eq!(
err.downcast_ref::<std::io::Error>().map(std::io::Error::kind),
Some(std::io::ErrorKind::InvalidData)
);
assert!(stream.next().await.is_none(), "stream must terminate after the error"); assert!(stream.next().await.is_none(), "stream must terminate after the error");
} }
@@ -12901,7 +12915,22 @@ mod tests {
.await .await
.expect("mismatched memory body must yield an item") .expect("mismatched memory body must yield an item")
.expect_err("an over-long memory body must fail the stream instead of serving mismatched bytes"); .expect_err("an over-long memory body must fail the stream instead of serving mismatched bytes");
assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); assert_eq!(
err.downcast_ref::<std::io::Error>().map(std::io::Error::kind),
Some(std::io::ErrorKind::InvalidData)
);
}
#[test]
fn memory_blob_preserves_exact_remaining_length() {
let blob = DefaultObjectUsecase::build_memory_bytes_blob(
Bytes::from_static(b"hello"),
5,
GET_MEMORY_BODY_SOURCE_BUFFERED_BODY,
GetObjectBodyLifecycle::disabled(),
);
assert_eq!(blob.remaining_length().exact(), Some(5));
} }
#[tokio::test] #[tokio::test]