tokio 1.0: AsyncRead/Seek with ReadBuf

Signed-off-by: Fabian Grünbichler <f.gruenbichler@proxmox.com>
This commit is contained in:
Fabian Grünbichler 2020-12-04 09:34:08 +01:00
parent d148958b67
commit 484172b5f8
3 changed files with 19 additions and 62 deletions

View File

@ -6,7 +6,7 @@ use std::io::SeekFrom;
use anyhow::Error; use anyhow::Error;
use futures::future::FutureExt; use futures::future::FutureExt;
use futures::ready; use futures::ready;
use tokio::io::{AsyncRead, AsyncSeek}; use tokio::io::{AsyncRead, AsyncSeek, ReadBuf};
use proxmox::sys::error::io_err_other; use proxmox::sys::error::io_err_other;
use proxmox::io_format_err; use proxmox::io_format_err;
@ -71,8 +71,8 @@ where
fn poll_read( fn poll_read(
self: Pin<&mut Self>, self: Pin<&mut Self>,
cx: &mut Context, cx: &mut Context,
buf: &mut [u8], buf: &mut ReadBuf,
) -> Poll<tokio::io::Result<usize>> { ) -> Poll<tokio::io::Result<()>> {
let this = Pin::get_mut(self); let this = Pin::get_mut(self);
loop { loop {
match &mut this.state { match &mut this.state {
@ -86,12 +86,12 @@ where
} else { } else {
match this.index.chunk_from_offset(this.position) { match this.index.chunk_from_offset(this.position) {
Some(res) => res, Some(res) => res,
None => return Poll::Ready(Ok(0)) None => return Poll::Ready(Ok(()))
} }
}; };
if idx >= this.index.index_count() { if idx >= this.index.index_count() {
return Poll::Ready(Ok(0)); return Poll::Ready(Ok(()));
} }
let info = this let info = this
@ -142,13 +142,13 @@ where
AsyncIndexReaderState::HaveData => { AsyncIndexReaderState::HaveData => {
let offset = this.current_chunk_offset as usize; let offset = this.current_chunk_offset as usize;
let len = this.read_buffer.len(); let len = this.read_buffer.len();
let n = if len - offset < buf.len() { let n = if len - offset < buf.remaining() {
len - offset len - offset
} else { } else {
buf.len() buf.remaining()
}; };
buf[0..n].copy_from_slice(&this.read_buffer[offset..(offset + n)]); buf.put_slice(&this.read_buffer[offset..(offset + n)]);
this.position += n as u64; this.position += n as u64;
if offset + n == len { if offset + n == len {
@ -158,7 +158,7 @@ where
this.state = AsyncIndexReaderState::HaveData; this.state = AsyncIndexReaderState::HaveData;
} }
return Poll::Ready(Ok(n)); return Poll::Ready(Ok(()));
} }
} }
} }
@ -172,9 +172,8 @@ where
{ {
fn start_seek( fn start_seek(
self: Pin<&mut Self>, self: Pin<&mut Self>,
_cx: &mut Context<'_>,
pos: SeekFrom, pos: SeekFrom,
) -> Poll<tokio::io::Result<()>> { ) -> tokio::io::Result<()> {
let this = Pin::get_mut(self); let this = Pin::get_mut(self);
this.seek_to_pos = match pos { this.seek_to_pos = match pos {
SeekFrom::Start(offset) => { SeekFrom::Start(offset) => {
@ -187,7 +186,7 @@ where
this.position as i64 + offset this.position as i64 + offset
} }
}; };
Poll::Ready(Ok(())) Ok(())
} }
fn poll_complete( fn poll_complete(

View File

@ -1,13 +1,12 @@
//! Generic AsyncRead/AsyncWrite utilities. //! Generic AsyncRead/AsyncWrite utilities.
use std::io; use std::io;
use std::mem::MaybeUninit;
use std::os::unix::io::{AsRawFd, RawFd}; use std::os::unix::io::{AsRawFd, RawFd};
use std::pin::Pin; use std::pin::Pin;
use std::task::{Context, Poll}; use std::task::{Context, Poll};
use futures::stream::{Stream, TryStream}; use futures::stream::{Stream, TryStream};
use tokio::io::{AsyncRead, AsyncWrite}; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio::net::TcpListener; use tokio::net::TcpListener;
use hyper::client::connect::Connection; use hyper::client::connect::Connection;
@ -20,8 +19,8 @@ impl<L: AsyncRead + Unpin, R: AsyncRead + Unpin> AsyncRead for EitherStream<L, R
fn poll_read( fn poll_read(
self: Pin<&mut Self>, self: Pin<&mut Self>,
cx: &mut Context, cx: &mut Context,
buf: &mut [u8], buf: &mut ReadBuf,
) -> Poll<Result<usize, io::Error>> { ) -> Poll<Result<(), io::Error>> {
match self.get_mut() { match self.get_mut() {
EitherStream::Left(ref mut s) => { EitherStream::Left(ref mut s) => {
Pin::new(s).poll_read(cx, buf) Pin::new(s).poll_read(cx, buf)
@ -31,31 +30,6 @@ impl<L: AsyncRead + Unpin, R: AsyncRead + Unpin> AsyncRead for EitherStream<L, R
} }
} }
} }
unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [MaybeUninit<u8>]) -> bool {
match *self {
EitherStream::Left(ref s) => s.prepare_uninitialized_buffer(buf),
EitherStream::Right(ref s) => s.prepare_uninitialized_buffer(buf),
}
}
fn poll_read_buf<B>(
self: Pin<&mut Self>,
cx: &mut Context,
buf: &mut B,
) -> Poll<Result<usize, io::Error>>
where
B: bytes::BufMut,
{
match self.get_mut() {
EitherStream::Left(ref mut s) => {
Pin::new(s).poll_read_buf(cx, buf)
}
EitherStream::Right(ref mut s) => {
Pin::new(s).poll_read_buf(cx, buf)
}
}
}
} }
impl<L: AsyncWrite + Unpin, R: AsyncWrite + Unpin> AsyncWrite for EitherStream<L, R> { impl<L: AsyncWrite + Unpin, R: AsyncWrite + Unpin> AsyncWrite for EitherStream<L, R> {
@ -95,24 +69,6 @@ impl<L: AsyncWrite + Unpin, R: AsyncWrite + Unpin> AsyncWrite for EitherStream<L
} }
} }
} }
fn poll_write_buf<B>(
self: Pin<&mut Self>,
cx: &mut Context,
buf: &mut B,
) -> Poll<Result<usize, io::Error>>
where
B: bytes::Buf,
{
match self.get_mut() {
EitherStream::Left(ref mut s) => {
Pin::new(s).poll_write_buf(cx, buf)
}
EitherStream::Right(ref mut s) => {
Pin::new(s).poll_write_buf(cx, buf)
}
}
}
} }
// we need this for crate::client::http_client: // we need this for crate::client::http_client:

View File

@ -3,7 +3,7 @@ use std::pin::Pin;
use std::task::{Context, Poll}; use std::task::{Context, Poll};
use std::sync::mpsc::Receiver; use std::sync::mpsc::Receiver;
use tokio::io::AsyncRead; use tokio::io::{AsyncRead, ReadBuf};
use futures::ready; use futures::ready;
use futures::stream::Stream; use futures::stream::Stream;
@ -69,8 +69,10 @@ impl<R: AsyncRead + Unpin> Stream for AsyncReaderStream<R> {
fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> { fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
let this = self.get_mut(); let this = self.get_mut();
match ready!(Pin::new(&mut this.reader).poll_read(cx, &mut this.buffer)) { let mut read_buf = ReadBuf::new(&mut this.buffer);
Ok(n) => { match ready!(Pin::new(&mut this.reader).poll_read(cx, &mut read_buf)) {
Ok(()) => {
let n = read_buf.filled().len();
if n == 0 { if n == 0 {
// EOF // EOF
Poll::Ready(None) Poll::Ready(None)