mirror of
https://github.com/rustfs/rustfs.git
synced 2026-08-30 08:49:26 +00:00
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:
@@ -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]
|
||||||
|
|||||||
Reference in New Issue
Block a user