feat: support zstd (de)compression

Author: crepererum

Cargo.toml

@@ -36,6 +36,7 @@ time = "0.3"
 tokio = { version = "1.14", default-features = false, features = ["io-util", "net", "rt", "sync", "time"] }
 tokio-rustls = { version = "0.23", optional = true }
 tracing = "0.1"
+zstd = { version = "0.10", optional = true }
 
 [dev-dependencies]
 assert_matches = "1.5"
@@ -45,7 +46,7 @@ futures = "0.3"
 proptest = "1"
 proptest-derive = "0.3"
 rustls-pemfile = "0.2"
-rdkafka = "0.28"
+rdkafka = { version = "0.28", default-features = false, features = ["libz", "tokio", "zstd"] }
 tokio = { version = "1.14", features = ["macros", "rt-multi-thread"] }
 tracing-log = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
@@ -57,6 +58,7 @@ default = ["transport-tls"]
 compression-gzip = ["flate2"]
 compression-lz4 = ["lz4"]
 compression-snappy = ["snap"]
+compression-zstd = ["zstd"]
 fuzzing = []
 transport-socks5 = ["async-socks5"]
 transport-tls = ["rustls", "tokio-rustls"]

README.md

@@ -93,6 +93,7 @@ For more advanced production and consumption, see [`crate::client::producer`] an
 - **`compression-gzip`:** Support compression and decompression of messages using [gzip].
 - **`compression-lz4`:** Support compression and decompression of messages using [LZ4].
 - **`compression-snappy`:** Support compression and decompression of messages using [Snappy].
+- **`compression-zstd`:** Support compression and decompression of messages using [zstd].
 - **`fuzzing`:** Exposes some internal data structures so that they can be used by our fuzzers. This is NOT a stable
   feature / API!
 - **`transport-socks5`:** Allow transport via SOCKS5 proxy.
@@ -259,3 +260,4 @@ e.g. by batching writes to multiple partitions in a single ProduceRequest
 [Redpanda]: https://vectorized.io/redpanda
 [rustls]: https://github.com/rustls/rustls
 [Snappy]: https://github.com/google/snappy
+[zstd]: https://github.com/facebook/zstd

src/client/partition.rs

@@ -34,6 +34,8 @@ pub enum Compression {
     Lz4,
     #[cfg(feature = "compression-snappy")]
     Snappy,
+    #[cfg(feature = "compression-zstd")]
+    Zstd,
 }
 
 impl Default for Compression {
@@ -375,6 +377,8 @@ fn build_produce_request(
                 Compression::Lz4 => RecordBatchCompression::Lz4,
                 #[cfg(feature = "compression-snappy")]
                 Compression::Snappy => RecordBatchCompression::Snappy,
+                #[cfg(feature = "compression-zstd")]
+                Compression::Zstd => RecordBatchCompression::Zstd,
             },
             timestamp_type: RecordBatchTimestampType::CreateTime,
             producer_id: -1,

src/protocol/record.rs

@@ -673,6 +673,17 @@ where
 
                 records
             }
+            #[cfg(feature = "compression-zstd")]
+            RecordBatchCompression::Zstd => {
+                use zstd::Decoder;
+
+                let mut decoder = Decoder::new(reader)?;
+                let records = Self::read_records(&mut decoder, is_control, n_records)?;
+
+                ensure_eof(&mut decoder, "Data left in zstd block")?;
+
+                records
+            }
             #[allow(unreachable_patterns)]
             _ => {
                 return Err(ReadError::Malformed(
@@ -843,6 +854,14 @@ where
 
                 writer.write_all(&output[..len])?;
             }
+            #[cfg(feature = "compression-zstd")]
+            RecordBatchCompression::Zstd => {
+                use zstd::Encoder;
+
+                let mut encoder = Encoder::new(writer, 0)?;
+                Self::write_records(&mut encoder, self.records)?;
+                encoder.finish()?;
+            }
             #[allow(unreachable_patterns)]
             _ => {
                 return Err(WriteError::Malformed(
@@ -1114,6 +1133,56 @@ mod tests {
         assert_eq!(actual2, expected);
     }
 
+    #[cfg(feature = "compression-zstd")]
+    #[test]
+    fn test_decode_fixture_zstd() {
+        // This data was obtained by watching rdkafka.
+        let data = [
+            b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x5d\x00\x00\x00\x00".to_vec(),
+            b"\x02\xa1\x6e\x4e\x95\x00\x04\x00\x00\x00\x00\x00\x00\x01\x7e\xbf".to_vec(),
+            b"\x78\xf3\xad\x00\x00\x01\x7e\xbf\x78\xf3\xad\xff\xff\xff\xff\xff".to_vec(),
+            b"\xff\xff\xff\xff\xff\xff\xff\xff\xff\x00\x00\x00\x01\x28\xb5\x2f".to_vec(),
+            b"\xfd\x00\x58\x1d\x01\x00\xe8\xfc\x01\x00\x00\x00\xc8\x01\x78\x16".to_vec(),
+            b"\x68\x65\x6c\x6c\x6f\x20\x6b\x61\x66\x6b\x61\x02\x06\x66\x6f\x6f".to_vec(),
+            b"\x06\x62\x61\x72\x01\x00\x20\x05\x5c".to_vec(),
+        ]
+        .concat();
+
+        let actual = RecordBatch::read(&mut Cursor::new(data)).unwrap();
+        let expected = RecordBatch {
+            base_offset: 0,
+            partition_leader_epoch: 0,
+            last_offset_delta: 0,
+            first_timestamp: 1643889882029,
+            max_timestamp: 1643889882029,
+            producer_id: -1,
+            producer_epoch: -1,
+            base_sequence: -1,
+            records: ControlBatchOrRecords::Records(vec![Record {
+                timestamp_delta: 0,
+                offset_delta: 0,
+                key: Some(vec![b'x'; 100]),
+                value: Some(b"hello kafka".to_vec()),
+                headers: vec![RecordHeader {
+                    key: "foo".to_owned(),
+                    value: b"bar".to_vec(),
+                }],
+            }]),
+            compression: RecordBatchCompression::Zstd,
+            is_transactional: false,
+            timestamp_type: RecordBatchTimestampType::CreateTime,
+        };
+        assert_eq!(actual, expected);
+
+        let mut data2 = vec![];
+        actual.write(&mut data2).unwrap();
+
+        // don't compare if the data is equal because compression encoder might work slightly differently, use another
+        // roundtrip instead
+        let actual2 = RecordBatch::read(&mut Cursor::new(data2)).unwrap();
+        assert_eq!(actual2, expected);
+    }
+
     #[test]
     fn test_decode_fixture_null_key() {
         // This data was obtained by watching rdkafka driven by IOx.

tests/produce_consume.rs

@@ -105,6 +105,30 @@ async fn test_produce_rskafka_consume_rskafka_snappy() {
     assert_produce_consume(produce_rskafka, consume_rskafka, Compression::Snappy).await;
 }
 
+#[cfg(feature = "compression-zstd")]
+#[tokio::test]
+async fn test_produce_rdkafka_consume_rdkafka_zstd() {
+    assert_produce_consume(produce_rdkafka, consume_rdkafka, Compression::Zstd).await;
+}
+
+#[cfg(feature = "compression-zstd")]
+#[tokio::test]
+async fn test_produce_rskafka_consume_rdkafka_zstd() {
+    assert_produce_consume(produce_rskafka, consume_rdkafka, Compression::Zstd).await;
+}
+
+#[cfg(feature = "compression-zstd")]
+#[tokio::test]
+async fn test_produce_rdkafka_consume_rskafka_zstd() {
+    assert_produce_consume(produce_rdkafka, consume_rskafka, Compression::Zstd).await;
+}
+
+#[cfg(feature = "compression-zstd")]
+#[tokio::test]
+async fn test_produce_rskafka_consume_rskafka_zstd() {
+    assert_produce_consume(produce_rskafka, consume_rskafka, Compression::Zstd).await;
+}
+
 async fn assert_produce_consume<F1, G1, F2, G2>(
     f_produce: F1,
     f_consume: F2,

tests/rdkafka_helper.rs

@@ -37,6 +37,10 @@ pub async fn produce(
         Compression::Snappy => {
             cfg.set("compression.codec", "snappy");
         }
+        #[cfg(feature = "compression-zstd")]
+        Compression::Zstd => {
+            cfg.set("compression.codec", "zstd");
+        }
     }
     let client: FutureProducer<_> = cfg.create().unwrap();