Refactor chunk size configuration

Author: hippalus

src/connectors/kafka/config.rs

@@ -640,6 +640,28 @@ impl ConsumerConfig {
     pub fn topics(&self) -> Vec<&str> {
         self.topics.iter().map(|t| t.as_str()).collect()
     }
+
+    pub fn buffer_config(&self) -> BufferConfig {
+        BufferConfig {
+            buffer_size: self.buffer_size,
+            buffer_timeout: self.buffer_timeout,
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct BufferConfig {
+    pub buffer_size: usize,
+    pub buffer_timeout: Duration,
+}
+
+impl Default for BufferConfig {
+    fn default() -> Self {
+        Self {
+            buffer_size: 10000,
+            buffer_timeout: Duration::from_millis(10000),
+        }
+    }
 }
 
 impl ProducerConfig {
@@ -874,7 +896,7 @@ impl Default for ConsumerConfig {
             topics: vec![],
             group_id: "parseable-connect-cg".to_string(),
             buffer_size: 10_000,
-            buffer_timeout: Duration::from_millis(5000),
+            buffer_timeout: Duration::from_millis(10000),
             group_instance_id: "parseable-cg-ii".to_string(),
             // NOTE: cooperative-sticky does not work well in rdkafka when using manual commit.
             // @see https://github.com/confluentinc/librdkafka/issues/4629

src/connectors/kafka/processor.rs

@@ -17,6 +17,7 @@
  */
 
 use crate::connectors::common::processor::Processor;
+use crate::connectors::kafka::config::BufferConfig;
 use crate::connectors::kafka::{ConsumerRecord, StreamConsumer, TopicPartition};
 use crate::event::format;
 use crate::event::format::EventFormat;
@@ -30,7 +31,6 @@ use rdkafka::consumer::{CommitMode, Consumer};
 use serde_json::Value;
 use std::collections::HashMap;
 use std::sync::Arc;
-use std::time::Duration;
 use tokio_stream::wrappers::ReceiverStream;
 use tracing::{debug, error, warn};
 
@@ -110,25 +110,25 @@ where
 {
     processor: Arc<P>,
     consumer: Arc<StreamConsumer>,
-    buffer_size: usize,
-    buffer_timeout: Duration,
+    buffer_config: BufferConfig,
 }
 
 impl<P> StreamWorker<P>
 where
     P: Processor<Vec<ConsumerRecord>, ()> + Send + Sync + 'static,
 {
-    pub fn new(
-        processor: Arc<P>,
-        consumer: Arc<StreamConsumer>,
-        buffer_size: usize,
-        buffer_timeout: Duration,
-    ) -> Self {
+    pub fn new(processor: Arc<P>, consumer: Arc<StreamConsumer>) -> Self {
+        let buffer_config = consumer
+            .context()
+            .config()
+            .consumer()
+            .expect("Consumer config is missing")
+            .buffer_config();
+
         Self {
             processor,
             consumer,
-            buffer_size,
-            buffer_timeout,
+            buffer_config,
         }
     }
 
@@ -139,8 +139,8 @@ where
     ) -> anyhow::Result<()> {
         let chunked_stream = tokio_stream::StreamExt::chunks_timeout(
             record_stream,
-            self.buffer_size,
-            self.buffer_timeout,
+            self.buffer_config.buffer_size,
+            self.buffer_config.buffer_timeout,
         );
 
         chunked_stream

src/connectors/kafka/sink.rs

@@ -23,7 +23,6 @@ use crate::connectors::kafka::ConsumerRecord;
 use anyhow::Result;
 use futures_util::StreamExt;
 use std::sync::Arc;
-use tokio::time::Duration;
 use tracing::{error, info};
 
 pub struct KafkaSinkConnector<P>
@@ -38,17 +37,10 @@ impl<P> KafkaSinkConnector<P>
 where
     P: Processor<Vec<ConsumerRecord>, ()> + Send + Sync + 'static,
 {
-    pub fn new(
-        kafka_streams: KafkaStreams,
-        processor: P,
-        buffer_size: usize,
-        buffer_timeout: Duration,
-    ) -> Self {
+    pub fn new(kafka_streams: KafkaStreams, processor: P) -> Self {
         let worker = Arc::new(StreamWorker::new(
             Arc::new(processor),
             kafka_streams.consumer(),
-            buffer_size,
-            buffer_timeout,
         ));
 
         Self {

src/connectors/mod.rs

@@ -31,7 +31,6 @@ use crate::option::{Mode, CONFIG};
 use actix_web_prometheus::PrometheusMetrics;
 use prometheus::Registry;
 use std::sync::Arc;
-use std::time::Duration;
 use tokio::sync::RwLock;
 use tracing::{info, warn};
 
@@ -93,8 +92,7 @@ where
     let stats = kafka_streams.statistics();
     registry.register(Box::new(KafkaMetricsCollector::new(stats)?))?;
 
-    let kafka_parseable_sink_connector =
-        KafkaSinkConnector::new(kafka_streams, processor, 10000, Duration::from_millis(5000));
+    let kafka_parseable_sink_connector = KafkaSinkConnector::new(kafka_streams, processor);
 
     rebalance_listener.start();
     kafka_parseable_sink_connector.run().await?;