Skip to content

Commit acf4aa1

Browse files
durchdurch
committed
fix: preserve response stream chunk remainders
1 parent 57525a4 commit acf4aa1

1 file changed

Lines changed: 92 additions & 73 deletions

File tree

‎s3/src/request/request_trait.rs‎

Lines changed: 92 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,10 @@ use http::header::{
2121
};
2222
use std::fmt::Write as _;
2323

24-
#[cfg(feature = "with-async-std")]
25-
use async_std::stream::Stream;
26-
27-
#[cfg(feature = "with-tokio")]
28-
use tokio_stream::Stream;
24+
#[cfg(any(feature = "with-tokio", feature = "with-async-std"))]
25+
use futures_util::Stream;
26+
#[cfg(any(feature = "with-tokio", feature = "with-async-std"))]
27+
use futures_util::stream::{self, StreamExt};
2928

3029
#[derive(Debug)]
3130

@@ -126,28 +125,39 @@ impl tokio::io::AsyncRead for ResponseDataStream {
126125
cx: &mut std::task::Context<'_>,
127126
buf: &mut tokio::io::ReadBuf<'_>,
128127
) -> std::task::Poll<std::io::Result<()>> {
129-
// Poll the stream for the next chunk of bytes
130-
match Stream::poll_next(self.bytes.as_mut(), cx) {
131-
std::task::Poll::Ready(Some(Ok(chunk))) => {
132-
// Write as much of the chunk as fits in the buffer
133-
let amt = std::cmp::min(chunk.len(), buf.remaining());
134-
buf.put_slice(&chunk[..amt]);
135-
136-
// AIDEV-NOTE: Bytes that don't fit in the buffer are discarded from this chunk.
137-
// This is expected AsyncRead behavior - consumers should use appropriately sized
138-
// buffers or wrap in BufReader for efficiency with small reads.
139-
140-
std::task::Poll::Ready(Ok(()))
141-
}
142-
std::task::Poll::Ready(Some(Err(error))) => {
143-
// Convert S3Error to io::Error
144-
std::task::Poll::Ready(Err(std::io::Error::other(error)))
145-
}
146-
std::task::Poll::Ready(None) => {
147-
// Stream is exhausted, signal EOF by returning Ok(()) with no bytes written
148-
std::task::Poll::Ready(Ok(()))
128+
if buf.remaining() == 0 {
129+
return std::task::Poll::Ready(Ok(()));
130+
}
131+
132+
loop {
133+
match Stream::poll_next(self.bytes.as_mut(), cx) {
134+
std::task::Poll::Ready(Some(Ok(chunk))) => {
135+
if chunk.is_empty() {
136+
continue;
137+
}
138+
139+
let amt = std::cmp::min(chunk.len(), buf.remaining());
140+
buf.put_slice(&chunk[..amt]);
141+
142+
if amt < chunk.len() {
143+
let remainder = chunk.slice(amt..);
144+
let previous_stream =
145+
std::mem::replace(&mut self.bytes, Box::pin(stream::empty()));
146+
self.bytes = Box::pin(
147+
stream::once(async move { Ok(remainder) }).chain(previous_stream),
148+
);
149+
}
150+
151+
return std::task::Poll::Ready(Ok(()));
152+
}
153+
std::task::Poll::Ready(Some(Err(error))) => {
154+
return std::task::Poll::Ready(Err(std::io::Error::other(error)));
155+
}
156+
std::task::Poll::Ready(None) => {
157+
return std::task::Poll::Ready(Ok(()));
158+
}
159+
std::task::Poll::Pending => return std::task::Poll::Pending,
149160
}
150-
std::task::Poll::Pending => std::task::Poll::Pending,
151161
}
152162
}
153163
}
@@ -159,28 +169,39 @@ impl async_std::io::Read for ResponseDataStream {
159169
cx: &mut std::task::Context<'_>,
160170
buf: &mut [u8],
161171
) -> std::task::Poll<std::io::Result<usize>> {
162-
// Poll the stream for the next chunk of bytes
163-
match Stream::poll_next(self.bytes.as_mut(), cx) {
164-
std::task::Poll::Ready(Some(Ok(chunk))) => {
165-
// Write as much of the chunk as fits in the buffer
166-
let amt = std::cmp::min(chunk.len(), buf.len());
167-
buf[..amt].copy_from_slice(&chunk[..amt]);
168-
169-
// AIDEV-NOTE: Bytes that don't fit in the buffer are discarded from this chunk.
170-
// This is expected AsyncRead behavior - consumers should use appropriately sized
171-
// buffers or wrap in BufReader for efficiency with small reads.
172-
173-
std::task::Poll::Ready(Ok(amt))
174-
}
175-
std::task::Poll::Ready(Some(Err(error))) => {
176-
// Convert S3Error to io::Error
177-
std::task::Poll::Ready(Err(std::io::Error::other(error)))
178-
}
179-
std::task::Poll::Ready(None) => {
180-
// Stream is exhausted, signal EOF by returning 0 bytes read
181-
std::task::Poll::Ready(Ok(0))
172+
if buf.is_empty() {
173+
return std::task::Poll::Ready(Ok(0));
174+
}
175+
176+
loop {
177+
match Stream::poll_next(self.bytes.as_mut(), cx) {
178+
std::task::Poll::Ready(Some(Ok(chunk))) => {
179+
if chunk.is_empty() {
180+
continue;
181+
}
182+
183+
let amt = std::cmp::min(chunk.len(), buf.len());
184+
buf[..amt].copy_from_slice(&chunk[..amt]);
185+
186+
if amt < chunk.len() {
187+
let remainder = chunk.slice(amt..);
188+
let previous_stream =
189+
std::mem::replace(&mut self.bytes, Box::pin(stream::empty()));
190+
self.bytes = Box::pin(
191+
stream::once(async move { Ok(remainder) }).chain(previous_stream),
192+
);
193+
}
194+
195+
return std::task::Poll::Ready(Ok(amt));
196+
}
197+
std::task::Poll::Ready(Some(Err(error))) => {
198+
return std::task::Poll::Ready(Err(std::io::Error::other(error)));
199+
}
200+
std::task::Poll::Ready(None) => {
201+
return std::task::Poll::Ready(Ok(0));
202+
}
203+
std::task::Poll::Pending => return std::task::Poll::Pending,
182204
}
183-
std::task::Poll::Pending => std::task::Poll::Pending,
184205
}
185206
}
186207
}
@@ -875,11 +896,9 @@ mod tests {
875896
}
876897

877898
#[tokio::test]
878-
async fn test_async_read_with_small_buffer() {
879-
// Create a stream with a large chunk
880-
let chunks = vec![Ok(Bytes::from(
881-
"This is a much longer string that won't fit in a small buffer",
882-
))];
899+
async fn test_async_read_with_small_buffer_preserves_large_chunk() {
900+
let expected = b"This is a much longer string that won't fit in a small buffer";
901+
let chunks = vec![Ok(Bytes::copy_from_slice(expected))];
883902

884903
let stream = stream::iter(chunks);
885904
let data_stream: DataStream = Box::pin(stream);
@@ -889,17 +908,18 @@ mod tests {
889908
status_code: 200,
890909
};
891910

892-
// Read with a small buffer - demonstrates that excess bytes are discarded per chunk
911+
let mut output = Vec::new();
893912
let mut buffer = [0u8; 10];
894-
let n = response_stream.read(&mut buffer).await.unwrap();
895913

896-
// We should only get the first 10 bytes
897-
assert_eq!(n, 10);
898-
assert_eq!(&buffer[..n], b"This is a ");
914+
loop {
915+
let n = response_stream.read(&mut buffer).await.unwrap();
916+
if n == 0 {
917+
break;
918+
}
919+
output.extend_from_slice(&buffer[..n]);
920+
}
899921

900-
// Next read should get 0 bytes (EOF) because the chunk was consumed
901-
let n = response_stream.read(&mut buffer).await.unwrap();
902-
assert_eq!(n, 0);
922+
assert_eq!(output, expected);
903923
}
904924

905925
#[tokio::test]
@@ -992,11 +1012,9 @@ mod async_std_tests {
9921012
}
9931013

9941014
#[async_std::test]
995-
async fn test_async_read_with_small_buffer() {
996-
// Create a stream with a large chunk
997-
let chunks = vec![Ok(Bytes::from(
998-
"This is a much longer string that won't fit in a small buffer",
999-
))];
1015+
async fn test_async_read_with_small_buffer_preserves_large_chunk() {
1016+
let expected = b"This is a much longer string that won't fit in a small buffer";
1017+
let chunks = vec![Ok(Bytes::copy_from_slice(expected))];
10001018

10011019
let stream = stream::iter(chunks);
10021020
let data_stream: DataStream = Box::pin(stream);
@@ -1006,17 +1024,18 @@ mod async_std_tests {
10061024
status_code: 200,
10071025
};
10081026

1009-
// Read with a small buffer - demonstrates that excess bytes are discarded per chunk
1027+
let mut output = Vec::new();
10101028
let mut buffer = [0u8; 10];
1011-
let n = response_stream.read(&mut buffer).await.unwrap();
10121029

1013-
// We should only get the first 10 bytes
1014-
assert_eq!(n, 10);
1015-
assert_eq!(&buffer[..n], b"This is a ");
1030+
loop {
1031+
let n = response_stream.read(&mut buffer).await.unwrap();
1032+
if n == 0 {
1033+
break;
1034+
}
1035+
output.extend_from_slice(&buffer[..n]);
1036+
}
10161037

1017-
// Next read should get 0 bytes (EOF) because the chunk was consumed
1018-
let n = response_stream.read(&mut buffer).await.unwrap();
1019-
assert_eq!(n, 0);
1038+
assert_eq!(output, expected);
10201039
}
10211040

10221041
#[async_std::test]

0 commit comments

Comments
 (0)