feat: support gzip compression

Author: crepererum

Cargo.toml

@@ -11,6 +11,7 @@ async-socks5 = { version = "0.5", optional = true }
 async-trait = "0.1"
 bytes = "1.1"
 crc = "2"
+flate2 = { version = "1", optional = true }
 futures = "0.3"
 integer-encoding = "3"
 parking_lot = "0.11"
@@ -38,9 +39,11 @@ uuid = { version = "0.8", features = ["v4"] }
 
 [features]
 default = ["transport-tls"]
+
+compression-gzip = ["flate2"]
 fuzzing = []
-transport-tls = ["rustls", "tokio-rustls"]
 transport-socks5 = ["async-socks5"]
+transport-tls = ["rustls", "tokio-rustls"]
 
 [workspace]
 members = [

README.md

@@ -25,7 +25,10 @@ It will be a good fit for workloads that:
 ```rust,no_run
 # async fn test() {
 use rskafka::{
-    client::ClientBuilder,
+    client::{
+        ClientBuilder,
+        partition::Compression,
+    },
     record::Record,
 };
 use time::OffsetDateTime;
@@ -63,7 +66,7 @@ let record = Record {
     ]),
     timestamp: OffsetDateTime::now_utc(),
 };
-partition_client.produce(vec![record]).await.unwrap();
+partition_client.produce(vec![record], Compression::default()).await.unwrap();
 
 // consume data
 let (records, high_watermark) = partition_client
@@ -82,10 +85,11 @@ For more advanced production and consumption, see [`crate::client::producer`] an
 
 ## Features
 
+- **`compression-gzip`:** Support compression and decompression of messages using [gzip].
 - **`fuzzing`:** Exposes some internal data structures so that they can be used by our fuzzers. This is NOT a stable
   feature / API!
-- **`transport-tls` (default):** Allows TLS transport via [rustls].
 - **`transport-socks5`:** Allow transport via SOCKS5 proxy.
+- **`transport-tls` (default):** Allows TLS transport via [rustls].
 
 ## Testing
 
@@ -215,6 +219,7 @@ e.g. by batching writes to multiple partitions in a single ProduceRequest
 
 [Apache Kafka]: https://kafka.apache.org/
 [cargo-fuzz]: https://github.com/rust-fuzz/cargo-fuzz
+[gzip]: https://en.wikipedia.org/wiki/Gzip
 [IOx]: https://github.com/influxdata/influxdb_iox/
 [LLDB]: https://lldb.llvm.org/
 [Redpanda]: https://vectorized.io/redpanda

src/client/partition.rs

@@ -23,6 +23,19 @@ use time::OffsetDateTime;
 use tokio::sync::Mutex;
 use tracing::{error, info};
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum Compression {
+    NoCompression,
+    #[cfg(feature = "compression-gzip")]
+    Gzip,
+}
+
+impl Default for Compression {
+    fn default() -> Self {
+        Self::NoCompression
+    }
+}
+
 /// Many operations must be performed on the leader for a partition
 ///
 /// Additionally a partition is the unit of concurrency within Kafka
@@ -62,7 +75,11 @@ impl PartitionClient {
     }
 
     /// Produce a batch of records to the partition
-    pub async fn produce(&self, records: Vec<Record>) -> Result<Vec<i64>> {
+    pub async fn produce(
+        &self,
+        records: Vec<Record>,
+        compression: Compression,
+    ) -> Result<Vec<i64>> {
         // skip request entirely if `records` is empty
         if records.is_empty() {
             return Ok(vec![]);
@@ -105,7 +122,11 @@ impl PartitionClient {
                 last_offset_delta: n as i32 - 1,
                 is_transactional: false,
                 base_sequence: -1,
-                compression: RecordBatchCompression::NoCompression,
+                compression: match compression {
+                    Compression::NoCompression => RecordBatchCompression::NoCompression,
+                    #[cfg(feature = "compression-gzip")]
+                    Compression::Gzip => RecordBatchCompression::Gzip,
+                },
                 timestamp_type: RecordBatchTimestampType::CreateTime,
                 producer_id: -1,
                 producer_epoch: -1,

src/client/producer.rs

@@ -207,6 +207,8 @@ use crate::client::producer::aggregator::TryPush;
 use crate::client::{error::Error as ClientError, partition::PartitionClient};
 use crate::record::Record;
 
+use super::partition::Compression;
+
 pub mod aggregator;
 mod broadcast;
 
@@ -230,6 +232,8 @@ pub struct BatchProducerBuilder {
     client: Arc<dyn ProducerClient>,
 
     linger: Duration,
+
+    compression: Compression,
 }
 
 impl BatchProducerBuilder {
@@ -243,6 +247,7 @@ impl BatchProducerBuilder {
         Self {
             client,
             linger: Duration::from_millis(5),
+            compression: Compression::default(),
         }
     }
 
@@ -251,12 +256,21 @@ impl BatchProducerBuilder {
         Self { linger, ..self }
     }
 
+    /// Sets compression.
+    pub fn with_compression(self, compression: Compression) -> Self {
+        Self {
+            compression,
+            ..self
+        }
+    }
+
     pub fn build<A>(self, aggregator: A) -> BatchProducer<A>
     where
         A: aggregator::Aggregator,
     {
         BatchProducer {
             linger: self.linger,
+            compression: self.compression,
             client: self.client,
             inner: Mutex::new(ProducerInner {
                 aggregator,
@@ -268,12 +282,20 @@ impl BatchProducerBuilder {
 
 // A trait wrapper to allow mocking
 trait ProducerClient: std::fmt::Debug + Send + Sync {
-    fn produce(&self, records: Vec<Record>) -> BoxFuture<'_, Result<Vec<i64>, ClientError>>;
+    fn produce(
+        &self,
+        records: Vec<Record>,
+        compression: Compression,
+    ) -> BoxFuture<'_, Result<Vec<i64>, ClientError>>;
 }
 
 impl ProducerClient for PartitionClient {
-    fn produce(&self, records: Vec<Record>) -> BoxFuture<'_, Result<Vec<i64>, ClientError>> {
-        Box::pin(self.produce(records))
+    fn produce(
+        &self,
+        records: Vec<Record>,
+        compression: Compression,
+    ) -> BoxFuture<'_, Result<Vec<i64>, ClientError>> {
+        Box::pin(self.produce(records, compression))
     }
 }
 
@@ -293,6 +315,8 @@ where
 {
     linger: Duration,
 
+    compression: Compression,
+
     client: Arc<dyn ProducerClient>,
 
     inner: Mutex<ProducerInner<A>>,
@@ -370,7 +394,7 @@ where
                 TryPush::NoCapacity(data) => {
                     debug!("Insufficient capacity in aggregator - flushing");
 
-                    Self::flush_impl(&mut inner, self.client.as_ref()).await;
+                    Self::flush_impl(&mut inner, self.client.as_ref(), self.compression).await;
                     match inner
                         .aggregator
                         .try_push(data)
@@ -413,7 +437,7 @@ where
         debug!("Linger expired - flushing");
 
         // Flush data
-        Self::flush_impl(&mut inner, self.client.as_ref()).await;
+        Self::flush_impl(&mut inner, self.client.as_ref(), self.compression).await;
 
         extract(&result_slot.now_or_never().expect("just flushed"), tag)
     }
@@ -422,12 +446,16 @@ where
     pub async fn flush(&self) {
         let mut inner = self.inner.lock().await;
         debug!("Manual flush");
-        Self::flush_impl(&mut inner, self.client.as_ref()).await;
+        Self::flush_impl(&mut inner, self.client.as_ref(), self.compression).await;
     }
 
     /// Flushes out the data from the aggregator, publishes the result to the result slot,
     /// and creates a fresh result slot for future writes to use
-    async fn flush_impl(inner: &mut ProducerInner<A>, client: &dyn ProducerClient) {
+    async fn flush_impl(
+        inner: &mut ProducerInner<A>,
+        client: &dyn ProducerClient,
+        compression: Compression,
+    ) {
         trace!("Flushing batch producer");
 
         let (output, status_deagg) = match inner.aggregator.flush() {
@@ -444,7 +472,7 @@ where
             return;
         }
 
-        let r = client.produce(output).await;
+        let r = client.produce(output, compression).await;
 
         // Reset result slot
         let slot = std::mem::take(&mut inner.result_slot);
@@ -483,7 +511,11 @@ mod tests {
     }
 
     impl ProducerClient for MockClient {
-        fn produce(&self, records: Vec<Record>) -> BoxFuture<'_, Result<Vec<i64>, ClientError>> {
+        fn produce(
+            &self,
+            records: Vec<Record>,
+            _compression: Compression,
+        ) -> BoxFuture<'_, Result<Vec<i64>, ClientError>> {
             Box::pin(async move {
                 tokio::time::sleep(self.delay).await;