refactor: switch to OrderedBatchProducerV2

src/log-store/src/kafka/background_producer.rs

@@ -23,7 +23,7 @@ use rskafka::client::producer::ProducerClient;
 use rskafka::record::Record;
 use snafu::ResultExt;
 use tokio::sync::mpsc::{self, Receiver, Sender};
-use tokio::sync::{oneshot, Mutex};
+use tokio::sync::oneshot;
 
 use crate::error::{self, Result};
 
@@ -183,7 +183,7 @@ impl BackgroundProducerWorker {
 }
 
 pub(crate) struct OrderedBatchProducer {
- sender: Mutex<Sender<ProduceRequest>>,
+ sender: Sender<ProduceRequest>,
  /// Used to control the [`BackgroundProducerWorker`].
  running: Arc<AtomicBool>,
  /// The handle of [`BackgroundProducerWorker`].
@@ -196,6 +196,16 @@ impl Drop for OrderedBatchProducer {
  }
 }
 
+pub(crate) struct ProduceResultHandle {
+ receiver: oneshot::Receiver<ProduceResultReceiver>,
+}
+
+impl ProduceResultHandle {
+ pub(crate) async fn wait(self) -> Result<Vec<i64>> {
+ self.receiver.await.unwrap().wait().await
+ }
+}
+
 impl OrderedBatchProducer {
  pub(crate) fn new(
  client: Arc<dyn ProducerClient>,
@@ -218,26 +228,22 @@ impl OrderedBatchProducer {
  };
 
  Self {
- sender: Mutex::new(tx),
+ sender: tx,
  running,
  handle: tokio::spawn(async move { worker.run().await }),
  }
  }
 
- pub(crate) async fn produce(
- &self,
- batch: Vec<Record>,
- ) -> Result<oneshot::Receiver<ProduceResultReceiver>> {
+ pub(crate) async fn produce(&self, batch: Vec<Record>) -> Result<ProduceResultHandle> {
  let receiver = {
  let (tx, rx) = oneshot::channel();
- let sender = self.sender.lock().await;
- sender
+ self.sender
  .send(ProduceRequest { batch, sender: tx })
  .await
  .expect("worker panic");
  rx
  };
 
- Ok(receiver)
+ Ok(ProduceResultHandle { receiver })
  }
 }

src/log-store/src/kafka/log_store.rs

@@ -42,9 +42,6 @@ use crate::kafka::util::record::{
 };
 use crate::metrics;
 
-/// The max flush queue size.
-pub(crate) const MAX_FLUSH_QUEUE_SIZE: usize = 512;
-
 /// A log store backed by Kafka.
 #[derive(Debug)]
 pub struct KafkaLogStore {
@@ -77,10 +74,8 @@ impl KafkaLogStore {
  })?;
  let producer_registry = ProducerRegistry::new(
  client,
- config.linger,
  config.max_batch_size.as_bytes() as usize,
  config.compression,
- MAX_FLUSH_QUEUE_SIZE,
  );
 
  Ok(Self {
@@ -205,17 +200,13 @@ impl LogStore for KafkaLogStore {
  for (region_id, records) in region_grouped_records {
  region_ids.push(region_id);
  let producer = region_grouped_producers.get(&region_id).unwrap();
- let mut receivers = Vec::with_capacity(records.len());
- for record in records {
- receivers.push(producer.produce(record).await?);
- }
- region_grouped_result_receivers.push(receivers)
+ region_grouped_result_receivers.push(producer.produce(records).await?)
  }
 
  let region_grouped_offsets = try_join_all(
  region_grouped_result_receivers
  .into_iter()
- .map(|handles| try_join_all(handles.into_iter().map(|handle| handle.wait()))),
+ .map(|handle| handle.wait()),
  )
  .await?;
 

src/log-store/src/kafka/producer_registry.rs

@@ -15,19 +15,17 @@
 use std::collections::HashMap;
 use std::fmt::Debug;
 use std::sync::Arc;
-use std::time::Duration;
 
 use rskafka::client::partition::{Compression, UnknownTopicHandling};
-use rskafka::client::producer::aggregator::RecordAggregator;
 use snafu::ResultExt;
 use store_api::logstore::provider::KafkaProvider;
 use tokio::sync::RwLock;
 
 use crate::error;
 use crate::error::Result;
-use crate::kafka::producer::OrderedBatchProducer;
+use crate::kafka::background_producer::OrderedBatchProducer;
 
-pub type OrderedBatchProducerRef = Arc<OrderedBatchProducer<RecordAggregator>>;
+pub type OrderedBatchProducerRef = Arc<OrderedBatchProducer>;
 
 // Each topic only has one partition for now.
 // The `DEFAULT_PARTITION` refers to the index of the partition.
@@ -39,12 +37,9 @@ pub(crate) const MIN_FLUSH_BATCH_SIZE: usize = 4 * 1024;
 /// The registry or [OrderedBatchProducer].
 pub struct ProducerRegistry {
  registry: RwLock<HashMap<Arc<KafkaProvider>, OrderedBatchProducerRef>>,
-
  client: rskafka::client::Client,
- linger: Duration,
  aggregator_batch_size: usize,
  compression: Compression,
- flush_queue_size: usize,
 }
 
 impl Debug for ProducerRegistry {
@@ -56,19 +51,15 @@ impl Debug for ProducerRegistry {
 impl ProducerRegistry {
  pub fn new(
  client: rskafka::client::Client,
- linger: Duration,
  aggregator_batch_size: usize,
  compression: Compression,
- flush_queue_size: usize,
  ) -> Self {
  let aggregator_batch_size = aggregator_batch_size.max(MIN_FLUSH_BATCH_SIZE);
  Self {
  registry: RwLock::new(HashMap::new()),
  client,
- linger,
  aggregator_batch_size,
  compression,
- flush_queue_size,
  }
  }
 
@@ -89,13 +80,12 @@ impl ProducerRegistry {
  partition: DEFAULT_PARTITION,
  })?;
 
- let aggregator = RecordAggregator::new(self.aggregator_batch_size);
  let producer = OrderedBatchProducer::new(
- aggregator,
  Arc::new(partition_client),
- self.linger,
  self.compression,
- self.flush_queue_size,
+ 128,
+ 64,
+ self.aggregator_batch_size,
  );
 
  Ok(Arc::new(producer))