2020-06-18 11:55:27 +00:00
|
|
|
use std::future::Future;
|
|
|
|
use std::task::{Poll, Context};
|
|
|
|
use std::pin::Pin;
|
2020-07-22 13:56:22 +00:00
|
|
|
use std::io::SeekFrom;
|
2020-06-18 11:55:27 +00:00
|
|
|
|
|
|
|
use anyhow::Error;
|
|
|
|
use futures::future::FutureExt;
|
|
|
|
use futures::ready;
|
2020-12-04 08:34:08 +00:00
|
|
|
use tokio::io::{AsyncRead, AsyncSeek, ReadBuf};
|
2020-06-18 11:55:27 +00:00
|
|
|
|
|
|
|
use proxmox::sys::error::io_err_other;
|
|
|
|
use proxmox::io_format_err;
|
|
|
|
|
|
|
|
use super::IndexFile;
|
|
|
|
use super::read_chunk::AsyncReadChunk;
|
2020-07-22 13:56:22 +00:00
|
|
|
use super::index::ChunkReadInfo;
|
2020-06-18 11:55:27 +00:00
|
|
|
|
2021-01-25 13:42:50 +00:00
|
|
|
type ReadFuture<S> = dyn Future<Output = Result<(S, Vec<u8>), Error>> + Send + 'static;
|
|
|
|
|
2020-10-14 12:10:28 +00:00
|
|
|
// FIXME: This enum may not be required?
|
|
|
|
// - Put the `WaitForData` case directly into a `read_future: Option<>`
|
|
|
|
// - make the read loop as follows:
|
|
|
|
// * if read_buffer is not empty:
|
|
|
|
// use it
|
|
|
|
// * else if read_future is there:
|
|
|
|
// poll it
|
|
|
|
// if read: move data to read_buffer
|
|
|
|
// * else
|
|
|
|
// create read future
|
|
|
|
#[allow(clippy::enum_variant_names)]
|
2020-06-18 11:55:27 +00:00
|
|
|
enum AsyncIndexReaderState<S> {
|
|
|
|
NoData,
|
2021-01-25 13:42:50 +00:00
|
|
|
WaitForData(Pin<Box<ReadFuture<S>>>),
|
2020-07-22 13:56:22 +00:00
|
|
|
HaveData,
|
2020-06-18 11:55:27 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
pub struct AsyncIndexReader<S, I: IndexFile> {
|
|
|
|
store: Option<S>,
|
|
|
|
index: I,
|
|
|
|
read_buffer: Vec<u8>,
|
2020-07-22 13:56:22 +00:00
|
|
|
current_chunk_offset: u64,
|
2020-06-18 11:55:27 +00:00
|
|
|
current_chunk_idx: usize,
|
2020-07-22 13:56:22 +00:00
|
|
|
current_chunk_info: Option<ChunkReadInfo>,
|
|
|
|
position: u64,
|
|
|
|
seek_to_pos: i64,
|
2020-06-18 11:55:27 +00:00
|
|
|
state: AsyncIndexReaderState<S>,
|
|
|
|
}
|
|
|
|
|
|
|
|
// ok because the only public interfaces operates on &mut Self
|
|
|
|
unsafe impl<S: Sync, I: IndexFile + Sync> Sync for AsyncIndexReader<S, I> {}
|
|
|
|
|
|
|
|
impl<S: AsyncReadChunk, I: IndexFile> AsyncIndexReader<S, I> {
|
|
|
|
pub fn new(index: I, store: S) -> Self {
|
|
|
|
Self {
|
|
|
|
store: Some(store),
|
|
|
|
index,
|
2020-07-03 07:45:25 +00:00
|
|
|
read_buffer: Vec::with_capacity(1024 * 1024),
|
2020-07-22 13:56:22 +00:00
|
|
|
current_chunk_offset: 0,
|
2020-06-18 11:55:27 +00:00
|
|
|
current_chunk_idx: 0,
|
2020-07-22 13:56:22 +00:00
|
|
|
current_chunk_info: None,
|
|
|
|
position: 0,
|
|
|
|
seek_to_pos: 0,
|
2020-06-18 11:55:27 +00:00
|
|
|
state: AsyncIndexReaderState::NoData,
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2020-07-03 07:45:25 +00:00
|
|
|
impl<S, I> AsyncRead for AsyncIndexReader<S, I>
|
|
|
|
where
|
|
|
|
S: AsyncReadChunk + Unpin + Sync + 'static,
|
|
|
|
I: IndexFile + Unpin,
|
2020-06-18 11:55:27 +00:00
|
|
|
{
|
|
|
|
fn poll_read(
|
|
|
|
self: Pin<&mut Self>,
|
|
|
|
cx: &mut Context,
|
2020-12-04 08:34:08 +00:00
|
|
|
buf: &mut ReadBuf,
|
|
|
|
) -> Poll<tokio::io::Result<()>> {
|
2020-06-18 11:55:27 +00:00
|
|
|
let this = Pin::get_mut(self);
|
|
|
|
loop {
|
|
|
|
match &mut this.state {
|
|
|
|
AsyncIndexReaderState::NoData => {
|
2020-07-22 13:56:22 +00:00
|
|
|
let (idx, offset) = if this.current_chunk_info.is_some() &&
|
|
|
|
this.position == this.current_chunk_info.as_ref().unwrap().range.end
|
|
|
|
{
|
|
|
|
// optimization for sequential chunk read
|
|
|
|
let next_idx = this.current_chunk_idx + 1;
|
|
|
|
(next_idx, 0)
|
|
|
|
} else {
|
|
|
|
match this.index.chunk_from_offset(this.position) {
|
|
|
|
Some(res) => res,
|
2020-12-04 08:34:08 +00:00
|
|
|
None => return Poll::Ready(Ok(()))
|
2020-07-22 13:56:22 +00:00
|
|
|
}
|
|
|
|
};
|
|
|
|
|
|
|
|
if idx >= this.index.index_count() {
|
2020-12-04 08:34:08 +00:00
|
|
|
return Poll::Ready(Ok(()));
|
2020-06-18 11:55:27 +00:00
|
|
|
}
|
|
|
|
|
2020-07-22 13:56:22 +00:00
|
|
|
let info = this
|
2020-06-18 11:55:27 +00:00
|
|
|
.index
|
2020-07-22 13:56:22 +00:00
|
|
|
.chunk_info(idx)
|
2021-01-19 13:04:46 +00:00
|
|
|
.ok_or_else(|| io_format_err!("could not get digest"))?;
|
2020-07-22 13:56:22 +00:00
|
|
|
|
|
|
|
this.current_chunk_offset = offset;
|
|
|
|
this.current_chunk_idx = idx;
|
|
|
|
let old_info = this.current_chunk_info.replace(info.clone());
|
|
|
|
|
|
|
|
if let Some(old_info) = old_info {
|
|
|
|
if old_info.digest == info.digest {
|
|
|
|
// hit, chunk is currently in cache
|
|
|
|
this.state = AsyncIndexReaderState::HaveData;
|
|
|
|
continue;
|
|
|
|
}
|
2020-06-18 11:55:27 +00:00
|
|
|
}
|
|
|
|
|
2020-07-22 13:56:22 +00:00
|
|
|
// miss, need to download new chunk
|
2020-07-03 05:36:23 +00:00
|
|
|
let store = match this.store.take() {
|
2020-06-18 11:55:27 +00:00
|
|
|
Some(store) => store,
|
|
|
|
None => {
|
|
|
|
return Poll::Ready(Err(io_format_err!("could not find store")));
|
2020-07-03 07:45:25 +00:00
|
|
|
}
|
2020-06-18 11:55:27 +00:00
|
|
|
};
|
|
|
|
|
|
|
|
let future = async move {
|
2020-07-22 13:56:22 +00:00
|
|
|
store.read_chunk(&info.digest)
|
2020-06-18 11:55:27 +00:00
|
|
|
.await
|
|
|
|
.map(move |x| (store, x))
|
|
|
|
};
|
|
|
|
|
|
|
|
this.state = AsyncIndexReaderState::WaitForData(future.boxed());
|
2020-07-03 07:45:25 +00:00
|
|
|
}
|
2020-06-18 11:55:27 +00:00
|
|
|
AsyncIndexReaderState::WaitForData(ref mut future) => {
|
|
|
|
match ready!(future.as_mut().poll(cx)) {
|
2020-10-14 12:10:28 +00:00
|
|
|
Ok((store, chunk_data)) => {
|
|
|
|
this.read_buffer = chunk_data;
|
2020-07-22 13:56:22 +00:00
|
|
|
this.state = AsyncIndexReaderState::HaveData;
|
2020-06-18 11:55:27 +00:00
|
|
|
this.store = Some(store);
|
2020-07-03 07:45:25 +00:00
|
|
|
}
|
2020-06-18 11:55:27 +00:00
|
|
|
Err(err) => {
|
|
|
|
return Poll::Ready(Err(io_err_other(err)));
|
2020-07-03 07:45:25 +00:00
|
|
|
}
|
2020-06-18 11:55:27 +00:00
|
|
|
};
|
2020-07-03 07:45:25 +00:00
|
|
|
}
|
2020-07-22 13:56:22 +00:00
|
|
|
AsyncIndexReaderState::HaveData => {
|
|
|
|
let offset = this.current_chunk_offset as usize;
|
2020-06-18 11:55:27 +00:00
|
|
|
let len = this.read_buffer.len();
|
2020-12-04 08:34:08 +00:00
|
|
|
let n = if len - offset < buf.remaining() {
|
2020-06-18 11:55:27 +00:00
|
|
|
len - offset
|
|
|
|
} else {
|
2020-12-04 08:34:08 +00:00
|
|
|
buf.remaining()
|
2020-06-18 11:55:27 +00:00
|
|
|
};
|
|
|
|
|
2020-12-04 08:34:08 +00:00
|
|
|
buf.put_slice(&this.read_buffer[offset..(offset + n)]);
|
2020-07-22 13:56:22 +00:00
|
|
|
this.position += n as u64;
|
|
|
|
|
2020-06-18 11:55:27 +00:00
|
|
|
if offset + n == len {
|
|
|
|
this.state = AsyncIndexReaderState::NoData;
|
|
|
|
} else {
|
2020-07-22 13:56:22 +00:00
|
|
|
this.current_chunk_offset += n as u64;
|
|
|
|
this.state = AsyncIndexReaderState::HaveData;
|
2020-06-18 11:55:27 +00:00
|
|
|
}
|
|
|
|
|
2020-12-04 08:34:08 +00:00
|
|
|
return Poll::Ready(Ok(()));
|
2020-07-03 07:45:25 +00:00
|
|
|
}
|
2020-06-18 11:55:27 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
2020-07-22 13:56:22 +00:00
|
|
|
|
|
|
|
impl<S, I> AsyncSeek for AsyncIndexReader<S, I>
|
|
|
|
where
|
|
|
|
S: AsyncReadChunk + Unpin + Sync + 'static,
|
|
|
|
I: IndexFile + Unpin,
|
|
|
|
{
|
|
|
|
fn start_seek(
|
|
|
|
self: Pin<&mut Self>,
|
|
|
|
pos: SeekFrom,
|
2020-12-04 08:34:08 +00:00
|
|
|
) -> tokio::io::Result<()> {
|
2020-07-22 13:56:22 +00:00
|
|
|
let this = Pin::get_mut(self);
|
|
|
|
this.seek_to_pos = match pos {
|
|
|
|
SeekFrom::Start(offset) => {
|
|
|
|
offset as i64
|
|
|
|
},
|
|
|
|
SeekFrom::End(offset) => {
|
|
|
|
this.index.index_bytes() as i64 + offset
|
|
|
|
},
|
|
|
|
SeekFrom::Current(offset) => {
|
|
|
|
this.position as i64 + offset
|
|
|
|
}
|
|
|
|
};
|
2020-12-04 08:34:08 +00:00
|
|
|
Ok(())
|
2020-07-22 13:56:22 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
fn poll_complete(
|
|
|
|
self: Pin<&mut Self>,
|
|
|
|
_cx: &mut Context<'_>,
|
|
|
|
) -> Poll<tokio::io::Result<u64>> {
|
|
|
|
let this = Pin::get_mut(self);
|
|
|
|
|
|
|
|
let index_bytes = this.index.index_bytes();
|
|
|
|
if this.seek_to_pos < 0 {
|
|
|
|
return Poll::Ready(Err(io_format_err!("cannot seek to negative values")));
|
|
|
|
} else if this.seek_to_pos > index_bytes as i64 {
|
|
|
|
this.position = index_bytes;
|
|
|
|
} else {
|
|
|
|
this.position = this.seek_to_pos as u64;
|
|
|
|
}
|
|
|
|
|
|
|
|
// even if seeking within one chunk, we need to go to NoData to
|
|
|
|
// recalculate the current_chunk_offset (data is cached anyway)
|
|
|
|
this.state = AsyncIndexReaderState::NoData;
|
|
|
|
|
|
|
|
Poll::Ready(Ok(this.position))
|
|
|
|
}
|
|
|
|
}
|