add zstd stream support

fairingrey · fairingrey · commit 04cb0281799e · 2019-05-11T23:07:02.000-07:00
diff --git a/Cargo.toml b/Cargo.toml
@@ -11,6 +11,8 @@ bytes = "0.4.12"
 flate2 = "1.0.7"
 futures-preview = "0.3.0-alpha.16"
 pin-project = "0.3.2"
+zstd = "0.4"
+zstd-safe = "1.4"
 
 [dev-dependencies]
 proptest = "0.9.3"
diff --git a/src/stream/mod.rs b/src/stream/mod.rs
@@ -13,8 +13,10 @@ mod deflate;
 mod flate;
 mod gzip;
 mod zlib;
+mod zstd;
 
-pub use brotli::{BrotliDecoder, BrotliEncoder};
-pub use deflate::{DeflateDecoder, DeflateEncoder};
-pub use gzip::{GzipDecoder, GzipEncoder};
-pub use zlib::{ZlibDecoder, ZlibEncoder};
+pub use self::brotli::{BrotliDecoder, BrotliEncoder};
+pub use self::deflate::{DeflateDecoder, DeflateEncoder};
+pub use self::gzip::{GzipDecoder, GzipEncoder};
+pub use self::zlib::{ZlibDecoder, ZlibEncoder};
+pub use self::zstd::{ZstdDecoder, ZstdEncoder};
diff --git a/src/stream/zstd.rs b/src/stream/zstd.rs
@@ -0,0 +1,200 @@
+use std::{
+    io::Result,
+    mem,
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+use bytes::{Bytes, BytesMut};
+use futures::{ready, stream::Stream};
+use pin_project::unsafe_project;
+use zstd::{
+    stream::raw::{Decoder, Encoder, Operation},
+    DEFAULT_COMPRESSION_LEVEL,
+};
+
+#[derive(Debug)]
+enum State {
+    Reading,
+    Writing(Bytes),
+    Flushing,
+    Done,
+    Invalid,
+}
+
+#[derive(Debug)]
+enum DeState {
+    Reading,
+    Writing(Bytes),
+    Done,
+    Invalid,
+}
+
+/// A zstd encoder, or compressor.
+///
+/// This structure implements a [`Stream`] interface and will read uncompressed data from an
+/// underlying stream and emit a stream of compressed data.
+#[unsafe_project(Unpin)]
+pub struct ZstdEncoder<S: Stream<Item = Result<Bytes>>> {
+    #[pin]
+    inner: S,
+    state: State,
+    output: BytesMut,
+    encoder: Encoder,
+}
+
+/// A zstd decoder, or decompressor.
+///
+/// This structure implements a [`Stream`] interface and will read compressed data from an
+/// underlying stream and emit a stream of uncompressed data.
+#[unsafe_project(Unpin)]
+pub struct ZstdDecoder<S: Stream<Item = Result<Bytes>>> {
+    #[pin]
+    inner: S,
+    state: DeState,
+    output: BytesMut,
+    decoder: Decoder,
+}
+
+impl<S: Stream<Item = Result<Bytes>>> ZstdEncoder<S> {
+    /// Creates a new encoder which will read uncompressed data from the given stream and emit a
+    /// compressed stream.
+    pub fn new(stream: S) -> ZstdEncoder<S> {
+        ZstdEncoder {
+            inner: stream,
+            state: State::Reading,
+            output: BytesMut::new(),
+            encoder: Encoder::new(DEFAULT_COMPRESSION_LEVEL).unwrap(),
+        }
+    }
+}
+
+impl<S: Stream<Item = Result<Bytes>>> ZstdDecoder<S> {
+    /// Creates a new decoder which will read compressed data from the given stream and emit an
+    /// uncompressed stream.
+    pub fn new(stream: S) -> ZstdDecoder<S> {
+        ZstdDecoder {
+            inner: stream,
+            state: DeState::Reading,
+            output: BytesMut::new(),
+            decoder: Decoder::new().unwrap(),
+        }
+    }
+}
+
+impl<S: Stream<Item = Result<Bytes>>> Stream for ZstdEncoder<S> {
+    type Item = Result<Bytes>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Result<Bytes>>> {
+        let mut this = self.project();
+
+        fn compress(
+            encoder: &mut Encoder,
+            input: &mut Bytes,
+            output: &mut BytesMut,
+        ) -> Result<Bytes> {
+            const OUTPUT_BUFFER_SIZE: usize = 8_000;
+
+            if output.len() < OUTPUT_BUFFER_SIZE {
+                output.resize(OUTPUT_BUFFER_SIZE, 0);
+            }
+
+            let status = encoder.run_on_buffers(input, output)?;
+            input.advance(status.bytes_read);
+            Ok(output.split_to(status.bytes_written).freeze())
+        }
+
+        #[allow(clippy::never_loop)] // https://github.com/rust-lang/rust-clippy/issues/4058
+        loop {
+            break match mem::replace(this.state, State::Invalid) {
+                State::Reading => {
+                    *this.state = State::Reading;
+                    *this.state = match ready!(this.inner.as_mut().poll_next(cx)) {
+                        Some(chunk) => State::Writing(chunk?),
+                        None => State::Flushing,
+                    };
+                    continue;
+                }
+                State::Writing(mut input) => {
+                    if input.is_empty() {
+                        *this.state = State::Reading;
+                        continue;
+                    }
+
+                    let chunk = compress(&mut this.encoder, &mut input, &mut this.output)?;
+
+                    *this.state = State::Writing(input);
+
+                    Poll::Ready(Some(Ok(chunk)))
+                }
+                State::Flushing => {
+                    let mut outbuffer = zstd_safe::OutBuffer::around(this.output);
+
+                    let bytes_left = this.encoder.flush(&mut outbuffer).unwrap();
+                    *this.state = if bytes_left == 0 {
+                        let _ = this.encoder.finish(&mut outbuffer, true);
+                        State::Done
+                    } else {
+                        State::Flushing
+                    };
+                    Poll::Ready(Some(Ok(outbuffer.as_slice().into())))
+                }
+                State::Done => Poll::Ready(None),
+                State::Invalid => panic!("ZstdEncoder reached invalid state"),
+            };
+        }
+    }
+}
+
+impl<S: Stream<Item = Result<Bytes>>> Stream for ZstdDecoder<S> {
+    type Item = Result<Bytes>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Result<Bytes>>> {
+        let mut this = self.project();
+
+        fn decompress(
+            decoder: &mut Decoder,
+            input: &mut Bytes,
+            output: &mut BytesMut,
+        ) -> Result<Bytes> {
+            const OUTPUT_BUFFER_SIZE: usize = 8_000;
+
+            if output.len() < OUTPUT_BUFFER_SIZE {
+                output.resize(OUTPUT_BUFFER_SIZE, 0);
+            }
+
+            let status = decoder.run_on_buffers(input, output)?;
+            dbg!(&status.remaining, &status.bytes_written, &status.bytes_read);
+            input.advance(status.bytes_read);
+            Ok(output.split_to(status.bytes_written).freeze())
+        }
+
+        #[allow(clippy::never_loop)] // https://github.com/rust-lang/rust-clippy/issues/4058
+        loop {
+            break match mem::replace(this.state, DeState::Invalid) {
+                DeState::Reading => {
+                    *this.state = DeState::Reading;
+                    *this.state = match ready!(this.inner.as_mut().poll_next(cx)) {
+                        Some(chunk) => DeState::Writing(chunk?),
+                        None => DeState::Done,
+                    };
+                    continue;
+                }
+                DeState::Writing(mut input) => {
+                    if input.is_empty() {
+                        *this.state = DeState::Reading;
+                        continue;
+                    }
+
+                    let chunk = decompress(&mut this.decoder, &mut input, &mut this.output)?;
+
+                    *this.state = DeState::Writing(input);
+
+                    Poll::Ready(Some(Ok(chunk)))
+                }
+                DeState::Done => Poll::Ready(None),
+                DeState::Invalid => panic!("ZstdDecoder reached invalid state"),
+            };
+        }
+    }
+}
diff --git a/tests/utils/mod.rs b/tests/utils/mod.rs
@@ -182,3 +182,26 @@ pub fn gzip_stream_decompress(input: impl Stream<Item = io::Result<Bytes>>) -> V
     pin_mut!(input);
     stream_to_vec(GzipDecoder::new(input))
 }
+
+pub fn zstd_compress(bytes: &[u8]) -> Vec<u8> {
+    use zstd::stream::read::Encoder;
+    use zstd::DEFAULT_COMPRESSION_LEVEL;
+    read_to_vec(Encoder::new(bytes, DEFAULT_COMPRESSION_LEVEL).unwrap())
+}
+
+pub fn zstd_decompress(bytes: &[u8]) -> Vec<u8> {
+    use zstd::stream::read::Decoder;
+    read_to_vec(Decoder::new(bytes).unwrap())
+}
+
+pub fn zstd_stream_compress(input: impl Stream<Item = io::Result<Bytes>>) -> Vec<u8> {
+    use async_compression::stream::ZstdEncoder;
+    pin_mut!(input);
+    stream_to_vec(ZstdEncoder::new(input))
+}
+
+pub fn zstd_stream_decompress(input: impl Stream<Item = io::Result<Bytes>>) -> Vec<u8> {
+    use async_compression::stream::ZstdDecoder;
+    pin_mut!(input);
+    stream_to_vec(ZstdDecoder::new(input))
+}
diff --git a/tests/zstd.rs b/tests/zstd.rs
@@ -0,0 +1,37 @@
+use std::iter::FromIterator;
+
+mod utils;
+
+#[test]
+fn zstd_stream_compress() {
+    let input = utils::InputStream::from([[1, 2, 3], [4, 5, 6]]);
+
+    let compressed = utils::zstd_stream_compress(input.stream());
+    let output = utils::zstd_decompress(&compressed);
+
+    assert_eq!(output, vec![1, 2, 3, 4, 5, 6]);
+}
+
+#[test]
+fn zstd_stream_compress_large() {
+    let input = vec![
+        Vec::from_iter((0..20_000).map(|_| rand::random())),
+        Vec::from_iter((0..20_000).map(|_| rand::random())),
+    ];
+    let input = utils::InputStream::from(input);
+
+    let compressed = utils::zstd_stream_compress(input.stream());
+    let output = utils::zstd_decompress(&compressed);
+
+    assert_eq!(output, input.bytes());
+}
+
+#[test]
+fn zstd_stream_decompress() {
+    let compressed = utils::zstd_compress(&[1, 2, 3, 4, 5, 6][..]);
+
+    let stream = utils::InputStream::from(vec![compressed]);
+    let output = utils::zstd_stream_decompress(stream.stream());
+
+    assert_eq!(output, vec![1, 2, 3, 4, 5, 6]);
+}
