add AsyncReaderStream
and replace AsyncIndexReader's stream implementation with that Signed-off-by: Dominik Csapak <d.csapak@proxmox.com>
This commit is contained in:
parent
3ddb14889a
commit
f386f512d0
@ -23,7 +23,7 @@ use crate::config::datastore;
|
||||
use crate::config::cached_user_info::CachedUserInfo;
|
||||
|
||||
use crate::server::WorkerTask;
|
||||
use crate::tools::{self, WrappedReaderStream};
|
||||
use crate::tools::{self, AsyncReaderStream, WrappedReaderStream};
|
||||
use crate::config::acl::{
|
||||
PRIV_DATASTORE_AUDIT,
|
||||
PRIV_DATASTORE_MODIFY,
|
||||
@ -842,7 +842,7 @@ fn download_file_decoded(
|
||||
|
||||
let chunk_reader = LocalChunkReader::new(datastore, None);
|
||||
let reader = AsyncIndexReader::new(index, chunk_reader);
|
||||
Body::wrap_stream(reader
|
||||
Body::wrap_stream(AsyncReaderStream::new(reader)
|
||||
.map_err(move |err| {
|
||||
eprintln!("error during streaming of '{:?}' - {}", path, err);
|
||||
err
|
||||
@ -854,7 +854,7 @@ fn download_file_decoded(
|
||||
|
||||
let chunk_reader = LocalChunkReader::new(datastore, None);
|
||||
let reader = AsyncIndexReader::new(index, chunk_reader);
|
||||
Body::wrap_stream(reader
|
||||
Body::wrap_stream(AsyncReaderStream::with_buffer_size(reader, 4*1024*1024)
|
||||
.map_err(move |err| {
|
||||
eprintln!("error during streaming of '{:?}' - {}", path, err);
|
||||
err
|
||||
|
@ -5,7 +5,7 @@ use std::pin::Pin;
|
||||
use anyhow::Error;
|
||||
use futures::future::FutureExt;
|
||||
use futures::ready;
|
||||
use tokio::{io::AsyncRead, stream::Stream};
|
||||
use tokio::io::AsyncRead;
|
||||
|
||||
use proxmox::sys::error::io_err_other;
|
||||
use proxmox::io_format_err;
|
||||
@ -125,62 +125,3 @@ I: IndexFile + Unpin
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, I> Stream for AsyncIndexReader<S, I>
|
||||
where
|
||||
S: AsyncReadChunk + Unpin + 'static,
|
||||
I: IndexFile + Unpin
|
||||
{
|
||||
type Item = Result<Vec<u8>, std::io::Error>;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
|
||||
let this = Pin::get_mut(self);
|
||||
loop {
|
||||
match &mut this.state {
|
||||
AsyncIndexReaderState::NoData => {
|
||||
if this.current_chunk_idx >= this.index.index_count() {
|
||||
return Poll::Ready(None);
|
||||
}
|
||||
|
||||
let digest = this
|
||||
.index
|
||||
.index_digest(this.current_chunk_idx)
|
||||
.ok_or(io_format_err!("could not get digest"))?
|
||||
.clone();
|
||||
|
||||
let mut store = match this.store.take() {
|
||||
Some(store) => store,
|
||||
None => {
|
||||
return Poll::Ready(Some(Err(io_format_err!("could not find store"))));
|
||||
},
|
||||
};
|
||||
|
||||
let future = async move {
|
||||
store.read_chunk(&digest)
|
||||
.await
|
||||
.map(move |x| (store, x))
|
||||
};
|
||||
|
||||
this.state = AsyncIndexReaderState::WaitForData(future.boxed());
|
||||
},
|
||||
AsyncIndexReaderState::WaitForData(ref mut future) => {
|
||||
match ready!(future.as_mut().poll(cx)) {
|
||||
Ok((store, chunk_data)) => {
|
||||
this.state = AsyncIndexReaderState::NoData;
|
||||
this.store = Some(store);
|
||||
this.current_chunk_idx += 1;
|
||||
return Poll::Ready(Some(Ok(chunk_data.clone())));
|
||||
},
|
||||
Err(err) => {
|
||||
return Poll::Ready(Some(Err(io_err_other(err))));
|
||||
},
|
||||
}
|
||||
},
|
||||
_ => {
|
||||
return Poll::Ready(Some(Err(io_format_err!("invalid state in stream"))));
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -3,7 +3,8 @@ use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
use std::sync::mpsc::Receiver;
|
||||
|
||||
|
||||
use tokio::io::AsyncRead;
|
||||
use futures::ready;
|
||||
use futures::stream::Stream;
|
||||
|
||||
use crate::tools::runtime::block_in_place;
|
||||
@ -42,6 +43,45 @@ impl<R: Read + Unpin> Stream for WrappedReaderStream<R> {
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrapper struct to convert an AsyncReader into a Stream
|
||||
pub struct AsyncReaderStream<R: AsyncRead + Unpin> {
|
||||
reader: R,
|
||||
buffer: Vec<u8>,
|
||||
}
|
||||
|
||||
impl <R: AsyncRead + Unpin> AsyncReaderStream<R> {
|
||||
|
||||
pub fn new(reader: R) -> Self {
|
||||
let mut buffer = Vec::with_capacity(64*1024);
|
||||
unsafe { buffer.set_len(buffer.capacity()); }
|
||||
Self { reader, buffer }
|
||||
}
|
||||
|
||||
pub fn with_buffer_size(reader: R, buffer_size: usize) -> Self {
|
||||
let mut buffer = Vec::with_capacity(buffer_size);
|
||||
unsafe { buffer.set_len(buffer.capacity()); }
|
||||
Self { reader, buffer }
|
||||
}
|
||||
}
|
||||
|
||||
impl<R: AsyncRead + Unpin> Stream for AsyncReaderStream<R> {
|
||||
type Item = Result<Vec<u8>, io::Error>;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
|
||||
let this = self.get_mut();
|
||||
match ready!(Pin::new(&mut this.reader).poll_read(cx, &mut this.buffer)) {
|
||||
Ok(n) => {
|
||||
if n == 0 {
|
||||
// EOF
|
||||
Poll::Ready(None)
|
||||
} else {
|
||||
Poll::Ready(Some(Ok(this.buffer[..n].to_vec())))
|
||||
}
|
||||
}
|
||||
Err(err) => Poll::Ready(Some(Err(err))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrapper struct to convert a channel Receiver into a Stream
|
||||
pub struct StdChannelStream<T>(pub Receiver<T>);
|
||||
|
Loading…
Reference in New Issue
Block a user