feat(stream): Add recovery feature for nats source (risingwavelabs#12073)

Signed-off-by: tabVersion <[email protected]>
Co-authored-by: tabVersion <[email protected]>

Author: yufansong

src/connector/src/common.rs

@@ -17,6 +17,7 @@ use std::collections::HashMap;
 use std::time::Duration;
 
 use anyhow::Ok;
+use async_nats::jetstream::consumer::DeliverPolicy;
 use async_nats::jetstream::{self};
 use aws_sdk_kinesis::Client as KinesisClient;
 use clickhouse::Client;
@@ -386,6 +387,39 @@ impl NatsCommon {
         Ok(subscription)
     }
 
+    pub(crate) async fn build_consumer(
+        &self,
+        split_id: i32,
+        start_sequence: Option<u64>,
+    ) -> anyhow::Result<
+        async_nats::jetstream::consumer::Consumer<async_nats::jetstream::consumer::pull::Config>,
+    > {
+        let context = self.build_context().await?;
+        let stream = self.build_or_get_stream(context.clone()).await?;
+        let name = format!("risingwave-consumer-{}-{}", self.subject, split_id);
+        let mut config = jetstream::consumer::pull::Config {
+            ack_policy: jetstream::consumer::AckPolicy::None,
+            ..Default::default()
+        };
+        match start_sequence {
+            Some(v) => {
+                let consumer = stream
+                    .get_or_create_consumer(&name, {
+                        config.deliver_policy = DeliverPolicy::ByStartSequence {
+                            start_sequence: v + 1,
+                        };
+                        config
+                    })
+                    .await?;
+                Ok(consumer)
+            }
+            None => {
+                let consumer = stream.get_or_create_consumer(&name, config).await?;
+                Ok(consumer)
+            }
+        }
+    }
+
     pub(crate) async fn build_or_get_stream(
         &self,
         jetstream: jetstream::Context,

src/connector/src/source/nats/enumerator/mod.rs

@@ -44,7 +44,8 @@ impl SplitEnumerator for NatsSplitEnumerator {
         // TODO: to simplify the logic, return 1 split for first version
         let nats_split = NatsSplit {
             subject: self.subject.clone(),
-            split_num: 1,
+            split_num: 0, // be the same as `from_nats_jetstream_message`
+            start_sequence: None,
         };
 
         Ok(vec![nats_split])

src/connector/src/source/nats/source/message.rs

@@ -12,21 +12,19 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::sync::Arc;
-
 use async_nats;
 
 use crate::source::base::SourceMessage;
 use crate::source::SourceMeta;
 
 impl SourceMessage {
-    pub fn from_nats_message(message: async_nats::Message) -> Self {
+    pub fn from_nats_jetstream_message(message: async_nats::jetstream::message::Message) -> Self {
         SourceMessage {
             key: None,
-            payload: Some(message.payload.to_vec()),
-            // Nats message doesn't have offset
-            offset: "".to_string(),
-            split_id: Arc::from(""),
+            payload: Some(message.message.payload.to_vec()),
+            // For nats jetstream, use sequence id as offset
+            offset: message.info().unwrap().stream_sequence.to_string(),
+            split_id: "0".into(),
             meta: SourceMeta::Empty,
         }
     }

src/connector/src/source/nats/source/reader.rs

@@ -13,19 +13,21 @@
 // limitations under the License.
 
 use anyhow::Result;
+use async_nats::jetstream::consumer;
 use async_trait::async_trait;
 use futures::StreamExt;
 use futures_async_stream::try_stream;
 
 use crate::parser::ParserConfig;
 use crate::source::common::{into_chunk_stream, CommonSplitReader};
+use crate::source::nats::split::NatsSplit;
 use crate::source::nats::NatsProperties;
 use crate::source::{
     BoxSourceWithStateStream, Column, SourceContextRef, SourceMessage, SplitImpl, SplitReader,
 };
 
 pub struct NatsSplitReader {
-    subscriber: async_nats::Subscriber,
+    consumer: consumer::Consumer<consumer::pull::Config>,
     properties: NatsProperties,
     parser_config: ParserConfig,
     source_ctx: SourceContextRef,
@@ -44,9 +46,16 @@ impl SplitReader for NatsSplitReader {
     ) -> Result<Self> {
         // TODO: to simplify the logic, return 1 split for first version
         assert!(splits.len() == 1);
-        let subscriber = properties.common.build_subscriber().await?;
+        let splits = splits
+            .into_iter()
+            .map(|split| split.into_nats().unwrap())
+            .collect::<Vec<NatsSplit>>();
+        let consumer = properties
+            .common
+            .build_consumer(0, splits[0].start_sequence)
+            .await?;
         Ok(Self {
-            subscriber,
+            consumer,
             properties,
             parser_config,
             source_ctx,
@@ -64,11 +73,12 @@ impl CommonSplitReader for NatsSplitReader {
     #[try_stream(ok = Vec<SourceMessage>, error = anyhow::Error)]
     async fn into_data_stream(self) {
         let capacity = self.source_ctx.source_ctrl_opts.chunk_size;
+        let messages = self.consumer.messages().await?;
         #[for_await]
-        for msgs in self.subscriber.ready_chunks(capacity) {
+        for msgs in messages.ready_chunks(capacity) {
             let mut msg_vec = Vec::with_capacity(capacity);
             for msg in msgs {
-                msg_vec.push(SourceMessage::from_nats_message(msg));
+                msg_vec.push(SourceMessage::from_nats_jetstream_message(msg?));
             }
             yield msg_vec;
         }

src/connector/src/source/nats/split.rs

@@ -25,7 +25,7 @@ pub struct NatsSplit {
     // TODO: to simplify the logic, return 1 split for first version. May use parallelism in
     // future.
     pub(crate) split_num: i32,
-    // nats does not provide offset
+    pub(crate) start_sequence: Option<u64>,
 }
 
 impl SplitMetaData for NatsSplit {
@@ -44,11 +44,16 @@ impl SplitMetaData for NatsSplit {
 }
 
 impl NatsSplit {
-    pub fn new(subject: String, split_num: i32) -> Self {
-        Self { subject, split_num }
+    pub fn new(subject: String, split_num: i32, start_sequence: Option<u64>) -> Self {
+        Self {
+            subject,
+            split_num,
+            start_sequence,
+        }
     }
 
-    pub fn update_with_offset(&mut self, _start_offset: String) -> anyhow::Result<()> {
+    pub fn update_with_offset(&mut self, start_sequence: String) -> anyhow::Result<()> {
+        self.start_sequence = Some(start_sequence.as_str().parse::<u64>().unwrap());
         Ok(())
     }
 }