Skip to content

Commit 7bce8f9

Browse files
durchdurch
authored
Merge pull request #457 from durch/fix/issue-437-response-stream-read
Fix ResponseDataStream small-buffer reads
2 parents 7017729 + d106d33 commit 7bce8f9

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
}
@@ -876,11 +897,9 @@ mod tests {
876897
}
877898

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

885904
let stream = stream::iter(chunks);
886905
let data_stream: DataStream = Box::pin(stream);
@@ -890,17 +909,18 @@ mod tests {
890909
status_code: 200,
891910
};
892911

893-
// Read with a small buffer - demonstrates that excess bytes are discarded per chunk
912+
let mut output = Vec::new();
894913
let mut buffer = [0u8; 10];
895-
let n = response_stream.read(&mut buffer).await.unwrap();
896914

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

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

906926
#[tokio::test]
@@ -993,11 +1013,9 @@ mod async_std_tests {
9931013
}
9941014

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

10021020
let stream = stream::iter(chunks);
10031021
let data_stream: DataStream = Box::pin(stream);
@@ -1007,17 +1025,18 @@ mod async_std_tests {
10071025
status_code: 200,
10081026
};
10091027

1010-
// Read with a small buffer - demonstrates that excess bytes are discarded per chunk
1028+
let mut output = Vec::new();
10111029
let mut buffer = [0u8; 10];
1012-
let n = response_stream.read(&mut buffer).await.unwrap();
10131030

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

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

10231042
#[async_std::test]

0 commit comments

Comments
 (0)