forked from Proxmox/proxmox
proxmox-compression: add streaming zstd encoder
similar to our DeflateEncoder, takes a Stream and implements it itself, so that we can use it as an adapter for async api calls Signed-off-by: Dominik Csapak <d.csapak@proxmox.com> Signed-off-by: Wolfgang Bumiller <w.bumiller@proxmox.com>
This commit is contained in:
parent
d107d0b2eb
commit
e7e4411f44
@ -18,6 +18,7 @@ futures = "0.3"
|
||||
tokio = { version = "1.6", features = [ "fs", "io-util"] }
|
||||
walkdir = "2"
|
||||
tar = "0.4"
|
||||
zstd = { version = "0.6", features = []}
|
||||
|
||||
proxmox-time = { path = "../proxmox-time", version = "1" }
|
||||
proxmox-io = { path = "../proxmox-io", version = "1", features = [ "tokio" ] }
|
||||
|
@ -3,3 +3,4 @@ pub use compression::*;
|
||||
|
||||
pub mod tar;
|
||||
pub mod zip;
|
||||
pub mod zstd;
|
||||
|
126
proxmox-compression/src/zstd.rs
Normal file
126
proxmox-compression/src/zstd.rs
Normal file
@ -0,0 +1,126 @@
|
||||
//! zstd helper
|
||||
use std::io;
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
use anyhow::{format_err, Error};
|
||||
use bytes::Bytes;
|
||||
use futures::ready;
|
||||
use futures::stream::Stream;
|
||||
use zstd::stream::raw::{Encoder, Operation, OutBuffer};
|
||||
|
||||
use proxmox_io::ByteBuffer;
|
||||
|
||||
const BUFFER_SIZE: usize = 8192;
|
||||
|
||||
#[derive(Eq, PartialEq)]
|
||||
enum EncoderState {
|
||||
Reading,
|
||||
Writing,
|
||||
Finishing,
|
||||
Finished,
|
||||
}
|
||||
|
||||
/// An async ZstdEncoder that implements [Stream] for another [Stream]
|
||||
///
|
||||
/// Useful for on-the-fly zstd compression in streaming api calls
|
||||
pub struct ZstdEncoder<'a, T> {
|
||||
inner: T,
|
||||
compressor: Encoder<'a>,
|
||||
buffer: ByteBuffer,
|
||||
input_buffer: Bytes,
|
||||
state: EncoderState,
|
||||
}
|
||||
|
||||
impl<'a, T> ZstdEncoder<'a, T> {
|
||||
/// Returns a new [ZstdEncoder] with default level 3
|
||||
pub fn new(inner: T) -> Result<Self, io::Error> {
|
||||
Self::with_quality(inner, 3)
|
||||
}
|
||||
|
||||
/// Returns a new [ZstdEncoder] with the given level
|
||||
pub fn with_quality(inner: T, level: i32) -> Result<Self, io::Error> {
|
||||
Ok(Self {
|
||||
inner,
|
||||
compressor: Encoder::new(level)?,
|
||||
buffer: ByteBuffer::with_capacity(BUFFER_SIZE),
|
||||
input_buffer: Bytes::new(),
|
||||
state: EncoderState::Reading,
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns the wrapped [Stream]
|
||||
pub fn into_inner(self) -> T {
|
||||
self.inner
|
||||
}
|
||||
|
||||
fn encode(&mut self, inbuf: &[u8]) -> Result<zstd::stream::raw::Status, io::Error> {
|
||||
let res = self
|
||||
.compressor
|
||||
.run_on_buffers(inbuf, self.buffer.get_free_mut_slice())?;
|
||||
self.buffer.add_size(res.bytes_written);
|
||||
|
||||
Ok(res)
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<usize, io::Error> {
|
||||
let mut outbuf = OutBuffer::around(self.buffer.get_free_mut_slice());
|
||||
let res = self.compressor.finish(&mut outbuf, true);
|
||||
let size = outbuf.pos;
|
||||
drop(outbuf);
|
||||
self.buffer.add_size(size);
|
||||
res
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, T, O> Stream for ZstdEncoder<'a, T>
|
||||
where
|
||||
T: Stream<Item = Result<O, Error>> + Unpin,
|
||||
O: Into<Bytes>,
|
||||
{
|
||||
type Item = Result<Bytes, Error>;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
let this = self.get_mut();
|
||||
|
||||
loop {
|
||||
match this.state {
|
||||
EncoderState::Reading => {
|
||||
if let Some(res) = ready!(Pin::new(&mut this.inner).poll_next(cx)) {
|
||||
let buf = res?;
|
||||
this.input_buffer = buf.into();
|
||||
this.state = EncoderState::Writing;
|
||||
} else {
|
||||
this.state = EncoderState::Finishing;
|
||||
}
|
||||
}
|
||||
EncoderState::Writing => {
|
||||
if this.input_buffer.is_empty() {
|
||||
return Poll::Ready(Some(Err(format_err!("empty input during write"))));
|
||||
}
|
||||
let mut buf = this.input_buffer.split_off(0);
|
||||
let status = this.encode(&buf[..])?;
|
||||
this.input_buffer = buf.split_off(status.bytes_read);
|
||||
if this.input_buffer.is_empty() {
|
||||
this.state = EncoderState::Reading;
|
||||
}
|
||||
if this.buffer.is_full() {
|
||||
let bytes = this.buffer.remove_data(this.buffer.len()).to_vec();
|
||||
return Poll::Ready(Some(Ok(bytes.into())));
|
||||
}
|
||||
}
|
||||
EncoderState::Finishing => {
|
||||
let remaining = this.finish()?;
|
||||
if remaining == 0 {
|
||||
this.state = EncoderState::Finished;
|
||||
}
|
||||
if !this.buffer.is_empty() {
|
||||
let bytes = this.buffer.remove_data(this.buffer.len()).to_vec();
|
||||
return Poll::Ready(Some(Ok(bytes.into())));
|
||||
}
|
||||
}
|
||||
EncoderState::Finished => return Poll::Ready(None),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
Loading…
Reference in New Issue
Block a user