refactor: impl ack and migrate to durable consumer for Nats

src/connector/src/error.rs

@@ -21,6 +21,7 @@ use crate::parser::AccessError;
 use crate::schema::schema_registry::{ConcurrentRequestError, WireFormatError};
 use crate::schema::InvalidOptionError;
 use crate::sink::SinkError;
+use crate::source::nats::NatsJetStreamError;
 
 def_anyhow_newtype! {
     pub ConnectorError,
@@ -61,6 +62,8 @@ def_anyhow_newtype! {
     async_nats::jetstream::consumer::pull::MessagesError => "Nats error",
     async_nats::jetstream::context::CreateStreamError => "Nats error",
     async_nats::jetstream::stream::ConsumerError => "Nats error",
+    NatsJetStreamError => "Nats error",
+
     icelake::Error => "Iceberg error",
     iceberg::Error => "IcebergV2 error",
     redis::RedisError => "Redis error",

src/connector/src/source/base.rs

@@ -40,6 +40,7 @@ use super::google_pubsub::GooglePubsubMeta;
 use super::kafka::KafkaMeta;
 use super::kinesis::KinesisMeta;
 use super::monitor::SourceMetrics;
+use super::nats::NatsJetStreamMeta;
 use super::nexmark::source::message::NexmarkMeta;
 use super::{AZBLOB_CONNECTOR, GCS_CONNECTOR, OPENDAL_S3_CONNECTOR, POSIX_FS_CONNECTOR};
 use crate::error::ConnectorResult as Result;
@@ -630,6 +631,7 @@ pub enum SourceMeta {
     GooglePubsub(GooglePubsubMeta),
     Datagen(DatagenMeta),
     DebeziumCdc(DebeziumCdcMeta),
+    NatsJetStream(NatsJetStreamMeta),
     // For the source that doesn't have meta data.
     Empty,
 }

src/connector/src/source/mod.rs

@@ -26,6 +26,8 @@ pub mod nats;
 pub mod nexmark;
 pub mod pulsar;
 
+use std::future::IntoFuture;
+
 pub use base::{UPSTREAM_SOURCE_KEY, *};
 pub(crate) use common::*;
 use google_cloud_pubsub::subscription::Subscription;
@@ -40,6 +42,8 @@ mod manager;
 pub mod reader;
 pub mod test_source;
 
+use async_nats::jetstream::consumer::AckPolicy as JetStreamAckPolicy;
+use async_nats::jetstream::context::Context as JetStreamContext;
 pub use manager::{SourceColumnDesc, SourceColumnType};
 use risingwave_common::array::{Array, ArrayRef};
 use thiserror_ext::AsReport;
@@ -77,6 +81,7 @@ pub fn should_copy_to_format_encode_options(key: &str, connector: &str) -> bool
 pub enum WaitCheckpointTask {
     CommitCdcOffset(Option<(SplitId, String)>),
     AckPubsubMessage(Subscription, Vec<ArrayRef>),
+    AckNatsJetStream(JetStreamContext, Vec<ArrayRef>, JetStreamAckPolicy),
 }
 
 impl WaitCheckpointTask {
@@ -123,6 +128,49 @@ impl WaitCheckpointTask {
                 }
                 ack(&subscription, ack_ids).await;
             }
+            WaitCheckpointTask::AckNatsJetStream(
+                ref context,
+                reply_subjects_arrs,
+                ref ack_policy,
+            ) => {
+                async fn ack(context: &JetStreamContext, reply_subject: String) {
+                    match context.publish(reply_subject.clone(), "".into()).await {
+                        Err(e) => {
+                            tracing::error!(error = %e.as_report(), subject = ?reply_subject, "failed to ack NATS JetStream message");
+                        }
+                        Ok(ack_future) => {
+                            if let Err(e) = ack_future.into_future().await {
+                                tracing::error!(error = %e.as_report(), subject = ?reply_subject, "failed to ack NATS JetStream message");
+                            }
+                        }
+                    }
+                }
+
+                let reply_subjects = reply_subjects_arrs
+                    .iter()
+                    .flat_map(|arr| {
+                        arr.as_utf8()
+                            .iter()
+                            .flatten()
+                            .map(|s| s.to_string())
+                            .collect::<Vec<String>>()
+                    })
+                    .collect::<Vec<String>>();
+
+                match ack_policy {
+                    JetStreamAckPolicy::None => (),
+                    JetStreamAckPolicy::Explicit => {
+                        for reply_subject in reply_subjects {
+                            ack(context, reply_subject).await;
+                        }
+                    }
+                    JetStreamAckPolicy::All => {
+                        if let Some(reply_subject) = reply_subjects.last() {
+                            ack(context, reply_subject.clone()).await;
+                        }
+                    }
+                }
+            }
         }
     }
 }

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

@@ -23,39 +23,55 @@ use async_nats::jetstream::consumer::pull::Config;
 use async_nats::jetstream::consumer::{AckPolicy, ReplayPolicy};
 use serde::Deserialize;
 use serde_with::{serde_as, DisplayFromStr};
+use strum_macros::Display;
+use thiserror::Error;
 use with_options::WithOptions;
 
 use crate::connector_common::NatsCommon;
+use crate::error::{ConnectorError, ConnectorResult};
 use crate::source::nats::enumerator::NatsSplitEnumerator;
+pub use crate::source::nats::source::NatsJetStreamMeta;
 use crate::source::nats::source::{NatsSplit, NatsSplitReader};
 use crate::source::SourceProperties;
 use crate::{
     deserialize_optional_string_seq_from_string, deserialize_optional_u64_seq_from_string,
 };
 
+#[derive(Debug, Clone, Error, Display)]
+pub enum NatsJetStreamError {
+    InvalidAckPolicy(String),
+    InvalidReplayPolicy(String),
+}
+
 pub const NATS_CONNECTOR: &str = "nats";
 
 pub struct AckPolicyWrapper;
 
 impl AckPolicyWrapper {
-    pub fn parse_str(s: &str) -> Result<AckPolicy, String> {
+    pub fn parse_str(s: &str) -> Result<AckPolicy, NatsJetStreamError> {
         match s {
             "none" => Ok(AckPolicy::None),
             "all" => Ok(AckPolicy::All),
             "explicit" => Ok(AckPolicy::Explicit),
-            _ => Err(format!("Invalid AckPolicy '{}'", s)),
+            _ => Err(NatsJetStreamError::InvalidAckPolicy(format!(
+                "Invalid AckPolicy '{}', expect `none`, `all`, and `explicit`",
+                s
+            ))),
         }
     }
 }
 
 pub struct ReplayPolicyWrapper;
 
 impl ReplayPolicyWrapper {
-    pub fn parse_str(s: &str) -> Result<ReplayPolicy, String> {
+    pub fn parse_str(s: &str) -> Result<ReplayPolicy, NatsJetStreamError> {
         match s {
             "instant" => Ok(ReplayPolicy::Instant),
             "original" => Ok(ReplayPolicy::Original),
-            _ => Err(format!("Invalid ReplayPolicy '{}'", s)),
+            _ => Err(NatsJetStreamError::InvalidReplayPolicy(format!(
+                "Invalid ReplayPolicy '{}', expect `instant` and `original`",
+                s
+            ))),
         }
     }
 }
@@ -257,6 +273,13 @@ impl NatsPropertiesConsumer {
             c.backoff = v.iter().map(|&x| Duration::from_secs(x)).collect()
         }
     }
+
+    pub fn get_ack_policy(&self) -> ConnectorResult<AckPolicy> {
+        match &self.ack_policy {
+            Some(policy) => Ok(AckPolicyWrapper::parse_str(policy).map_err(ConnectorError::from)?),
+            None => Ok(AckPolicy::None),
+        }
+    }
 }
 
 impl SourceProperties for NatsProperties {

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

@@ -22,6 +22,12 @@ pub struct NatsMessage {
     pub split_id: SplitId,
     pub sequence_number: String,
     pub payload: Vec<u8>,
+    pub reply_subject: Option<String>,
+}
+
+#[derive(Clone, Debug)]
+pub struct NatsJetStreamMeta {
+    pub reply_subject: Option<String>,
 }
 
 impl From<NatsMessage> for SourceMessage {
@@ -30,9 +36,12 @@ impl From<NatsMessage> for SourceMessage {
             key: None,
             payload: Some(message.payload),
             // For nats jetstream, use sequence id as offset
+            // DEPRECATED: no longer use sequence id as offset, let nats broker handle failover
             offset: message.sequence_number,
             split_id: message.split_id,
-            meta: SourceMeta::Empty,
+            meta: SourceMeta::NatsJetStream(NatsJetStreamMeta {
+                reply_subject: message.reply_subject,
+            }),
         }
     }
 }
@@ -43,6 +52,10 @@ impl NatsMessage {
             split_id,
             sequence_number: message.info().unwrap().stream_sequence.to_string(),
             payload: message.message.payload.to_vec(),
+            reply_subject: message
+                .message
+                .reply
+                .map(|subject| subject.as_str().to_string()),
         }
     }
 }

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

@@ -15,6 +15,7 @@
 mod message;
 mod reader;
 
+pub use message::*;
 pub use reader::*;
 
 pub use crate::source::nats::split::*;

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

@@ -127,6 +127,11 @@ impl SourceReader {
                 prop.subscription_client().await?,
                 vec![],
             )),
+            ConnectorProperties::Nats(prop) => Some(WaitCheckpointTask::AckNatsJetStream(
+                prop.common.build_context().await?,
+                vec![],
+                prop.nats_properties_consumer.get_ack_policy()?,
+            )),
             _ => None,
         })
     }

src/stream/src/executor/source/source_executor.rs

@@ -752,12 +752,14 @@ struct WaitCheckpointTaskBuilder {
 
 impl WaitCheckpointTaskBuilder {
     fn update_task_on_chunk(&mut self, offset_col: ArrayRef) {
-        #[expect(clippy::single_match)]
         match &mut self.building_task {
             WaitCheckpointTask::AckPubsubMessage(_, arrays) => {
                 arrays.push(offset_col);
             }
-            _ => {}
+            WaitCheckpointTask::AckNatsJetStream(_, arrays, _) => {
+                arrays.push(offset_col);
+            }
+            WaitCheckpointTask::CommitCdcOffset(_) => {}
         }
     }