refactor(source): implement the common split reader as a generic function (risingwavelabs#12059)

Author: wenym1

src/connector/src/macros.rs

@@ -180,57 +180,3 @@ macro_rules! impl_connector_properties {
         }
     }
 }
-
-#[macro_export]
-macro_rules! impl_common_split_reader_logic {
-    ($reader:ty, $props:ty) => {
-        impl $reader {
-            #[try_stream(boxed, ok = $crate::source::StreamChunkWithState, error = risingwave_common::error::RwError)]
-            pub(crate) async fn into_chunk_stream(self) {
-                let parser_config = self.parser_config.clone();
-                let actor_id = self.source_ctx.source_info.actor_id.to_string();
-                let source_id = self.source_ctx.source_info.source_id.to_string();
-                let metrics = self.source_ctx.metrics.clone();
-                let source_ctx = self.source_ctx.clone();
-
-                let data_stream = self.into_data_stream();
-
-                let data_stream = data_stream
-                .inspect_ok(move |data_batch| {
-                    let mut by_split_id = std::collections::HashMap::new();
-
-                    for msg in data_batch {
-                        by_split_id
-                            .entry(msg.split_id.as_ref())
-                            .or_insert_with(Vec::new)
-                            .push(msg);
-                    }
-
-                    for (split_id, msgs) in by_split_id {
-                        metrics
-                            .partition_input_count
-                            .with_label_values(&[&actor_id, &source_id, split_id])
-                            .inc_by(msgs.len() as u64);
-
-                        let sum_bytes = msgs
-                            .iter()
-                            .flat_map(|msg| msg.payload.as_ref().map(|p| p.len() as u64))
-                            .sum();
-
-                        metrics
-                            .partition_input_bytes
-                            .with_label_values(&[&actor_id, &source_id, &split_id])
-                            .inc_by(sum_bytes);
-                    }
-                }).boxed();
-
-                let parser =
-                    $crate::parser::ByteStreamSourceParserImpl::create(parser_config, source_ctx).await?;
-                #[for_await]
-                for msg_batch in parser.into_stream(data_stream) {
-                    yield msg_batch?;
-                }
-            }
-        }
-    };
-}

src/connector/src/source/cdc/source/reader.rs

@@ -16,22 +16,20 @@ use std::str::FromStr;
 
 use anyhow::{anyhow, Result};
 use async_trait::async_trait;
-use futures::{pin_mut, StreamExt, TryStreamExt};
+use futures::pin_mut;
 use futures_async_stream::try_stream;
 use risingwave_common::util::addr::HostAddr;
 use risingwave_pb::connector_service::GetEventStreamResponse;
 
-use crate::impl_common_split_reader_logic;
 use crate::parser::ParserConfig;
 use crate::source::base::SourceMessage;
 use crate::source::cdc::CdcProperties;
+use crate::source::common::{into_chunk_stream, CommonSplitReader};
 use crate::source::{
     BoxSourceWithStateStream, Column, SourceContextRef, SplitId, SplitImpl, SplitMetaData,
     SplitReader,
 };
 
-impl_common_split_reader_logic!(CdcSplitReader, CdcProperties);
-
 pub struct CdcSplitReader {
     source_id: u64,
     start_offset: Option<String>,
@@ -90,12 +88,14 @@ impl SplitReader for CdcSplitReader {
     }
 
     fn into_stream(self) -> BoxSourceWithStateStream {
-        self.into_chunk_stream()
+        let parser_config = self.parser_config.clone();
+        let source_context = self.source_ctx.clone();
+        into_chunk_stream(self, parser_config, source_context)
     }
 }
 
-impl CdcSplitReader {
-    #[try_stream(boxed, ok = Vec<SourceMessage>, error = anyhow::Error)]
+impl CommonSplitReader for CdcSplitReader {
+    #[try_stream(ok = Vec<SourceMessage>, error = anyhow::Error)]
     async fn into_data_stream(self) {
         let cdc_client = self.source_ctx.connector_client.clone().ok_or_else(|| {
             anyhow!("connector node endpoint not specified or unable to connect to connector node")

src/connector/src/source/common.rs

@@ -0,0 +1,76 @@
+// Copyright 2023 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use futures::{Stream, StreamExt, TryStreamExt};
+use futures_async_stream::try_stream;
+use risingwave_common::error::RwError;
+
+use crate::parser::ParserConfig;
+use crate::source::{SourceContextRef, SourceMessage, SplitReader, StreamChunkWithState};
+
+pub(crate) trait CommonSplitReader: SplitReader + 'static {
+    fn into_data_stream(
+        self,
+    ) -> impl Stream<Item = Result<Vec<SourceMessage>, anyhow::Error>> + Send;
+}
+
+#[try_stream(boxed, ok = StreamChunkWithState, error = RwError)]
+pub(crate) async fn into_chunk_stream(
+    reader: impl CommonSplitReader + Send,
+    parser_config: ParserConfig,
+    source_ctx: SourceContextRef,
+) {
+    let actor_id = source_ctx.source_info.actor_id.to_string();
+    let source_id = source_ctx.source_info.source_id.to_string();
+    let metrics = source_ctx.metrics.clone();
+
+    let data_stream = reader.into_data_stream();
+
+    let data_stream = data_stream
+        .inspect_ok(move |data_batch| {
+            let mut by_split_id = std::collections::HashMap::new();
+
+            for msg in data_batch {
+                by_split_id
+                    .entry(msg.split_id.as_ref())
+                    .or_insert_with(Vec::new)
+                    .push(msg);
+            }
+
+            for (split_id, msgs) in by_split_id {
+                metrics
+                    .partition_input_count
+                    .with_label_values(&[&actor_id, &source_id, split_id])
+                    .inc_by(msgs.len() as u64);
+
+                let sum_bytes = msgs
+                    .iter()
+                    .flat_map(|msg| msg.payload.as_ref().map(|p| p.len() as u64))
+                    .sum();
+
+                metrics
+                    .partition_input_bytes
+                    .with_label_values(&[&actor_id, &source_id, split_id])
+                    .inc_by(sum_bytes);
+            }
+        })
+        .boxed();
+
+    let parser =
+        crate::parser::ByteStreamSourceParserImpl::create(parser_config, source_ctx).await?;
+    #[for_await]
+    for msg_batch in parser.into_stream(data_stream) {
+        yield msg_batch?;
+    }
+}

src/connector/src/source/datagen/source/reader.rs

@@ -16,23 +16,20 @@ use std::collections::HashMap;
 
 use anyhow::{anyhow, Result};
 use async_trait::async_trait;
-use futures::{StreamExt, TryStreamExt};
-use futures_async_stream::try_stream;
+use futures::{Stream, StreamExt, TryStreamExt};
 use risingwave_common::field_generator::{FieldGeneratorImpl, VarcharProperty};
 
 use super::generator::DatagenEventGenerator;
-use crate::impl_common_split_reader_logic;
 use crate::parser::{EncodingProperties, ParserConfig, ProtocolProperties};
+use crate::source::common::{into_chunk_stream, CommonSplitReader};
 use crate::source::data_gen_util::spawn_data_generation_stream;
 use crate::source::datagen::source::SEQUENCE_FIELD_KIND;
 use crate::source::datagen::{DatagenProperties, DatagenSplit, FieldDesc};
 use crate::source::{
-    BoxSourceStream, BoxSourceWithStateStream, Column, DataType, SourceContextRef, SplitId,
+    BoxSourceWithStateStream, Column, DataType, SourceContextRef, SourceMessage, SplitId,
     SplitImpl, SplitMetaData, SplitReader,
 };
 
-impl_common_split_reader_logic!(DatagenSplitReader, DatagenProperties);
-
 pub struct DatagenSplitReader {
     generator: DatagenEventGenerator,
     assigned_split: DatagenSplit,
@@ -170,13 +167,17 @@ impl SplitReader for DatagenSplitReader {
                 )
                 .boxed()
             }
-            _ => self.into_chunk_stream(),
+            _ => {
+                let parser_config = self.parser_config.clone();
+                let source_context = self.source_ctx.clone();
+                into_chunk_stream(self, parser_config, source_context)
+            }
         }
     }
 }
 
-impl DatagenSplitReader {
-    pub(crate) fn into_data_stream(self) -> BoxSourceStream {
+impl CommonSplitReader for DatagenSplitReader {
+    fn into_data_stream(self) -> impl Stream<Item = Result<Vec<SourceMessage>, anyhow::Error>> {
         // Will buffer at most 4 event chunks.
         const BUFFER_SIZE: usize = 4;
         spawn_data_generation_stream(self.generator.into_msg_stream(), BUFFER_SIZE).boxed()

src/connector/src/source/google_pubsub/source/reader.rs

@@ -15,16 +15,15 @@
 use anyhow::{anyhow, ensure, Context, Result};
 use async_trait::async_trait;
 use chrono::{NaiveDateTime, TimeZone, Utc};
-use futures::{StreamExt, TryStreamExt};
 use futures_async_stream::try_stream;
 use google_cloud_pubsub::client::Client;
 use google_cloud_pubsub::subscription::{SeekTo, Subscription};
 use risingwave_common::bail;
 use tonic::Code;
 
 use super::TaggedReceivedMessage;
-use crate::impl_common_split_reader_logic;
 use crate::parser::ParserConfig;
+use crate::source::common::{into_chunk_stream, CommonSplitReader};
 use crate::source::google_pubsub::PubsubProperties;
 use crate::source::{
     BoxSourceWithStateStream, Column, SourceContextRef, SourceMessage, SplitId, SplitImpl,
@@ -33,8 +32,6 @@ use crate::source::{
 
 const PUBSUB_MAX_FETCH_MESSAGES: usize = 1024;
 
-impl_common_split_reader_logic!(PubsubSplitReader, PubsubProperties);
-
 pub struct PubsubSplitReader {
     subscription: Subscription,
     stop_offset: Option<NaiveDateTime>,
@@ -44,8 +41,8 @@ pub struct PubsubSplitReader {
     source_ctx: SourceContextRef,
 }
 
-impl PubsubSplitReader {
-    #[try_stream(boxed, ok = Vec<SourceMessage>, error = anyhow::Error)]
+impl CommonSplitReader for PubsubSplitReader {
+    #[try_stream(ok = Vec<SourceMessage>, error = anyhow::Error)]
     async fn into_data_stream(self) {
         loop {
             let pull_result = self
@@ -172,6 +169,8 @@ impl SplitReader for PubsubSplitReader {
     }
 
     fn into_stream(self) -> BoxSourceWithStateStream {
-        self.into_chunk_stream()
+        let parser_config = self.parser_config.clone();
+        let source_context = self.source_ctx.clone();
+        into_chunk_stream(self, parser_config, source_context)
     }
 }

src/connector/src/source/kafka/source/reader.rs

@@ -19,16 +19,16 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
 use anyhow::{anyhow, Result};
 use async_trait::async_trait;
-use futures::{StreamExt, TryStreamExt};
+use futures::StreamExt;
 use futures_async_stream::try_stream;
 use rdkafka::config::RDKafkaLogLevel;
 use rdkafka::consumer::{Consumer, StreamConsumer};
 use rdkafka::error::KafkaError;
 use rdkafka::{ClientConfig, Message, Offset, TopicPartitionList};
 
-use crate::impl_common_split_reader_logic;
 use crate::parser::ParserConfig;
 use crate::source::base::SourceMessage;
+use crate::source::common::{into_chunk_stream, CommonSplitReader};
 use crate::source::kafka::{
     KafkaProperties, KafkaSplit, PrivateLinkConsumerContext, KAFKA_ISOLATION_LEVEL,
 };
@@ -37,8 +37,6 @@ use crate::source::{
     SplitReader,
 };
 
-impl_common_split_reader_logic!(KafkaSplitReader, KafkaProperties);
-
 pub struct KafkaSplitReader {
     consumer: StreamConsumer<PrivateLinkConsumerContext>,
     offsets: HashMap<SplitId, (Option<i64>, Option<i64>)>,
@@ -159,7 +157,9 @@ impl SplitReader for KafkaSplitReader {
     }
 
     fn into_stream(self) -> BoxSourceWithStateStream {
-        self.into_chunk_stream()
+        let parser_config = self.parser_config.clone();
+        let source_context = self.source_ctx.clone();
+        into_chunk_stream(self, parser_config, source_context)
     }
 }
 
@@ -176,9 +176,11 @@ impl KafkaSplitReader {
             ])
             .set(offset);
     }
+}
 
-    #[try_stream(boxed, ok = Vec<SourceMessage>, error = anyhow::Error)]
-    pub async fn into_data_stream(self) {
+impl CommonSplitReader for KafkaSplitReader {
+    #[try_stream(ok = Vec<SourceMessage>, error = anyhow::Error)]
+    async fn into_data_stream(self) {
         if self.offsets.values().all(|(start_offset, stop_offset)| {
             match (start_offset, stop_offset) {
                 (Some(start), Some(stop)) if (*start + 1) >= *stop => true,