feat: support configure compression level

examples/batching.rs

@@ -2,10 +2,7 @@
 #[macro_use]
 extern crate serde;
 use futures::{future::join_all, TryStreamExt};
-use pulsar::{
-    message::proto, message::proto::command_subscribe::SubType, message::Payload, producer,
-    Consumer, DeserializeMessage, Error as PulsarError, Pulsar, SerializeMessage, TokioExecutor,
-};
+use pulsar::{message::proto::command_subscribe::SubType, message::Payload, producer, Consumer, DeserializeMessage, Error as PulsarError, Pulsar, SerializeMessage, TokioExecutor, compression};
 
 #[derive(Debug, Serialize, Deserialize)]
 struct TestData {
@@ -42,10 +39,10 @@ async fn main() -> Result<(), pulsar::Error> {
         .with_name("my-producer2".to_string())
         .with_options(producer::ProducerOptions {
             batch_size: Some(4),
-            //compression: Some(proto::CompressionType::Lz4),
-            //compression: Some(proto::CompressionType::Zlib),
-            //compression: Some(proto::CompressionType::Zstd),
-            compression: Some(proto::CompressionType::Snappy),
+            // compression: Some(compression::Compression::Lz4(compression::CompressionLz4::default())),
+            // compression: Some(compression::Compression::Zlib(compression::CompressionZlib::default())),
+            // compression: Some(compression::Compression::Zstd(compression::CompressionZstd::default())),
+            compression: Some(compression::Compression::Snappy),
             ..Default::default()
         })
         .build()

src/compression.rs

@@ -0,0 +1,68 @@
+//! Compression strategy configs
+
+use lz4::block::CompressionMode;
+
+/// Wrapper of supported compression algorithms
+#[derive(Clone, Debug)]
+pub enum Compression {
+    None,
+    Lz4(CompressionLz4),
+    Zlib(CompressionZlib),
+    Zstd(CompressionZstd),
+    Snappy,
+}
+
+/// Options of the [lz4](https://lz4.github.io/lz4/) algorithm
+#[derive(Debug)]
+pub struct CompressionLz4 {
+    /// compression mode of lz4 to be used
+    pub mode: CompressionMode,
+}
+
+/// Options of the [zlib](https://www.zlib.net/) algorithm
+#[derive(Default, Clone, Copy, Debug)]
+pub struct CompressionZlib {
+    /// compression level of zlib to be used (0-9)
+    pub level: flate2::Compression,
+}
+
+/// Options of the [zstd](http://facebook.github.io/zstd/zstd_manual.html) algorithm
+#[derive(Clone, Copy, Debug)]
+pub struct CompressionZstd {
+    /// compression level of zstd to be used ([`zstd::compression_level_range()`])
+    pub level: i32,
+}
+
+impl Default for Compression {
+    fn default() -> Self {
+        Compression::None
+    }
+}
+
+impl Default for CompressionLz4 {
+    fn default() -> Self {
+        CompressionLz4 {
+            mode: CompressionMode::DEFAULT,
+        }
+    }
+}
+
+impl Default for CompressionZstd {
+    fn default() -> Self {
+        CompressionZstd {
+            level: zstd::DEFAULT_COMPRESSION_LEVEL,
+        }
+    }
+}
+
+impl Clone for CompressionLz4 {
+    fn clone(&self) -> Self {
+        CompressionLz4 {
+            mode: match self.mode {
+                CompressionMode::HIGHCOMPRESSION(i) => CompressionMode::HIGHCOMPRESSION(i),
+                CompressionMode::FAST(i) => CompressionMode::FAST(i),
+                CompressionMode::DEFAULT => CompressionMode::DEFAULT,
+            }
+        }
+    }
+}

src/consumer.rs

@@ -1115,12 +1115,23 @@ impl<Exe: Executor> ConsumerEngine<Exe> {
         message: CommandMessage,
         mut payload: Payload,
     ) -> Result<(), Error> {
-        let compression = payload.metadata.compression;
+        let compression = match payload.metadata.compression {
+            None => proto::CompressionType::None,
+            Some(compression) => {
+                proto::CompressionType::from_i32(compression)
+                    .ok_or_else(|| {
+                        error!("unknown compression type: {}", compression);
+                        Error::Consumer(ConsumerError::Io(std::io::Error::new(
+                            std::io::ErrorKind::Other,
+                            format!("unknown compression type: {}", compression),
+                        )))
+                    })?
+            }
+        };
 
         let payload = match compression {
-            None | Some(0) => payload,
-            // LZ4
-            Some(1) => {
+            proto::CompressionType::None => payload,
+            proto::CompressionType::Lz4 => {
                 #[cfg(not(feature = "lz4"))]
                 {
                     return Err(Error::Consumer(ConsumerError::Io(std::io::Error::new(
@@ -1142,8 +1153,7 @@ impl<Exe: Executor> ConsumerEngine<Exe> {
                     payload
                 }
             }
-            // zlib
-            Some(2) => {
+            proto::CompressionType::Zlib => {
                 #[cfg(not(feature = "flate2"))]
                 {
                     return Err(Error::Consumer(ConsumerError::Io(std::io::Error::new(
@@ -1166,8 +1176,7 @@ impl<Exe: Executor> ConsumerEngine<Exe> {
                     payload
                 }
             }
-            // zstd
-            Some(3) => {
+            proto::CompressionType::Zstd => {
                 #[cfg(not(feature = "zstd"))]
                 {
                     return Err(Error::Consumer(ConsumerError::Io(std::io::Error::new(
@@ -1186,8 +1195,7 @@ impl<Exe: Executor> ConsumerEngine<Exe> {
                     payload
                 }
             }
-            // Snappy
-            Some(4) => {
+            proto::CompressionType::Snappy => {
                 #[cfg(not(feature = "snap"))]
                 {
                     return Err(Error::Consumer(ConsumerError::Io(std::io::Error::new(
@@ -1209,13 +1217,6 @@ impl<Exe: Executor> ConsumerEngine<Exe> {
                     payload
                 }
             }
-            Some(i) => {
-                error!("unknown compression type: {}", i);
-                return Err(Error::Consumer(ConsumerError::Io(std::io::Error::new(
-                    std::io::ErrorKind::Other,
-                    format!("unknown compression type: {}", i),
-                ))));
-            }
         };
 
         match payload.metadata.num_messages_in_batch {

src/lib.rs

@@ -180,6 +180,7 @@ pub use producer::{MultiTopicProducer, Producer, ProducerOptions};
 
 pub mod authentication;
 mod client;
+pub mod compression;
 mod connection;
 mod connection_manager;
 pub mod consumer;

src/producer.rs

@@ -10,11 +10,12 @@ use crate::client::SerializeMessage;
 use crate::connection::{Connection, SerialId};
 use crate::error::{ConnectionError, ProducerError};
 use crate::executor::Executor;
-use crate::message::proto::{self, CommandSendReceipt, CompressionType, EncryptionKeys, Schema};
+use crate::message::proto::{self, CommandSendReceipt, EncryptionKeys, Schema};
 use crate::message::BatchedMessage;
 use crate::{Error, Pulsar};
 use futures::task::{Context, Poll};
 use futures::Future;
+use crate::compression::{Compression};
 
 type ProducerId = u64;
 type ProducerName = String;
@@ -127,7 +128,7 @@ pub struct ProducerOptions {
     /// batch message size
     pub batch_size: Option<u32>,
     /// algorithm used to compress the messages
-    pub compression: Option<proto::CompressionType>,
+    pub compression: Option<Compression>,
     /// producer access mode: shared = 0, exclusive = 1, waitforexclusive =2, exclusivewithoutfencing =3
     pub access_mode: Option<i32>,
 }
@@ -402,7 +403,7 @@ struct TopicProducer<Exe: Executor> {
     //putting it in a mutex because we must send multiple messages at once
     // while we might be pushing more messages from elsewhere
     batch: Option<Mutex<Batch>>,
-    compression: Option<proto::CompressionType>,
+    compression: Option<Compression>,
     drop_signal: oneshot::Sender<()>,
     options: ProducerOptions,
 }
@@ -422,26 +423,26 @@ impl<Exe: Executor> TopicProducer<Exe> {
 
         let topic = topic.clone();
         let batch_size = options.batch_size;
-        let compression = options.compression;
+        let compression = options.compression.clone();
 
         match compression {
-            None | Some(CompressionType::None) => {}
-            Some(CompressionType::Lz4) => {
+            None | Some(Compression::None) => {}
+            Some(Compression::Lz4(..)) => {
                 #[cfg(not(feature = "lz4"))]
                 return Err(Error::Custom("cannot create a producer with LZ4 compression because the 'lz4' cargo feature is not active".to_string()));
             }
-            Some(CompressionType::Zlib) => {
+            Some(Compression::Zlib(..)) => {
                 #[cfg(not(feature = "flate2"))]
                 return Err(Error::Custom("cannot create a producer with zlib compression because the 'flate2' cargo feature is not active".to_string()));
             }
-            Some(CompressionType::Zstd) => {
+            Some(Compression::Zstd(..)) => {
                 #[cfg(not(feature = "zstd"))]
                 return Err(Error::Custom("cannot create a producer with zstd compression because the 'zstd' cargo feature is not active".to_string()));
             }
-            Some(CompressionType::Snappy) => {
+            Some(Compression::Snappy) => {
                 #[cfg(not(feature = "snap"))]
                 return Err(Error::Custom("cannot create a producer with Snappy compression because the 'snap' cargo feature is not active".to_string()));
-            } //Some() => unimplemented!(),
+            }
         };
 
         let producer_name: ProducerName;
@@ -725,16 +726,16 @@ impl<Exe: Executor> TopicProducer<Exe> {
         &mut self,
         mut message: ProducerMessage,
     ) -> Result<proto::CommandSendReceipt, Error> {
-        let compressed_message = match self.compression {
-            None | Some(CompressionType::None) => message,
-            Some(CompressionType::Lz4) => {
+        let compressed_message = match self.compression.clone() {
+            None | Some(Compression::None) => message,
+            Some(Compression::Lz4(compression)) => {
                 #[cfg(not(feature = "lz4"))]
                 return unimplemented!();
 
                 #[cfg(feature = "lz4")]
                 {
                     let compressed_payload: Vec<u8> =
-                        lz4::block::compress(&message.payload[..], None, false)
+                        lz4::block::compress(&message.payload[..], Some(compression.mode), false)
                             .map_err(ProducerError::Io)?;
 
                     message.uncompressed_size = Some(message.payload.len() as u32);
@@ -743,14 +744,14 @@ impl<Exe: Executor> TopicProducer<Exe> {
                     message
                 }
             }
-            Some(CompressionType::Zlib) => {
+            Some(Compression::Zlib(compression)) => {
                 #[cfg(not(feature = "flate2"))]
                 return unimplemented!();
 
                 #[cfg(feature = "flate2")]
                 {
                     let mut e =
-                        flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::default());
+                        flate2::write::ZlibEncoder::new(Vec::new(), compression.level);
                     e.write_all(&message.payload[..])
                         .map_err(ProducerError::Io)?;
                     let compressed_payload = e.finish().map_err(ProducerError::Io)?;
@@ -761,21 +762,21 @@ impl<Exe: Executor> TopicProducer<Exe> {
                     message
                 }
             }
-            Some(CompressionType::Zstd) => {
+            Some(Compression::Zstd(compression)) => {
                 #[cfg(not(feature = "zstd"))]
                 return unimplemented!();
 
                 #[cfg(feature = "zstd")]
                 {
                     let compressed_payload =
-                        zstd::encode_all(&message.payload[..], 0).map_err(ProducerError::Io)?;
+                        zstd::encode_all(&message.payload[..], compression.level).map_err(ProducerError::Io)?;
                     message.uncompressed_size = Some(message.payload.len() as u32);
                     message.payload = compressed_payload;
                     message.compression = Some(3);
                     message
                 }
             }
-            Some(CompressionType::Snappy) => {
+            Some(Compression::Snappy) => {
                 #[cfg(not(feature = "snap"))]
                 return unimplemented!();