Spaces:
Runtime error
Runtime error
File size: 2,852 Bytes
9552aa0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{self, AsyncRead, ReadBuf};
pub struct ReadToEndFuture<R> {
reader: R,
buffer: Vec<u8>,
}
impl<R> ReadToEndFuture<R> {
pub fn new(reader: R) -> Self {
Self {
reader,
buffer: Vec::new(),
}
}
}
impl<R> Future for ReadToEndFuture<R>
where
R: AsyncRead + Unpin,
{
type Output = io::Result<Vec<u8>>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut buf = [0; 1024];
let mut read_buf = ReadBuf::new(&mut buf);
loop {
match Pin::new(&mut self.reader).poll_read(cx, &mut read_buf) {
Poll::Ready(Ok(())) => {
let n = read_buf.filled().len();
if n == 0 {
return Poll::Ready(Ok(self.buffer.clone()));
}
self.buffer.extend_from_slice(read_buf.filled());
read_buf.clear();
}
Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
Poll::Pending => return Poll::Pending,
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{self, AsyncRead};
struct MockReader {
data: Vec<u8>,
position: usize,
}
impl MockReader {
fn new(data: &[u8]) -> Self {
Self {
data: data.to_vec(),
position: 0,
}
}
}
impl AsyncRead for MockReader {
fn poll_read(
mut self: Pin<&mut Self>,
_cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
if self.position >= self.data.len() {
return Poll::Ready(Ok(()));
}
let end = (self.position + buf.remaining()).min(self.data.len());
buf.put_slice(&self.data[self.position..end]);
self.position = end;
Poll::Ready(Ok(()))
}
}
#[tokio::test]
async fn test_read_to_end_empty_reader() {
let reader = MockReader::new(&[]);
let future = ReadToEndFuture::new(reader);
let result = future.await;
assert_eq!(result.unwrap(), Vec::<u8>::new());
}
#[tokio::test]
async fn test_read_to_end_non_empty_reader() {
let reader = MockReader::new(b"hello world");
let future = ReadToEndFuture::new(reader);
let result = future.await;
assert_eq!(result.unwrap(), b"hello world");
}
struct ErrorReader;
impl AsyncRead for ErrorReader {
fn poll_read(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
_buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
Poll::Ready(Err(io::Error::other("read error")))
}
}
#[tokio::test]
async fn test_read_to_end_error() {
let reader = ErrorReader;
let future = ReadToEndFuture::new(reader);
let result = future.await;
assert!(result.is_err());
assert_eq!(result.unwrap_err().kind(), io::ErrorKind::Other);
}
}
|