feat: Configurable option for compression (#151)

Author: vrongmeal

src/append_session.rs

@@ -29,6 +29,7 @@ use crate::{
 async fn connect(
     stream_client: &StreamClient,
     frame_signal: FrameSignal,
+    compression: bool,
 ) -> Result<
     (
         mpsc::Sender<types::AppendInput>,
@@ -44,6 +45,7 @@ async fn connect(
             .frame_monitoring_stream_service_client(frame_signal.clone()),
         &stream_client.stream,
         ReceiverStream::new(input_rx),
+        compression,
     );
 
     Ok((input_tx, stream_client.inner.send(service_req).await?))
@@ -214,6 +216,7 @@ async fn session_inner<S>(
     frame_signal: FrameSignal,
     stream_client: StreamClient,
     output_tx: mpsc::Sender<Result<types::AppendOutput, ClientError>>,
+    compression: bool,
 ) -> Result<(), ClientError>
 where
     S: 'static + Send + Unpin + futures::Stream<Item = types::AppendInput>,
@@ -230,7 +233,8 @@ where
 
     assert!(*inflight_size <= stream_client.inner.config.max_append_inflight_bytes);
 
-    let (s2_input_tx, mut s2_ack_stream) = connect(&stream_client, frame_signal.clone()).await?;
+    let (s2_input_tx, mut s2_ack_stream) =
+        connect(&stream_client, frame_signal.clone(), compression).await?;
     let batch_ack_deadline = stream_client.inner.config.request_timeout;
 
     if !inflight.is_empty() {
@@ -346,6 +350,7 @@ pub(crate) async fn manage_session<S>(
     stream_client: StreamClient,
     input: S,
     output_tx: mpsc::Sender<Result<types::AppendOutput, ClientError>>,
+    compression: bool,
 ) where
     S: 'static + Send + Unpin + futures::Stream<Item = types::AppendInput>,
 {
@@ -367,6 +372,7 @@ pub(crate) async fn manage_session<S>(
             frame_signal.clone(),
             stream_client.clone(),
             output_tx.clone(),
+            compression,
         )
         .await
         {

src/client.rs

@@ -246,6 +246,7 @@ pub struct ClientConfig {
     pub(crate) retry_backoff_duration: Duration,
     pub(crate) max_attempts: usize,
     pub(crate) max_append_inflight_bytes: u64,
+    pub(crate) compression: bool,
 }
 
 impl ClientConfig {
@@ -263,6 +264,7 @@ impl ClientConfig {
             retry_backoff_duration: Duration::from_millis(100),
             max_attempts: 3,
             max_append_inflight_bytes: 100 * MIB_BYTES,
+            compression: false,
         }
     }
 
@@ -353,6 +355,15 @@ impl ClientConfig {
             ..self
         }
     }
+
+    /// Configure compression for requests and responses.
+    /// Disabled by default.
+    pub fn with_compression(self, compression: bool) -> Self {
+        Self {
+            compression,
+            ..self
+        }
+    }
 }
 
 /// Error from client operations.
@@ -616,6 +627,7 @@ impl StreamClient {
                 self.inner.stream_service_client(),
                 &self.stream,
                 req,
+                self.inner.config.compression,
             ))
             .await
     }
@@ -625,8 +637,12 @@ impl StreamClient {
         &self,
         req: types::ReadSessionRequest,
     ) -> Result<Streaming<types::ReadOutput>, ClientError> {
-        let request =
-            ReadSessionServiceRequest::new(self.inner.stream_service_client(), &self.stream, req);
+        let request = ReadSessionServiceRequest::new(
+            self.inner.stream_service_client(),
+            &self.stream,
+            req,
+            self.inner.config.compression,
+        );
         self.inner
             .send_retryable(request.clone())
             .await
@@ -653,6 +669,7 @@ impl StreamClient {
                 frame_signal,
                 &self.stream,
                 req,
+                self.inner.config.compression,
             ))
             .await
     }
@@ -671,6 +688,7 @@ impl StreamClient {
             self.clone(),
             req,
             response_tx,
+            self.inner.config.compression,
         ));
 
         Ok(Box::pin(ReceiverStream::new(response_rx)))
@@ -817,6 +835,7 @@ impl ClientInner {
     pub(crate) fn stream_service_client(&self) -> StreamServiceClient<Channel> {
         StreamServiceClient::new(self.channel.clone())
     }
+
     pub(crate) fn frame_monitoring_stream_service_client(
         &self,
         frame_signal: FrameSignal,

src/service/stream.rs

@@ -64,12 +64,16 @@ pub struct ReadServiceRequest {
 
 impl ReadServiceRequest {
     pub fn new(
-        client: StreamServiceClient<Channel>,
+        mut client: StreamServiceClient<Channel>,
         stream: impl Into<String>,
         req: types::ReadRequest,
+        compression: bool,
     ) -> Self {
+        if compression {
+            client = client.accept_compressed(CompressionEncoding::Zstd);
+        }
         Self {
-            client: client.accept_compressed(CompressionEncoding::Zstd),
+            client,
             stream: stream.into(),
             req,
         }
@@ -111,12 +115,16 @@ pub struct ReadSessionServiceRequest {
 
 impl ReadSessionServiceRequest {
     pub fn new(
-        client: StreamServiceClient<Channel>,
+        mut client: StreamServiceClient<Channel>,
         stream: impl Into<String>,
         req: types::ReadSessionRequest,
+        compression: bool,
     ) -> Self {
+        if compression {
+            client = client.accept_compressed(CompressionEncoding::Zstd);
+        }
         Self {
-            client: client.accept_compressed(CompressionEncoding::Zstd),
+            client,
             stream: stream.into(),
             req,
         }
@@ -181,14 +189,18 @@ pub struct AppendServiceRequest {
 
 impl AppendServiceRequest {
     pub fn new(
-        client: StreamServiceClient<RequestFrameMonitor>,
+        mut client: StreamServiceClient<RequestFrameMonitor>,
         append_retry_policy: AppendRetryPolicy,
         frame_signal: FrameSignal,
         stream: impl Into<String>,
         req: types::AppendInput,
+        compression: bool,
     ) -> Self {
+        if compression {
+            client = client.send_compressed(CompressionEncoding::Zstd);
+        }
         Self {
-            client: client.send_compressed(CompressionEncoding::Zstd),
+            client,
             append_retry_policy,
             frame_signal,
             stream: stream.into(),
@@ -259,12 +271,16 @@ where
     S: Send + futures::Stream<Item = types::AppendInput> + Unpin,
 {
     pub fn new(
-        client: StreamServiceClient<RequestFrameMonitor>,
+        mut client: StreamServiceClient<RequestFrameMonitor>,
         stream: impl Into<String>,
         req: S,
+        compression: bool,
     ) -> Self {
+        if compression {
+            client = client.send_compressed(CompressionEncoding::Zstd);
+        }
         Self {
-            client: client.send_compressed(CompressionEncoding::Zstd),
+            client,
             stream: stream.into(),
             req: Some(req),
         }