move channel/stream helpers to pbs-tools
pbs_tools ::blocking: std/async wrapping with block_in_place ::stream: stream <-> AsyncRead/AsyncWrite wrapping Signed-off-by: Wolfgang Bumiller <w.bumiller@proxmox.com>
This commit is contained in:
Wolfgang Bumiller
@ -32,6 +32,7 @@ walkdir = "2"
|
||||
proxmox = { version = "0.11.5", default-features = false, features = [ "tokio" ] }
|
||||
|
||||
pbs-buildcfg = { path = "../pbs-buildcfg" }
|
||||
pbs-runtime = { path = "../pbs-runtime" }
|
||||
|
||||
[dev-dependencies]
|
||||
tokio = { version = "1.6", features = [ "macros" ] }
|
||||
|
||||
99
pbs-tools/src/blocking.rs
Normal file
99
pbs-tools/src/blocking.rs
Normal file
@ -0,0 +1,99 @@
|
||||
//! Async wrappers for blocking I/O (adding `block_in_place` around channels/readers)
|
||||
|
||||
use std::io::{self, Read};
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
use std::sync::mpsc::Receiver;
|
||||
|
||||
use futures::stream::Stream;
|
||||
|
||||
use pbs_runtime::block_in_place;
|
||||
|
||||
/// Wrapper struct to convert a Reader into a Stream
|
||||
pub struct WrappedReaderStream<R: Read + Unpin> {
|
||||
reader: R,
|
||||
buffer: Vec<u8>,
|
||||
}
|
||||
|
||||
impl <R: Read + Unpin> WrappedReaderStream<R> {
|
||||
|
||||
pub fn new(reader: R) -> Self {
|
||||
let mut buffer = Vec::with_capacity(64*1024);
|
||||
unsafe { buffer.set_len(buffer.capacity()); }
|
||||
Self { reader, buffer }
|
||||
}
|
||||
}
|
||||
|
||||
impl<R: Read + Unpin> Stream for WrappedReaderStream<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 block_in_place(|| this.reader.read(&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>);
|
||||
|
||||
impl<T> Stream for StdChannelStream<T> {
|
||||
type Item = T;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Option<Self::Item>> {
|
||||
match block_in_place(|| self.0.recv()) {
|
||||
Ok(data) => Poll::Ready(Some(data)),
|
||||
Err(_) => Poll::Ready(None),// channel closed
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use std::io;
|
||||
|
||||
use anyhow::Error;
|
||||
use futures::stream::TryStreamExt;
|
||||
|
||||
#[test]
|
||||
fn test_wrapped_stream_reader() -> Result<(), Error> {
|
||||
pbs_runtime::main(async {
|
||||
run_wrapped_stream_reader_test().await
|
||||
})
|
||||
}
|
||||
|
||||
struct DummyReader(usize);
|
||||
|
||||
impl io::Read for DummyReader {
|
||||
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
|
||||
self.0 += 1;
|
||||
|
||||
if self.0 >= 10 {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
unsafe {
|
||||
std::ptr::write_bytes(buf.as_mut_ptr(), 0, buf.len());
|
||||
}
|
||||
|
||||
Ok(buf.len())
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_wrapped_stream_reader_test() -> Result<(), Error> {
|
||||
let mut reader = super::WrappedReaderStream::new(DummyReader(0));
|
||||
while let Some(_data) = reader.try_next().await? {
|
||||
// just waiting
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@ -1,5 +1,6 @@
|
||||
pub mod acl;
|
||||
pub mod auth;
|
||||
pub mod blocking;
|
||||
pub mod borrow;
|
||||
pub mod broadcast_future;
|
||||
pub mod cert;
|
||||
@ -7,12 +8,14 @@ pub mod compression;
|
||||
pub mod format;
|
||||
pub mod fs;
|
||||
pub mod json;
|
||||
pub mod lru_cache;
|
||||
pub mod nom;
|
||||
pub mod ops;
|
||||
pub mod percent_encoding;
|
||||
pub mod process_locker;
|
||||
pub mod sha;
|
||||
pub mod str;
|
||||
pub mod stream;
|
||||
pub mod sync;
|
||||
pub mod ticket;
|
||||
pub mod tokio;
|
||||
@ -20,7 +23,6 @@ pub mod xattr;
|
||||
pub mod zip;
|
||||
|
||||
pub mod async_lru_cache;
|
||||
pub mod lru_cache;
|
||||
|
||||
mod command;
|
||||
pub use command::{command_output, command_output_as_string, run_command};
|
||||
|
||||
229
pbs-tools/src/stream.rs
Normal file
229
pbs-tools/src/stream.rs
Normal file
@ -0,0 +1,229 @@
|
||||
//! Wrappers between async readers and streams.
|
||||
|
||||
use std::io::{self, Read};
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
use anyhow::{Error, Result};
|
||||
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
|
||||
use tokio::sync::mpsc::Sender;
|
||||
use futures::ready;
|
||||
use futures::future::FutureExt;
|
||||
use futures::stream::Stream;
|
||||
|
||||
use proxmox::io_format_err;
|
||||
use proxmox::tools::byte_buffer::ByteBuffer;
|
||||
use proxmox::sys::error::io_err_other;
|
||||
|
||||
use pbs_runtime::block_in_place;
|
||||
|
||||
/// Wrapper struct to convert a Reader into a Stream
|
||||
pub struct WrappedReaderStream<R: Read + Unpin> {
|
||||
reader: R,
|
||||
buffer: Vec<u8>,
|
||||
}
|
||||
|
||||
impl <R: Read + Unpin> WrappedReaderStream<R> {
|
||||
|
||||
pub fn new(reader: R) -> Self {
|
||||
let mut buffer = Vec::with_capacity(64*1024);
|
||||
unsafe { buffer.set_len(buffer.capacity()); }
|
||||
Self { reader, buffer }
|
||||
}
|
||||
}
|
||||
|
||||
impl<R: Read + Unpin> Stream for WrappedReaderStream<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 block_in_place(|| this.reader.read(&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 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();
|
||||
let mut read_buf = ReadBuf::new(&mut this.buffer);
|
||||
match ready!(Pin::new(&mut this.reader).poll_read(cx, &mut read_buf)) {
|
||||
Ok(()) => {
|
||||
let n = read_buf.filled().len();
|
||||
if n == 0 {
|
||||
// EOF
|
||||
Poll::Ready(None)
|
||||
} else {
|
||||
Poll::Ready(Some(Ok(this.buffer[..n].to_vec())))
|
||||
}
|
||||
}
|
||||
Err(err) => Poll::Ready(Some(Err(err))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use std::io;
|
||||
|
||||
use anyhow::Error;
|
||||
use futures::stream::TryStreamExt;
|
||||
|
||||
#[test]
|
||||
fn test_wrapped_stream_reader() -> Result<(), Error> {
|
||||
pbs_runtime::main(async {
|
||||
run_wrapped_stream_reader_test().await
|
||||
})
|
||||
}
|
||||
|
||||
struct DummyReader(usize);
|
||||
|
||||
impl io::Read for DummyReader {
|
||||
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
|
||||
self.0 += 1;
|
||||
|
||||
if self.0 >= 10 {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
unsafe {
|
||||
std::ptr::write_bytes(buf.as_mut_ptr(), 0, buf.len());
|
||||
}
|
||||
|
||||
Ok(buf.len())
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_wrapped_stream_reader_test() -> Result<(), Error> {
|
||||
let mut reader = super::WrappedReaderStream::new(DummyReader(0));
|
||||
while let Some(_data) = reader.try_next().await? {
|
||||
// just waiting
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrapper around tokio::sync::mpsc::Sender, which implements Write
|
||||
pub struct AsyncChannelWriter {
|
||||
sender: Option<Sender<Result<Vec<u8>, Error>>>,
|
||||
buf: ByteBuffer,
|
||||
state: WriterState,
|
||||
}
|
||||
|
||||
type SendResult = io::Result<Sender<Result<Vec<u8>>>>;
|
||||
|
||||
enum WriterState {
|
||||
Ready,
|
||||
Sending(Pin<Box<dyn Future<Output = SendResult> + Send + 'static>>),
|
||||
}
|
||||
|
||||
impl AsyncChannelWriter {
|
||||
pub fn new(sender: Sender<Result<Vec<u8>, Error>>, buf_size: usize) -> Self {
|
||||
Self {
|
||||
sender: Some(sender),
|
||||
buf: ByteBuffer::with_capacity(buf_size),
|
||||
state: WriterState::Ready,
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_write_impl(
|
||||
&mut self,
|
||||
cx: &mut Context,
|
||||
buf: &[u8],
|
||||
flush: bool,
|
||||
) -> Poll<io::Result<usize>> {
|
||||
loop {
|
||||
match &mut self.state {
|
||||
WriterState::Ready => {
|
||||
if flush {
|
||||
if self.buf.is_empty() {
|
||||
return Poll::Ready(Ok(0));
|
||||
}
|
||||
} else {
|
||||
let free_size = self.buf.free_size();
|
||||
if free_size > buf.len() || self.buf.is_empty() {
|
||||
let count = free_size.min(buf.len());
|
||||
self.buf.get_free_mut_slice()[..count].copy_from_slice(&buf[..count]);
|
||||
self.buf.add_size(count);
|
||||
return Poll::Ready(Ok(count));
|
||||
}
|
||||
}
|
||||
|
||||
let sender = match self.sender.take() {
|
||||
Some(sender) => sender,
|
||||
None => return Poll::Ready(Err(io_err_other("no sender"))),
|
||||
};
|
||||
|
||||
let data = self.buf.remove_data(self.buf.len()).to_vec();
|
||||
let future = async move {
|
||||
sender
|
||||
.send(Ok(data))
|
||||
.await
|
||||
.map(move |_| sender)
|
||||
.map_err(|err| io_format_err!("could not send: {}", err))
|
||||
};
|
||||
|
||||
self.state = WriterState::Sending(future.boxed());
|
||||
}
|
||||
WriterState::Sending(ref mut future) => match ready!(future.as_mut().poll(cx)) {
|
||||
Ok(sender) => {
|
||||
self.sender = Some(sender);
|
||||
self.state = WriterState::Ready;
|
||||
}
|
||||
Err(err) => return Poll::Ready(Err(err)),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AsyncWrite for AsyncChannelWriter {
|
||||
fn poll_write(self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<io::Result<usize>> {
|
||||
let this = self.get_mut();
|
||||
this.poll_write_impl(cx, buf, false)
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
|
||||
let this = self.get_mut();
|
||||
match ready!(this.poll_write_impl(cx, &[], true)) {
|
||||
Ok(_) => Poll::Ready(Ok(())),
|
||||
Err(err) => Poll::Ready(Err(err)),
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
|
||||
self.poll_flush(cx)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user