From mboxrd@z Thu Jan 1 00:00:00 1970 Return-Path: Received: from firstgate.proxmox.com (firstgate.proxmox.com [212.224.123.68]) (using TLSv1.3 with cipher TLS_AES_256_GCM_SHA384 (256/256 bits) key-exchange X25519 server-signature RSA-PSS (2048 bits)) (No client certificate requested) by lists.proxmox.com (Postfix) with ESMTPS id 1437CCA7E for ; Tue, 12 Apr 2022 13:04:22 +0200 (CEST) Received: from firstgate.proxmox.com (localhost [127.0.0.1]) by firstgate.proxmox.com (Proxmox) with ESMTP id 12439DC4D for ; Tue, 12 Apr 2022 13:04:22 +0200 (CEST) Received: from proxmox-new.maurer-it.com (proxmox-new.maurer-it.com [94.136.29.106]) (using TLSv1.3 with cipher TLS_AES_256_GCM_SHA384 (256/256 bits) key-exchange X25519 server-signature RSA-PSS (2048 bits)) (No client certificate requested) by firstgate.proxmox.com (Proxmox) with ESMTPS id 674D1DC1C for ; Tue, 12 Apr 2022 13:04:19 +0200 (CEST) Received: from proxmox-new.maurer-it.com (localhost.localdomain [127.0.0.1]) by proxmox-new.maurer-it.com (Proxmox) with ESMTP id 3FF55413F8 for ; Tue, 12 Apr 2022 13:04:19 +0200 (CEST) From: Dominik Csapak To: pbs-devel@lists.proxmox.com Date: Tue, 12 Apr 2022 13:04:14 +0200 Message-Id: <20220412110418.3360746-3-d.csapak@proxmox.com> X-Mailer: git-send-email 2.30.2 In-Reply-To: <20220412110418.3360746-1-d.csapak@proxmox.com> References: <20220412110418.3360746-1-d.csapak@proxmox.com> MIME-Version: 1.0 Content-Transfer-Encoding: 8bit X-SPAM-LEVEL: Spam detection results: 0 AWL 0.139 Adjusted score from AWL reputation of From: address BAYES_00 -1.9 Bayes spam probability is 0 to 1% KAM_DMARC_STATUS 0.01 Test Rule for DKIM or SPF Failure with Strict Alignment SPF_HELO_NONE 0.001 SPF: HELO does not publish an SPF Record SPF_PASS -0.001 SPF: sender matches SPF record T_SCC_BODY_TEXT_LINE -0.01 - URIBL_BLOCKED 0.001 ADMINISTRATOR NOTICE: The query to URIBL was blocked. See http://wiki.apache.org/spamassassin/DnsBlocklists#dnsbl-block for more information. [zstd.rs, lib.rs] Subject: [pbs-devel] [PATCH proxmox 2/2] proxmox-compression: add streaming zstd encoder X-BeenThere: pbs-devel@lists.proxmox.com X-Mailman-Version: 2.1.29 Precedence: list List-Id: Proxmox Backup Server development discussion List-Unsubscribe: , List-Archive: List-Post: List-Help: List-Subscribe: , X-List-Received-Date: Tue, 12 Apr 2022 11:04:22 -0000 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 --- proxmox-compression/Cargo.toml | 1 + proxmox-compression/src/lib.rs | 1 + proxmox-compression/src/zstd.rs | 126 ++++++++++++++++++++++++++++++++ 3 files changed, 128 insertions(+) create mode 100644 proxmox-compression/src/zstd.rs diff --git a/proxmox-compression/Cargo.toml b/proxmox-compression/Cargo.toml index c3f7f49..5ca67b2 100644 --- a/proxmox-compression/Cargo.toml +++ b/proxmox-compression/Cargo.toml @@ -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" ] } diff --git a/proxmox-compression/src/lib.rs b/proxmox-compression/src/lib.rs index e9dd113..1fcfb97 100644 --- a/proxmox-compression/src/lib.rs +++ b/proxmox-compression/src/lib.rs @@ -3,3 +3,4 @@ pub use compression::*; pub mod tar; pub mod zip; +pub mod zstd; diff --git a/proxmox-compression/src/zstd.rs b/proxmox-compression/src/zstd.rs new file mode 100644 index 0000000..0b480f6 --- /dev/null +++ b/proxmox-compression/src/zstd.rs @@ -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::with_quality(inner, 3) + } + + /// Returns a new [ZstdEncoder] with the given level + pub fn with_quality(inner: T, level: i32) -> Result { + 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 { + 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 { + 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> + Unpin, + O: Into, +{ + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + 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), + } + } + } +} -- 2.30.2