feat: use inner message

rust/connector/src/base.rs

@@ -1,5 +1,7 @@
 use anyhow::Result;
 use async_trait::async_trait;
+use bytes::Bytes;
+use serde::{Deserialize, Serialize};
 
 pub enum SourceOffset {
     Number(i64),
@@ -12,15 +14,22 @@ pub trait SourceMessage {
     fn serialize(&self) -> Result<String>;
 }
 
+#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
+pub struct InnerMessage {
+    pub payload: Option<Bytes>,
+    pub offset: String,
+    pub split_id: String,
+}
+
 pub trait SourceSplit {
     fn id(&self) -> String;
     fn to_string(&self) -> Result<String>;
 }
 
 #[async_trait]
 pub trait SourceReader: Sized {
-    async fn next(&mut self) -> Result<Option<Vec<Vec<u8>>>>;
-    async fn assign_split<'a>(&mut self, split: &'a [u8]) -> Result<()>;
+    async fn next(&mut self) -> Result<Option<Vec<InnerMessage>>>;
+    async fn assign_split<'a>(&'a mut self, split: &'a [u8]) -> Result<()>;
 }
 
 #[async_trait]

rust/connector/src/kafka/source/message.rs

@@ -1,9 +1,10 @@
 use anyhow::anyhow;
+use bytes::Bytes;
 use rdkafka::message::BorrowedMessage;
 use rdkafka::Message;
 use serde::{Deserialize, Serialize};
 
-use crate::base::{SourceMessage, SourceOffset};
+use crate::base::{InnerMessage, SourceMessage, SourceOffset};
 
 #[derive(Clone, Serialize, Deserialize)]
 pub struct KafkaMessage {
@@ -27,13 +28,12 @@ impl SourceMessage for KafkaMessage {
     }
 }
 
-impl<'a> From<BorrowedMessage<'a>> for KafkaMessage {
+impl<'a> From<BorrowedMessage<'a>> for InnerMessage {
     fn from(message: BorrowedMessage<'a>) -> Self {
-        KafkaMessage {
-            partition: message.partition(),
-            offset: message.offset(),
-            key: message.key().map(|key| key.to_vec()),
-            payload: message.payload().map(|m| m.to_vec()),
+        InnerMessage {
+            payload: message.payload().map(Bytes::copy_from_slice),
+            offset: message.offset().to_string(),
+            split_id: message.partition().to_string(),
         }
     }
 }

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

@@ -10,8 +10,7 @@ use rdkafka::consumer::stream_consumer::StreamPartitionQueue;
 use rdkafka::consumer::{Consumer, DefaultConsumerContext, StreamConsumer};
 use rdkafka::{ClientConfig, Message, Offset, TopicPartitionList};
 
-use crate::base::{SourceMessage, SourceReader};
-use crate::kafka::source::message::KafkaMessage;
+use crate::base::{InnerMessage, SourceReader};
 use crate::kafka::split::{KafkaOffset, KafkaSplit};
 
 const KAFKA_MAX_FETCH_MESSAGES: usize = 1024;
@@ -25,7 +24,7 @@ pub struct KafkaSplitReader {
 
 #[async_trait]
 impl SourceReader for KafkaSplitReader {
-    async fn next(&mut self) -> Result<Option<Vec<Vec<u8>>>> {
+    async fn next(&mut self) -> Result<Option<Vec<InnerMessage>>> {
         let mut stream = self
             .partition_queue
             .stream()
@@ -55,35 +54,35 @@ impl SourceReader for KafkaSplitReader {
                 }
             }
 
-            ret.push(KafkaMessage::from(msg).serialize()?.into_bytes());
+            ret.push(InnerMessage::from(msg));
         }
 
         Ok(Some(ret))
     }
 
-    async fn assign_split<'a>(&mut self, split: &'a [u8]) -> Result<()> {
-        let split: KafkaSplit = serde_json::from_str(from_utf8(split)?)?;
+    async fn assign_split<'a>(&'a mut self, split: &'a [u8]) -> Result<()> {
+        let kafka_split: KafkaSplit = serde_json::from_str(from_utf8(split)?)?;
         let mut tpl = TopicPartitionList::new();
 
-        let offset = match split.start_offset {
+        let offset = match kafka_split.start_offset {
             KafkaOffset::None | KafkaOffset::Earliest => Offset::Beginning,
             KafkaOffset::Latest => Offset::End,
             KafkaOffset::Offset(offset) => Offset::Offset(offset),
             KafkaOffset::Timestamp(_) => unimplemented!(),
         };
 
-        tpl.add_partition_offset(self.topic.as_str(), split.partition, offset)
+        tpl.add_partition_offset(self.topic.as_str(), kafka_split.partition, offset)
             .map_err(|e| anyhow!(e))?;
 
         self.consumer.assign(&tpl).map_err(|e| anyhow!(e))?;
 
         let partition_queue = self
             .consumer
-            .split_partition_queue(self.topic.as_str(), split.partition)
+            .split_partition_queue(self.topic.as_str(), kafka_split.partition)
             .ok_or_else(|| anyhow!("Failed to split partition queue"))?;
 
         self.partition_queue = partition_queue;
-        self.assigned_split = split;
+        self.assigned_split = kafka_split;
 
         Ok(())
     }

rust/connector/src/kinesis/source/message.rs

@@ -1,8 +1,9 @@
 use anyhow::anyhow;
 use aws_sdk_kinesis::model::Record;
+use bytes::Bytes;
 use serde::{Deserialize, Serialize};
 
-use crate::base::{SourceMessage, SourceOffset};
+use crate::base::{InnerMessage, SourceMessage, SourceOffset};
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct KinesisMessage {
@@ -26,6 +27,19 @@ impl SourceMessage for KinesisMessage {
     }
 }
 
+impl From<KinesisMessage> for InnerMessage {
+    fn from(msg: KinesisMessage) -> Self {
+        InnerMessage {
+            payload: msg
+                .payload
+                .as_ref()
+                .map(|payload| Bytes::copy_from_slice(payload)),
+            offset: msg.sequence_number.clone(),
+            split_id: msg.shard_id,
+        }
+    }
+}
+
 impl KinesisMessage {
     pub fn new(shard_id: String, message: Record) -> Self {
         KinesisMessage {

rust/connector/src/kinesis/source/reader.rs

@@ -11,7 +11,7 @@ use aws_sdk_kinesis::Client as kinesis_client;
 use aws_smithy_types::DateTime;
 use http::uri::Uri;
 
-use crate::base::{SourceMessage, SourceReader};
+use crate::base::{InnerMessage, SourceReader};
 use crate::kinesis::config::AwsConfigInfo;
 use crate::kinesis::source::message::KinesisMessage;
 use crate::kinesis::source::state::KinesisSplitReaderState;
@@ -28,7 +28,7 @@ pub struct KinesisSplitReader {
 
 #[async_trait]
 impl SourceReader for KinesisSplitReader {
-    async fn next(&mut self) -> Result<Option<Vec<Vec<u8>>>> {
+    async fn next(&mut self) -> Result<Option<Vec<InnerMessage>>> {
         if self.assigned_split.is_none() {
             return Err(anyhow::Error::msg(
                 "you should call `assign_split` before calling `next`".to_string(),
@@ -78,7 +78,7 @@ impl SourceReader for KinesisSplitReader {
                 continue;
             }
 
-            let mut record_collection: Vec<Vec<u8>> = Vec::new();
+            let mut record_collection: Vec<InnerMessage> = Vec::new();
             for record in records {
                 if !is_stopping(
                     record.sequence_number.as_ref().unwrap(),
@@ -87,17 +87,16 @@ impl SourceReader for KinesisSplitReader {
                     return Ok(Some(record_collection));
                 }
                 self.latest_sequence_num = record.sequence_number().unwrap().to_string();
-                record_collection.push(
-                    KinesisMessage::new(self.shard_id.clone(), record)
-                        .serialize()?
-                        .into_bytes(),
-                );
+                record_collection.push(InnerMessage::from(KinesisMessage::new(
+                    self.shard_id.clone(),
+                    record,
+                )));
             }
             return Ok(Some(record_collection));
         }
     }
 
-    async fn assign_split<'a>(&mut self, split: &'a [u8]) -> Result<()> {
+    async fn assign_split<'a>(&'a mut self, split: &'a [u8]) -> Result<()> {
         let split: KinesisSplit = serde_json::from_str(from_utf8(split)?)?;
         self.shard_id = split.shard_id.clone();
         match &split.start_position.clone() {

rust/connector/src/pulsar/source/message.rs

@@ -1,12 +1,14 @@
 use anyhow::anyhow;
+use pulsar::consumer::Message;
 use serde::{Deserialize, Serialize};
 
-use crate::base::{SourceMessage, SourceOffset};
+use crate::base::{InnerMessage, SourceMessage, SourceOffset};
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct PulsarMessage {
     pub payload: Option<Vec<u8>>,
     pub offset: i64,
+    pub split_id: String,
 }
 
 impl SourceMessage for PulsarMessage {
@@ -23,11 +25,12 @@ impl SourceMessage for PulsarMessage {
     }
 }
 
-impl PulsarMessage {
-    pub fn new(message: pulsar::consumer::Message<Vec<u8>>) -> PulsarMessage {
-        PulsarMessage {
-            payload: Some(message.payload.data),
-            offset: message.message_id.id.entry_id as i64,
+impl From<Message<Vec<u8>>> for InnerMessage {
+    fn from(msg: Message<Vec<u8>>) -> Self {
+        InnerMessage {
+            payload: Some(bytes::Bytes::from(msg.payload.data)),
+            offset: msg.message_id.id.entry_id.to_string(),
+            split_id: msg.topic,
         }
     }
 }

rust/connector/src/pulsar/source/reader.rs

@@ -7,8 +7,7 @@ use async_trait::async_trait;
 use futures::StreamExt;
 use pulsar::{Consumer, Pulsar, SubType, TokioExecutor};
 
-use crate::base::{SourceMessage, SourceReader};
-use crate::pulsar::source::message::PulsarMessage;
+use crate::base::{InnerMessage, SourceReader};
 use crate::pulsar::split::{PulsarOffset, PulsarSplit};
 
 pub struct PulsarSplitReader {
@@ -21,7 +20,7 @@ const PULSAR_MAX_FETCH_MESSAGES: u32 = 1024;
 
 #[async_trait]
 impl SourceReader for PulsarSplitReader {
-    async fn next(&mut self) -> anyhow::Result<Option<Vec<Vec<u8>>>> {
+    async fn next(&mut self) -> anyhow::Result<Option<Vec<InnerMessage>>> {
         let mut stream = self
             .consumer
             .borrow_mut()
@@ -56,13 +55,13 @@ impl SourceReader for PulsarSplitReader {
                 break;
             }
 
-            ret.push(PulsarMessage::new(msg).serialize()?.into_bytes());
+            ret.push(InnerMessage::from(msg));
         }
 
         Ok(Some(ret))
     }
 
-    async fn assign_split<'a>(&mut self, split: &'a [u8]) -> anyhow::Result<()> {
+    async fn assign_split<'a>(&'a mut self, split: &'a [u8]) -> anyhow::Result<()> {
         let split: PulsarSplit = serde_json::from_str(from_utf8(split)?)?;
         let consumer: Consumer<Vec<u8>, TokioExecutor> = self
             .pulsar