feat(new sink): Add possibility to use nats jetstream in nats sink

src/sinks/nats/config.rs

@@ -1,3 +1,4 @@
+use bytes::Bytes;
 use futures_util::TryFutureExt;
 use snafu::ResultExt;
 use vector_lib::codecs::JsonSerializerConfig;
@@ -76,6 +77,10 @@ pub struct NatsSinkConfig {
     #[configurable(derived)]
     #[serde(default)]
     pub(super) request: TowerRequestConfig<NatsTowerRequestConfigDefaults>,
+
+    #[configurable(derived)]
+    #[serde(default)]
+    pub(super) jetstream: bool,
 }
 
 fn default_name() -> String {
@@ -93,6 +98,7 @@ impl GenerateConfig for NatsSinkConfig {
             tls: None,
             url: "nats://127.0.0.1:4222".into(),
             request: Default::default(),
+            jetstream: Default::default(),
         })
         .unwrap()
     }
@@ -130,8 +136,56 @@ impl NatsSinkConfig {
 
         options.connect(&self.url).await.context(ConnectSnafu)
     }
+
+    pub(super) async fn publisher(&self) -> Result<NatsPublisher, NatsError> {
+        let connection = self.connect().await?;
+
+        match self.jetstream {
+            true => Ok(NatsPublisher::JetStream(async_nats::jetstream::new(
+                connection,
+            ))),
+            false => Ok(NatsPublisher::Core(connection)),
+        }
+    }
 }
 
 async fn healthcheck(config: NatsSinkConfig) -> crate::Result<()> {
     config.connect().map_ok(|_| ()).map_err(|e| e.into()).await
 }
+
+pub enum NatsPublisher {
+    Core(async_nats::Client),
+    JetStream(async_nats::jetstream::Context),
+}
+
+impl NatsPublisher {
+    pub(super) async fn publish<S: async_nats::subject::ToSubject>(
+        &self,
+        subject: S,
+        payload: Bytes,
+    ) -> Result<(), NatsError> {
+        match self {
+            NatsPublisher::Core(client) => {
+                client
+                    .publish(subject, payload)
+                    .await
+                    .map_err(|e| NatsError::PublishError {
+                        source: Box::new(e),
+                    })?;
+                client.flush().await.map_err(|e| NatsError::PublishError {
+                    source: Box::new(e),
+                })
+            }
+            NatsPublisher::JetStream(jetstream) => {
+                let ack = jetstream.publish(subject, payload).await.map_err(|e| {
+                    NatsError::PublishError {
+                        source: Box::new(e),
+                    }
+                })?;
+                ack.await.map(|_| ()).map_err(|e| NatsError::PublishError {
+                    source: Box::new(e),
+                })
+            }
+        }
+    }
+}

src/sinks/nats/mod.rs

@@ -27,4 +27,6 @@ enum NatsError {
     Connect { source: async_nats::ConnectError },
     #[snafu(display("NATS Server Error: {}", source))]
     ServerError { source: async_nats::Error },
+    #[snafu(display("NATS Publish Error: {}", source))]
+    PublishError { source: async_nats::Error },
 }

src/sinks/nats/service.rs

@@ -7,11 +7,11 @@ use futures_util::TryFutureExt;
 
 use crate::sinks::prelude::*;
 
-use super::{request_builder::NatsRequest, NatsError};
+use super::{config::NatsPublisher, request_builder::NatsRequest, NatsError};
 
 #[derive(Clone)]
 pub(super) struct NatsService {
-    pub(super) connection: Arc<async_nats::Client>,
+    pub(super) publisher: Arc<NatsPublisher>,
 }
 
 pub(super) struct NatsResponse {
@@ -44,13 +44,12 @@ impl Service<NatsRequest> for NatsService {
     }
 
     fn call(&mut self, req: NatsRequest) -> Self::Future {
-        let connection = Arc::clone(&self.connection);
+        let publisher = Arc::clone(&self.publisher);
 
         Box::pin(async move {
-            match connection
+            match publisher
                 .publish(req.subject, req.bytes)
                 .map_err(async_nats::Error::from)
-                .and_then(|_| connection.flush().map_err(Into::into))
                 .await
             {
                 Err(error) => Err(NatsError::ServerError { source: error }),

src/sinks/nats/sink.rs

@@ -5,7 +5,7 @@ use snafu::ResultExt;
 use crate::sinks::prelude::*;
 
 use super::{
-    config::{NatsSinkConfig, NatsTowerRequestConfigDefaults},
+    config::{NatsPublisher, NatsSinkConfig, NatsTowerRequestConfigDefaults},
     request_builder::{NatsEncoder, NatsRequestBuilder},
     service::{NatsResponse, NatsService},
     EncodingSnafu, NatsError,
@@ -20,7 +20,7 @@ pub(super) struct NatsSink {
     request: TowerRequestConfig<NatsTowerRequestConfigDefaults>,
     transformer: Transformer,
     encoder: Encoder<()>,
-    connection: Arc<async_nats::Client>,
+    publisher: Arc<NatsPublisher>,
     subject: Template,
 }
 
@@ -42,7 +42,7 @@ impl NatsSink {
     }
 
     pub(super) async fn new(config: NatsSinkConfig) -> Result<Self, NatsError> {
-        let connection = Arc::new(config.connect().await?);
+        let publisher = Arc::new(config.publisher().await?);
         let transformer = config.encoding.transformer();
         let serializer = config.encoding.build().context(EncodingSnafu)?;
         let encoder = Encoder::<()>::new(serializer);
@@ -51,9 +51,9 @@ impl NatsSink {
 
         Ok(NatsSink {
             request,
-            connection,
             transformer,
             encoder,
+            publisher,
             subject,
         })
     }
@@ -71,7 +71,7 @@ impl NatsSink {
         let service = ServiceBuilder::new()
             .settings(request, NatsRetryLogic)
             .service(NatsService {
-                connection: Arc::clone(&self.connection),
+                publisher: Arc::clone(&self.publisher),
             });
 
         input