all lists on lists.proxmox.com
 help / color / mirror / Atom feed
From: Dominik Csapak <d.csapak@proxmox.com>
To: pbs-devel@lists.proxmox.com
Subject: [pbs-devel] [PATCH proxmox 2/2] proxmox-compression: add streaming zstd encoder
Date: Tue, 12 Apr 2022 13:04:14 +0200	[thread overview]
Message-ID: <20220412110418.3360746-3-d.csapak@proxmox.com> (raw)
In-Reply-To: <20220412110418.3360746-1-d.csapak@proxmox.com>

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>
---
 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, 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),
+            }
+        }
+    }
+}
-- 
2.30.2





  parent reply	other threads:[~2022-04-12 11:04 UTC|newest]

Thread overview: 10+ messages / expand[flat|nested]  mbox.gz  Atom feed  top
2022-04-12 11:04 [pbs-devel] [PATCH proxmox/widget-toolkit/proxmox-backup] add tar.zst support for file download Dominik Csapak
2022-04-12 11:04 ` [pbs-devel] [PATCH proxmox 1/2] proxmox-compression: add async tar builder Dominik Csapak
2022-04-13  7:36   ` [pbs-devel] applied-series: " Wolfgang Bumiller
2022-04-12 11:04 ` Dominik Csapak [this message]
2022-04-12 11:04 ` [pbs-devel] [PATCH widget-toolkit 1/1] window/FileBrowser: add optional 'tar.zst' button Dominik Csapak
2022-04-13  8:37   ` [pbs-devel] applied: " Wolfgang Bumiller
2022-04-12 11:04 ` [pbs-devel] [PATCH proxmox-backup 1/3] pbs-client: add 'create_tar' helper function Dominik Csapak
2022-04-13  8:34   ` [pbs-devel] applied-series: " Wolfgang Bumiller
2022-04-12 11:04 ` [pbs-devel] [PATCH proxmox-backup 2/3] api: admin/datastore: add tar support for pxar_file_download Dominik Csapak
2022-04-12 11:04 ` [pbs-devel] [PATCH proxmox-backup 3/3] ui: datastore/Content: enable tar download in ui Dominik Csapak

Reply instructions:

You may reply publicly to this message via plain-text email
using any one of the following methods:

* Save the following mbox file, import it into your mail client,
  and reply-to-all from there: mbox

  Avoid top-posting and favor interleaved quoting:
  https://en.wikipedia.org/wiki/Posting_style#Interleaved_style

* Reply using the --to, --cc, and --in-reply-to
  switches of git-send-email(1):

  git send-email \
    --in-reply-to=20220412110418.3360746-3-d.csapak@proxmox.com \
    --to=d.csapak@proxmox.com \
    --cc=pbs-devel@lists.proxmox.com \
    /path/to/YOUR_REPLY

  https://kernel.org/pub/software/scm/git/docs/git-send-email.html

* If your mail client supports setting the In-Reply-To header
  via mailto: links, try the mailto: link
Be sure your reply has a Subject: header at the top and a blank line before the message body.
This is an external index of several public inboxes,
see mirroring instructions on how to clone and mirror
all data and code used by this external index.
Service provided by Proxmox Server Solutions GmbH | Privacy | Legal