Skip to main content

dynamo_runtime/transports/
nats.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! NATS transport
5//!
6//! The following environment variables are used to configure the NATS client:
7//!
8//! - `NATS_SERVER`: the NATS server address
9//!
10//! For authentication, the following environment variables are used and prioritized in the following order:
11//!
12//! - `NATS_AUTH_USERNAME`: the username for authentication
13//! - `NATS_AUTH_PASSWORD`: the password for authentication
14//! - `NATS_AUTH_TOKEN`: the token for authentication
15//! - `NATS_AUTH_NKEY`: the nkey for authentication
16//! - `NATS_AUTH_CREDENTIALS_FILE`: the path to the credentials file
17//!
18//! Note: `NATS_AUTH_USERNAME` and `NATS_AUTH_PASSWORD` must be used together.
19//!
20//! ## TLS
21//!
22//! A custom TLS config is applied when `NATS_TLS_CA_CERT_PATH` is set,
23//! `NATS_TLS_INSECURE` is truthy, or a client certificate is configured. When
24//! only the `tls://` URL scheme is used without explicit TLS env vars,
25//! async-nats handles TLS natively with system roots.
26//!
27//! - `NATS_TLS_CA_CERT_PATH`: path to the CA cert PEM used to verify the server
28//! - `NATS_TLS_CLIENT_CERT_PATH`: client cert PEM for mutual TLS (optional)
29//! - `NATS_TLS_CLIENT_KEY_PATH`: client key PEM for mutual TLS (optional)
30//! - `NATS_TLS_INSECURE`: set to a truthy value to skip certificate verification (dev only)
31use crate::metrics::MetricsHierarchy;
32use crate::protocols::EndpointId;
33
34use anyhow::Result;
35use async_nats::connection::State;
36use async_nats::{Subscriber, client, jetstream};
37use async_trait::async_trait;
38use bytes::Bytes;
39use derive_builder::Builder;
40use futures::{StreamExt, TryStreamExt};
41use prometheus::{Counter, Gauge, Histogram, HistogramOpts, IntCounter, IntGauge, Opts, Registry};
42use serde::de::DeserializeOwned;
43use serde::{Deserialize, Serialize};
44use std::path::{Path, PathBuf};
45use std::sync::atomic::Ordering;
46use tokio::fs::File as TokioFile;
47use tokio::io::AsyncRead;
48use tokio::time;
49use url::Url;
50use validator::{Validate, ValidationError};
51
52use crate::config::environment_names::nats as env_nats;
53pub use crate::slug::Slug;
54use tracing as log;
55
56use super::utils::build_in_runtime;
57
58pub const URL_PREFIX: &str = "nats://";
59
60#[derive(Clone)]
61pub struct Client {
62    client: client::Client,
63    js_ctx: jetstream::Context,
64}
65
66impl Client {
67    /// Create a NATS [`ClientOptionsBuilder`].
68    pub fn builder() -> ClientOptionsBuilder {
69        ClientOptionsBuilder::default()
70    }
71
72    /// Returns a reference to the underlying [`async_nats::client::Client`] instance
73    pub fn client(&self) -> &client::Client {
74        &self.client
75    }
76
77    /// Returns a reference to the underlying [`async_nats::jetstream::Context`] instance
78    pub fn jetstream(&self) -> &jetstream::Context {
79        &self.js_ctx
80    }
81
82    /// host:port of NATS
83    pub fn addr(&self) -> String {
84        let info = self.client.server_info();
85        format!("{}:{}", info.host, info.port)
86    }
87
88    /// fetch the list of streams
89    pub async fn list_streams(&self) -> Result<Vec<String>> {
90        let names = self.js_ctx.stream_names();
91        let stream_names: Vec<String> = names.try_collect().await?;
92        Ok(stream_names)
93    }
94
95    /// fetch the list of consumers for a given stream
96    pub async fn list_consumers(&self, stream_name: &str) -> Result<Vec<String>> {
97        let stream = self.js_ctx.get_stream(stream_name).await?;
98        let consumers: Vec<String> = stream.consumer_names().try_collect().await?;
99        Ok(consumers)
100    }
101
102    pub async fn stream_info(&self, stream_name: &str) -> Result<jetstream::stream::State> {
103        let mut stream = self.js_ctx.get_stream(stream_name).await?;
104        let info = stream.info().await?;
105        Ok(info.state.clone())
106    }
107
108    pub async fn get_stream(&self, name: &str) -> Result<jetstream::stream::Stream> {
109        let stream = self.js_ctx.get_stream(name).await?;
110        Ok(stream)
111    }
112
113    /// Issues a broadcast request for all services with the provided `service_name` to report their
114    /// current stats. Each service will only respond once. The service may have customized the reply
115    /// so the caller should select which endpoint and what concrete data model should be used to
116    /// extract the details.
117    ///
118    /// Note: Because each endpoint will only reply once, the caller must drop the subscription after
119    /// some time or it will await forever.
120    pub async fn scrape_service(&self, service_name: &str) -> Result<Subscriber> {
121        let subject = format!("$SRV.STATS.{}", service_name);
122        let reply_subject = format!("_INBOX.{}", nuid::next());
123        let subscription = self.client.subscribe(reply_subject.clone()).await?;
124
125        // Publish the request with the reply-to subject
126        self.client
127            .publish_with_reply(subject, reply_subject, "".into())
128            .await?;
129
130        Ok(subscription)
131    }
132
133    /// Helper method to get or optionally create an object store bucket
134    ///
135    /// # Arguments
136    /// * `bucket_name` - The name of the bucket to retrieve
137    /// * `create_if_not_found` - If true, creates the bucket when it doesn't exist
138    ///
139    /// # Returns
140    /// The object store bucket or an error
141    async fn get_or_create_bucket(
142        &self,
143        bucket_name: &str,
144        create_if_not_found: bool,
145    ) -> anyhow::Result<jetstream::object_store::ObjectStore> {
146        let context = self.jetstream();
147
148        match context.get_object_store(bucket_name).await {
149            Ok(bucket) => Ok(bucket),
150            Err(err) if err.to_string().contains("stream not found") => {
151                // err.source() is GetStreamError, which has a kind() which
152                // is GetStreamErrorKind::JetStream which wraps a jetstream::Error
153                // which has code 404. Phew. So yeah check the string for now.
154
155                if create_if_not_found {
156                    tracing::debug!("Creating NATS bucket {bucket_name}");
157                    context
158                        .create_object_store(jetstream::object_store::Config {
159                            bucket: bucket_name.to_string(),
160                            ..Default::default()
161                        })
162                        .await
163                        .map_err(|e| anyhow::anyhow!("Failed creating bucket / object store: {e}"))
164                } else {
165                    anyhow::bail!(
166                        "NATS get_object_store bucket does not exist: {bucket_name}. {err}."
167                    );
168                }
169            }
170            Err(err) => {
171                anyhow::bail!("NATS get_object_store error: {err}");
172            }
173        }
174    }
175
176    /// Upload file to NATS at this URL
177    pub async fn object_store_upload(&self, filepath: &Path, nats_url: &Url) -> anyhow::Result<()> {
178        let mut disk_file = TokioFile::open(filepath).await?;
179
180        let (bucket_name, key) = url_to_bucket_and_key(nats_url)?;
181        let bucket = self.get_or_create_bucket(&bucket_name, true).await?;
182
183        let key_meta = async_nats::jetstream::object_store::ObjectMetadata {
184            name: key.to_string(),
185            ..Default::default()
186        };
187        bucket.put(key_meta, &mut disk_file).await.map_err(|e| {
188            anyhow::anyhow!("Failed uploading to bucket / object store {bucket_name}/{key}: {e}")
189        })?;
190
191        Ok(())
192    }
193
194    /// Download file from NATS at this URL
195    pub async fn object_store_download(
196        &self,
197        nats_url: &Url,
198        filepath: &Path,
199    ) -> anyhow::Result<()> {
200        let mut disk_file = TokioFile::create(filepath).await?;
201
202        let (bucket_name, key) = url_to_bucket_and_key(nats_url)?;
203        let bucket = self.get_or_create_bucket(&bucket_name, false).await?;
204
205        let mut obj_reader = bucket.get(&key).await.map_err(|e| {
206            anyhow::anyhow!(
207                "Failed downloading from bucket / object store {bucket_name}/{key}: {e}"
208            )
209        })?;
210        let _bytes_copied = tokio::io::copy(&mut obj_reader, &mut disk_file).await?;
211
212        Ok(())
213    }
214
215    /// Delete a bucket and all it's contents from the NATS object store
216    pub async fn object_store_delete_bucket(&self, bucket_name: &str) -> anyhow::Result<()> {
217        let context = self.jetstream();
218        match context.delete_object_store(&bucket_name).await {
219            Ok(_) => Ok(()),
220            Err(err) if err.to_string().contains("stream not found") => {
221                tracing::trace!(bucket_name, "NATS bucket already gone");
222                Ok(())
223            }
224            Err(err) => Err(anyhow::anyhow!("NATS get_object_store error: {err}")),
225        }
226    }
227
228    /// Upload a serializable struct to NATS object store using bincode
229    pub async fn object_store_upload_data<T>(&self, data: &T, nats_url: &Url) -> anyhow::Result<()>
230    where
231        T: Serialize,
232    {
233        // Serialize the data using bincode (more efficient binary format)
234        let binary_data = bincode::serialize(data)
235            .map_err(|e| anyhow::anyhow!("Failed to serialize data with bincode: {e}"))?;
236
237        let (bucket_name, key) = url_to_bucket_and_key(nats_url)?;
238        let bucket = self.get_or_create_bucket(&bucket_name, true).await?;
239
240        let key_meta = async_nats::jetstream::object_store::ObjectMetadata {
241            name: key.to_string(),
242            ..Default::default()
243        };
244
245        // Upload the serialized bytes
246        let mut cursor = std::io::Cursor::new(binary_data);
247        bucket.put(key_meta, &mut cursor).await.map_err(|e| {
248            anyhow::anyhow!("Failed uploading to bucket / object store {bucket_name}/{key}: {e}")
249        })?;
250
251        Ok(())
252    }
253
254    /// Download and deserialize a struct from NATS object store using bincode
255    pub async fn object_store_download_data<T>(&self, nats_url: &Url) -> anyhow::Result<T>
256    where
257        T: DeserializeOwned,
258    {
259        let (bucket_name, key) = url_to_bucket_and_key(nats_url)?;
260        let bucket = self.get_or_create_bucket(&bucket_name, false).await?;
261
262        let mut obj_reader = bucket.get(&key).await.map_err(|e| {
263            anyhow::anyhow!(
264                "Failed downloading from bucket / object store {bucket_name}/{key}: {e}"
265            )
266        })?;
267
268        // Read all bytes into memory
269        let mut buffer = Vec::new();
270        tokio::io::copy(&mut obj_reader, &mut buffer)
271            .await
272            .map_err(|e| anyhow::anyhow!("Failed reading object data: {e}"))?;
273        tracing::debug!("Downloaded {} bytes from {bucket_name}/{key}", buffer.len());
274
275        // Deserialize from bincode
276        let data = bincode::deserialize(&buffer)
277            .map_err(|e| anyhow::anyhow!("Failed to deserialize data with bincode: {e}"))?;
278
279        Ok(data)
280    }
281}
282
283/// NATS client options
284///
285/// This object uses the builder pattern with default values that are evaluates
286/// from the environment variables if they are not explicitly set by the builder.
287#[derive(Debug, Clone, Builder, Validate)]
288pub struct ClientOptions {
289    #[builder(setter(into), default = "default_server()")]
290    #[validate(custom(function = "validate_nats_server"))]
291    server: String,
292
293    #[builder(default)]
294    auth: NatsAuth,
295
296    /// Path to PEM CA certificate for TLS. When set, TLS is required and
297    /// `NATS_SERVER` must use the `tls://` scheme.
298    #[builder(default = "default_nats_tls_ca_cert_path()")]
299    tls_ca_cert_path: Option<PathBuf>,
300
301    /// Path to PEM client certificate presented to the NATS server for mutual
302    /// TLS (mTLS). Must be set together with `tls_client_key_path`.
303    #[builder(default = "default_nats_tls_client_cert_path()")]
304    tls_client_cert_path: Option<PathBuf>,
305
306    /// Path to PEM client private key for mutual TLS (mTLS).
307    #[builder(default = "default_nats_tls_client_key_path()")]
308    tls_client_key_path: Option<PathBuf>,
309
310    /// Skip TLS certificate verification. For development only.
311    #[builder(default = "default_nats_tls_insecure()")]
312    tls_insecure: bool,
313}
314
315fn default_server() -> String {
316    if let Ok(server) = std::env::var(env_nats::NATS_SERVER) {
317        return server;
318    }
319
320    "nats://localhost:4222".to_string()
321}
322
323fn validate_nats_server(server: &str) -> Result<(), ValidationError> {
324    if server.starts_with("nats://") || server.starts_with("tls://") {
325        Ok(())
326    } else {
327        Err(ValidationError::new(
328            "server must start with 'nats://' or 'tls://'",
329        ))
330    }
331}
332
333fn default_nats_tls_ca_cert_path() -> Option<PathBuf> {
334    std::env::var(env_nats::tls::NATS_TLS_CA_CERT_PATH)
335        .ok()
336        .map(PathBuf::from)
337}
338
339fn default_nats_tls_client_cert_path() -> Option<PathBuf> {
340    std::env::var(env_nats::tls::NATS_TLS_CLIENT_CERT_PATH)
341        .ok()
342        .map(PathBuf::from)
343}
344
345fn default_nats_tls_client_key_path() -> Option<PathBuf> {
346    std::env::var(env_nats::tls::NATS_TLS_CLIENT_KEY_PATH)
347        .ok()
348        .map(PathBuf::from)
349}
350
351fn default_nats_tls_insecure() -> bool {
352    crate::config::env_is_truthy(env_nats::tls::NATS_TLS_INSECURE)
353}
354
355// TODO(jthomson04): We really shouldn't be hardcoding this.
356const NATS_WORKER_THREADS: usize = 4;
357
358impl ClientOptions {
359    /// Create a new [`ClientOptionsBuilder`]
360    pub fn builder() -> ClientOptionsBuilder {
361        ClientOptionsBuilder::default()
362    }
363
364    /// Validate the config and attempt to connection to the NATS server
365    pub async fn connect(self) -> Result<Client> {
366        self.validate()?;
367
368        // Client cert and key must be set together to present a client identity.
369        if self.tls_client_cert_path.is_some() != self.tls_client_key_path.is_some() {
370            anyhow::bail!(
371                "Both {} and {} must be set together to enable NATS mTLS",
372                env_nats::tls::NATS_TLS_CLIENT_CERT_PATH,
373                env_nats::tls::NATS_TLS_CLIENT_KEY_PATH,
374            );
375        }
376
377        // A client identity requires a CA (or insecure) so the server can still
378        // be verified; otherwise the root store would be empty and verification
379        // would fail with an opaque error.
380        if self.tls_client_cert_path.is_some()
381            && self.tls_ca_cert_path.is_none()
382            && !self.tls_insecure
383        {
384            anyhow::bail!(
385                "{} requires {} (or {}) to also be set",
386                env_nats::tls::NATS_TLS_CLIENT_CERT_PATH,
387                env_nats::tls::NATS_TLS_CA_CERT_PATH,
388                env_nats::tls::NATS_TLS_INSECURE,
389            );
390        }
391
392        let custom_tls = self.tls_ca_cert_path.is_some()
393            || self.tls_insecure
394            || self.tls_client_cert_path.is_some();
395        let tls_url = self.server.starts_with("tls://");
396
397        // Custom TLS settings imply an encrypted connection, so the server URL
398        // must use the tls:// scheme. Reject the mismatch up front with a clear
399        // error instead of silently forcing TLS onto a nats:// URL.
400        if custom_tls && !tls_url {
401            anyhow::bail!(
402                "NATS TLS is configured (NATS_TLS_CA_CERT_PATH, NATS_TLS_INSECURE, or a client \
403                 certificate) but NATS_SERVER does not use the 'tls://' scheme: {}",
404                self.server
405            );
406        }
407
408        let mut options = match self.auth {
409            NatsAuth::UserPass(username, password) => {
410                async_nats::ConnectOptions::with_user_and_password(username, password)
411            }
412            NatsAuth::Token(token) => async_nats::ConnectOptions::with_token(token),
413            NatsAuth::NKey(nkey) => async_nats::ConnectOptions::with_nkey(nkey),
414            NatsAuth::CredentialsFile(path) => {
415                async_nats::ConnectOptions::with_credentials_file(path).await?
416            }
417        };
418
419        // Install the ring crypto provider as the process-level default, but
420        // only when this connection actually uses TLS. async-nats calls
421        // ClientConfig::builder() internally for a tls:// URL and panics if no
422        // provider is installed; both ring and aws-lc-rs are compiled in (via
423        // async-nats and kube respectively), so rustls 0.23 cannot auto-detect.
424        // Gating on TLS avoids clobbering another component's provider choice
425        // (e.g. the HTTPS frontend's aws-lc-rs) for plaintext nats:// use.
426        // Silently ignored if a provider is already installed.
427        if custom_tls || tls_url {
428            let _ = rustls::crypto::ring::default_provider().install_default();
429        }
430
431        // Apply a custom TLS config when a CA cert or insecure mode is explicitly
432        // configured. When only a tls:// URL is used without explicit TLS env vars,
433        // let async-nats handle TLS natively (it uses its own rustls setup with
434        // system roots).
435        if custom_tls {
436            let tls_config = crate::tls_utils::client_tls_config(
437                self.tls_ca_cert_path.as_deref(),
438                self.tls_insecure,
439                self.tls_client_cert_path.as_deref(),
440                self.tls_client_key_path.as_deref(),
441            )?;
442            options = options.tls_client_config(tls_config).require_tls(true);
443        } else if tls_url {
444            // tls:// URL implies TLS but no custom CA — async-nats will use its
445            // built-in rustls with system roots. Just require TLS on the connection.
446            options = options.require_tls(true);
447        }
448
449        // 0 is treated as unset — Duration::from_secs(0) would time out every request immediately.
450        let request_timeout = std::env::var(env_nats::DYN_NATS_REQUEST_TIMEOUT_SECS)
451            .ok()
452            .and_then(|v| v.parse::<u64>().ok())
453            .filter(|&secs| secs > 0)
454            .map(time::Duration::from_secs);
455        let options = match request_timeout {
456            Some(timeout) => options.request_timeout(Some(timeout)),
457            None => options,
458        };
459
460        let (client, _) = build_in_runtime(
461            async move {
462                options
463                    .connect(self.server)
464                    .await
465                    .map_err(|e| anyhow::anyhow!("Failed to connect to NATS: {e}. Verify NATS server is running and accessible."))
466            },
467            NATS_WORKER_THREADS,
468        )
469        .await?;
470
471        let js_ctx = jetstream::new(client.clone());
472
473        Ok(Client { client, js_ctx })
474    }
475}
476
477impl Default for ClientOptions {
478    fn default() -> Self {
479        ClientOptions {
480            server: default_server(),
481            auth: NatsAuth::default(),
482            tls_ca_cert_path: default_nats_tls_ca_cert_path(),
483            tls_client_cert_path: default_nats_tls_client_cert_path(),
484            tls_client_key_path: default_nats_tls_client_key_path(),
485            tls_insecure: default_nats_tls_insecure(),
486        }
487    }
488}
489
490#[derive(Clone, Eq, PartialEq)]
491pub enum NatsAuth {
492    UserPass(String, String),
493    Token(String),
494    NKey(String),
495    CredentialsFile(PathBuf),
496}
497
498impl std::fmt::Debug for NatsAuth {
499    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
500        match self {
501            NatsAuth::UserPass(user, _pass) => {
502                write!(f, "UserPass({}, <redacted>)", user)
503            }
504            NatsAuth::Token(_token) => write!(f, "Token(<redacted>)"),
505            NatsAuth::NKey(_nkey) => write!(f, "NKey(<redacted>)"),
506            NatsAuth::CredentialsFile(path) => write!(f, "CredentialsFile({:?})", path),
507        }
508    }
509}
510
511impl Default for NatsAuth {
512    fn default() -> Self {
513        if let (Ok(username), Ok(password)) = (
514            std::env::var(env_nats::auth::NATS_AUTH_USERNAME),
515            std::env::var(env_nats::auth::NATS_AUTH_PASSWORD),
516        ) {
517            return NatsAuth::UserPass(username, password);
518        }
519
520        if let Ok(token) = std::env::var(env_nats::auth::NATS_AUTH_TOKEN) {
521            return NatsAuth::Token(token);
522        }
523
524        if let Ok(nkey) = std::env::var(env_nats::auth::NATS_AUTH_NKEY) {
525            return NatsAuth::NKey(nkey);
526        }
527
528        if let Ok(path) = std::env::var(env_nats::auth::NATS_AUTH_CREDENTIALS_FILE) {
529            return NatsAuth::CredentialsFile(PathBuf::from(path));
530        }
531
532        NatsAuth::UserPass("user".to_string(), "user".to_string())
533    }
534}
535
536/// Extract NATS bucket and key from a nats URL of the form:
537/// `nats://host[:port]/bucket/key`
538pub fn url_to_bucket_and_key(url: &Url) -> anyhow::Result<(String, String)> {
539    let Some(mut path_segments) = url.path_segments() else {
540        anyhow::bail!("No path in NATS URL: {url}");
541    };
542    let Some(bucket) = path_segments.next() else {
543        anyhow::bail!("No bucket in NATS URL: {url}");
544    };
545    let Some(key) = path_segments.next() else {
546        anyhow::bail!("No key in NATS URL: {url}");
547    };
548    Ok((bucket.to_string(), key.to_string()))
549}
550
551/// A queue implementation using NATS JetStream
552pub struct NatsQueue {
553    /// The name of the stream to use for the queue
554    stream_name: String,
555    /// The NATS server URL
556    nats_server: String,
557    /// Timeout for dequeue operations in seconds
558    dequeue_timeout: time::Duration,
559    /// The NATS client
560    client: Option<Client>,
561    /// The subject pattern used for this queue
562    subject: String,
563    /// The subscriber for pull-based consumption
564    subscriber: Option<jetstream::consumer::PullConsumer>,
565    /// Optional consumer name for broadcast pattern (if None, uses "worker-group")
566    consumer_name: Option<String>,
567    /// Message stream for efficient message consumption
568    message_stream: Option<jetstream::consumer::pull::Stream>,
569}
570
571impl NatsQueue {
572    /// Create a new NatsQueue with the default "worker-group" consumer
573    pub fn new(stream_name: String, nats_server: String, dequeue_timeout: time::Duration) -> Self {
574        // Sanitize stream name to remove path separators (like in Python version)
575        // rupei: are we sure NATs stream name accepts '_'?
576        let sanitized_stream_name = Slug::slugify(&stream_name).to_string();
577        let subject = format!("{sanitized_stream_name}.*");
578
579        Self {
580            stream_name: sanitized_stream_name,
581            nats_server,
582            dequeue_timeout,
583            client: None,
584            subject,
585            subscriber: None,
586            consumer_name: Some("worker-group".to_string()),
587            message_stream: None,
588        }
589    }
590
591    /// Create a new NatsQueue without a consumer (publisher-only mode)
592    pub fn new_without_consumer(
593        stream_name: String,
594        nats_server: String,
595        dequeue_timeout: time::Duration,
596    ) -> Self {
597        let sanitized_stream_name = Slug::slugify(&stream_name).to_string();
598        let subject = format!("{sanitized_stream_name}.*");
599
600        Self {
601            stream_name: sanitized_stream_name,
602            nats_server,
603            dequeue_timeout,
604            client: None,
605            subject,
606            subscriber: None,
607            consumer_name: None,
608            message_stream: None,
609        }
610    }
611
612    /// Create a new NatsQueue with a specific consumer name for broadcast pattern
613    /// Each consumer with a unique name will receive all messages independently
614    pub fn new_with_consumer(
615        stream_name: String,
616        nats_server: String,
617        dequeue_timeout: time::Duration,
618        consumer_name: String,
619    ) -> Self {
620        let sanitized_stream_name = Slug::slugify(&stream_name).to_string();
621        let subject = format!("{sanitized_stream_name}.*");
622
623        Self {
624            stream_name: sanitized_stream_name,
625            nats_server,
626            dequeue_timeout,
627            client: None,
628            subject,
629            subscriber: None,
630            consumer_name: Some(consumer_name),
631            message_stream: None,
632        }
633    }
634
635    /// Connect to the NATS server and set up the stream and consumer
636    pub async fn connect(&mut self) -> Result<()> {
637        self.connect_with_reset(false).await
638    }
639
640    /// Connect to the NATS server and set up the stream and consumer, optionally resetting the stream
641    pub async fn connect_with_reset(&mut self, reset_stream: bool) -> Result<()> {
642        if self.client.is_none() {
643            // Create a new client
644            let client_options = Client::builder().server(self.nats_server.clone()).build()?;
645
646            let client = client_options.connect().await?;
647
648            // messages older than a hour in the stream will be automatically purged
649            let max_age = std::env::var(env_nats::stream::DYN_NATS_STREAM_MAX_AGE)
650                .ok()
651                .and_then(|s| s.parse::<u64>().ok())
652                .map(time::Duration::from_secs)
653                .unwrap_or_else(|| time::Duration::from_secs(60 * 60));
654
655            let stream_config = jetstream::stream::Config {
656                name: self.stream_name.clone(),
657                subjects: vec![self.subject.clone()],
658                max_age,
659                ..Default::default()
660            };
661
662            // Get or create the stream
663            let stream = client
664                .jetstream()
665                .get_or_create_stream(stream_config)
666                .await?;
667
668            log::debug!("Stream {} is ready", self.stream_name);
669
670            // If reset_stream is true, purge all messages from the stream
671            if reset_stream {
672                match stream.purge().await {
673                    Ok(purge_info) => {
674                        log::info!(
675                            "Successfully purged {} messages from NATS stream {}",
676                            purge_info.purged,
677                            self.stream_name
678                        );
679                    }
680                    Err(e) => {
681                        log::warn!("Failed to purge NATS stream '{}': {e}", self.stream_name);
682                    }
683                }
684            }
685
686            // Create persistent subscriber only if consumer_name is set
687            if let Some(ref consumer_name) = self.consumer_name {
688                let consumer_config = jetstream::consumer::pull::Config {
689                    durable_name: Some(consumer_name.clone()),
690                    inactive_threshold: std::time::Duration::from_secs(300), // 5 minutes
691                    ..Default::default()
692                };
693
694                let subscriber = stream.create_consumer(consumer_config).await?;
695
696                // Create the message stream for efficient consumption
697                let message_stream = subscriber.messages().await?;
698
699                self.subscriber = Some(subscriber);
700                self.message_stream = Some(message_stream);
701            }
702
703            self.client = Some(client);
704        }
705
706        Ok(())
707    }
708
709    /// Ensure we have an active connection
710    pub async fn ensure_connection(&mut self) -> Result<()> {
711        if self.client.is_none() {
712            self.connect().await?;
713        }
714        Ok(())
715    }
716
717    /// Close the connection when done
718    pub async fn close(&mut self) -> Result<()> {
719        self.message_stream = None;
720        self.subscriber = None;
721        self.client = None;
722        Ok(())
723    }
724
725    /// Shutdown the consumer by deleting it from the stream and closing the connection
726    /// This permanently removes the consumer from the server
727    ///
728    /// If `consumer_name` is provided, that specific consumer will be deleted instead of the
729    /// current consumer. This allows deletion of other consumers on the same stream.
730    pub async fn shutdown(&mut self, consumer_name: Option<String>) -> Result<()> {
731        // Determine which consumer to delete
732        let target_consumer = consumer_name.as_ref().or(self.consumer_name.as_ref());
733
734        // Warn if deleting our own consumer via explicit parameter
735        if let Some(ref passed_name) = consumer_name
736            && self.consumer_name.as_ref() == Some(passed_name)
737        {
738            log::warn!(
739                "Deleting our own consumer '{}' via explicit consumer_name parameter. \
740                Consider calling shutdown without arguments instead.",
741                passed_name
742            );
743        }
744
745        if let (Some(client), Some(consumer_to_delete)) = (&self.client, target_consumer) {
746            // Get the stream and delete the consumer
747            let stream = client.jetstream().get_stream(&self.stream_name).await?;
748            stream
749                .delete_consumer(consumer_to_delete)
750                .await
751                .map_err(|e| {
752                    anyhow::anyhow!("Failed to delete consumer {}: {}", consumer_to_delete, e)
753                })?;
754            log::debug!(
755                "Deleted consumer {} from stream {}",
756                consumer_to_delete,
757                self.stream_name
758            );
759        } else {
760            log::debug!(
761                "Cannot shutdown consumer: client or target consumer is None (client: {:?}, target_consumer: {:?})",
762                self.client.is_some(),
763                target_consumer.is_some()
764            );
765        }
766
767        // Only close the connection if we deleted our own consumer
768        if consumer_name.is_none() {
769            self.close().await
770        } else {
771            Ok(())
772        }
773    }
774
775    /// Count the number of consumers for the stream
776    pub async fn count_consumers(&mut self) -> Result<usize> {
777        self.ensure_connection().await?;
778
779        if let Some(client) = &self.client {
780            let mut stream = client.jetstream().get_stream(&self.stream_name).await?;
781            let info = stream.info().await?;
782            Ok(info.state.consumer_count)
783        } else {
784            Err(anyhow::anyhow!("Client not connected"))
785        }
786    }
787
788    /// List all consumer names for the stream
789    pub async fn list_consumers(&mut self) -> Result<Vec<String>> {
790        self.ensure_connection().await?;
791
792        if let Some(client) = &self.client {
793            client.list_consumers(&self.stream_name).await
794        } else {
795            Err(anyhow::anyhow!("Client not connected"))
796        }
797    }
798
799    /// Enqueue a task using the provided data
800    pub async fn enqueue_task(&mut self, task_data: Bytes) -> Result<()> {
801        self.ensure_connection().await?;
802
803        if let Some(client) = &self.client {
804            let subject = format!("{}.queue", self.stream_name);
805            client.jetstream().publish(subject, task_data).await?;
806            Ok(())
807        } else {
808            Err(anyhow::anyhow!("Client not connected"))
809        }
810    }
811
812    /// Dequeue and return a task as raw bytes
813    pub async fn dequeue_task(&mut self, timeout: Option<time::Duration>) -> Result<Option<Bytes>> {
814        self.ensure_connection().await?;
815
816        let Some(ref mut stream) = self.message_stream else {
817            return Err(anyhow::anyhow!("Message stream not initialized"));
818        };
819
820        let timeout_duration = timeout.unwrap_or(self.dequeue_timeout);
821
822        // Try to get next message from the stream with timeout
823        let message = tokio::time::timeout(timeout_duration, stream.next()).await;
824
825        match message {
826            Ok(Some(Ok(msg))) => {
827                msg.ack()
828                    .await
829                    .map_err(|e| anyhow::anyhow!("Failed to ack message: {}", e))?;
830                Ok(Some(msg.payload.clone()))
831            }
832
833            Ok(Some(Err(e))) => Err(anyhow::anyhow!("Failed to get message from stream: {}", e)),
834
835            Ok(None) => Err(anyhow::anyhow!("Message stream ended unexpectedly")),
836
837            // Timeout - no messages available
838            Err(_) => Ok(None),
839        }
840    }
841
842    /// Get the number of messages currently in the queue
843    pub async fn get_queue_size(&mut self) -> Result<u64> {
844        self.ensure_connection().await?;
845
846        if let Some(client) = &self.client {
847            // Get consumer info to get pending messages count
848            let stream = client.jetstream().get_stream(&self.stream_name).await?;
849            let consumer_name = self
850                .consumer_name
851                .clone()
852                .unwrap_or_else(|| "worker-group".to_string());
853            let mut consumer: jetstream::consumer::PullConsumer = stream
854                .get_consumer(&consumer_name)
855                .await
856                .map_err(|e| anyhow::anyhow!("Failed to get consumer: {}", e))?;
857            let info = consumer.info().await?;
858
859            Ok(info.num_pending)
860        } else {
861            Err(anyhow::anyhow!("Client not connected"))
862        }
863    }
864
865    /// Get the total number of messages currently in the stream
866    pub async fn get_stream_messages(&mut self) -> Result<u64> {
867        self.ensure_connection().await?;
868
869        if let Some(client) = &self.client {
870            let mut stream = client.jetstream().get_stream(&self.stream_name).await?;
871            let info = stream.info().await?;
872            Ok(info.state.messages)
873        } else {
874            Err(anyhow::anyhow!("Client not connected"))
875        }
876    }
877
878    /// Purge messages from the stream up to (but not including) the specified sequence number
879    /// This permanently removes messages and affects all consumers of the stream
880    pub async fn purge_up_to_sequence(&self, sequence: u64) -> Result<()> {
881        if let Some(client) = &self.client {
882            let stream = client.jetstream().get_stream(&self.stream_name).await?;
883
884            // NOTE: this purge excludes the sequence itself
885            // https://docs.rs/nats/latest/nats/jetstream/struct.PurgeRequest.html
886            stream.purge().sequence(sequence).await.map_err(|e| {
887                anyhow::anyhow!("Failed to purge stream up to sequence {}: {}", sequence, e)
888            })?;
889
890            log::debug!(
891                "Purged stream {} up to sequence {}",
892                self.stream_name,
893                sequence
894            );
895            Ok(())
896        } else {
897            Err(anyhow::anyhow!("Client not connected"))
898        }
899    }
900
901    /// Purge messages from the stream up to the minimum acknowledged sequence across all consumers
902    /// This finds the lowest acknowledged sequence number across all consumers and purges up to that point
903    pub async fn purge_acknowledged(&mut self) -> Result<()> {
904        self.ensure_connection().await?;
905
906        let Some(client) = &self.client else {
907            return Err(anyhow::anyhow!("Client not connected"));
908        };
909
910        let stream = client.jetstream().get_stream(&self.stream_name).await?;
911
912        // Get all consumer names for the stream
913        let consumer_names: Vec<String> = stream
914            .consumer_names()
915            .try_collect()
916            .await
917            .map_err(|e| anyhow::anyhow!("Failed to list consumers: {}", e))?;
918
919        if consumer_names.is_empty() {
920            log::debug!("No consumers found for stream {}", self.stream_name);
921            return Ok(());
922        }
923
924        // Find the minimum acknowledged sequence across all consumers
925        let mut min_ack_sequence = u64::MAX;
926
927        for consumer_name in &consumer_names {
928            let mut consumer: jetstream::consumer::PullConsumer = stream
929                .get_consumer(consumer_name)
930                .await
931                .map_err(|e| anyhow::anyhow!("Failed to get consumer {}: {}", consumer_name, e))?;
932
933            let info = consumer.info().await.map_err(|e| {
934                anyhow::anyhow!("Failed to get consumer info for {}: {}", consumer_name, e)
935            })?;
936
937            // The ack_floor contains the stream sequence of the highest contiguously acknowledged message
938            // If stream_sequence is 0, it means no messages have been acknowledged yet
939            if info.ack_floor.stream_sequence > 0 {
940                min_ack_sequence = min_ack_sequence.min(info.ack_floor.stream_sequence);
941                log::debug!(
942                    "Consumer {} has ack_floor at sequence {}",
943                    consumer_name,
944                    info.ack_floor.stream_sequence
945                );
946            }
947        }
948
949        // Only purge if we found a valid minimum acknowledged sequence
950        if min_ack_sequence < u64::MAX && min_ack_sequence > 0 {
951            // Purge up to (but not including) the minimum acknowledged sequence + 1
952            // We add 1 because we want to include the minimum acknowledged message in the purge
953            let purge_sequence = min_ack_sequence + 1;
954
955            self.purge_up_to_sequence(purge_sequence).await?;
956
957            log::debug!(
958                "Purged stream {} up to acknowledged sequence {} (purged up to sequence {})",
959                self.stream_name,
960                min_ack_sequence,
961                purge_sequence
962            );
963        } else {
964            log::debug!(
965                "No messages to purge for stream {} (min_ack_sequence: {})",
966                self.stream_name,
967                min_ack_sequence
968            );
969        }
970
971        Ok(())
972    }
973}
974
975impl NatsQueue {
976    pub fn event_subject(&self) -> String {
977        self.stream_name.clone()
978    }
979
980    pub async fn publish_event(
981        &self,
982        event_name: impl AsRef<str> + Send + Sync,
983        event: &(impl Serialize + Send + Sync),
984    ) -> Result<()> {
985        let bytes = serde_json::to_vec(event)?;
986        self.publish_event_bytes(event_name, bytes).await
987    }
988
989    pub async fn publish_event_bytes(
990        &self,
991        event_name: impl AsRef<str> + Send + Sync,
992        bytes: Vec<u8>,
993    ) -> Result<()> {
994        let subject = format!("{}.{}", self.event_subject(), event_name.as_ref());
995
996        // Note: enqueue_task requires &mut self, but EventPublisher requires &self
997        // We need to ensure the client is connected and use it directly
998        if let Some(client) = &self.client {
999            client.jetstream().publish(subject, bytes.into()).await?;
1000            Ok(())
1001        } else {
1002            Err(anyhow::anyhow!("Client not connected"))
1003        }
1004    }
1005}
1006
1007/// The NATS subject / inbox to talk to an instance on.
1008/// TODO: Do we need to sanitize the names?
1009pub fn instance_subject(endpoint_id: &EndpointId, instance_id: u64) -> String {
1010    format!(
1011        "{}_{}.{}-{:x}",
1012        endpoint_id.namespace, endpoint_id.component, endpoint_id.name, instance_id,
1013    )
1014}
1015
1016#[cfg(test)]
1017mod tests {
1018
1019    use super::*;
1020    use figment::Jail;
1021    use serde::{Deserialize, Serialize};
1022
1023    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1024    struct TestData {
1025        id: u32,
1026        name: String,
1027        values: Vec<f64>,
1028    }
1029
1030    #[test]
1031    #[allow(clippy::result_large_err)]
1032    fn test_client_options_builder() {
1033        Jail::expect_with(|_jail| {
1034            let opts = ClientOptions::builder().build();
1035            assert!(opts.is_ok());
1036            Ok(())
1037        });
1038
1039        Jail::expect_with(|jail| {
1040            jail.set_env(env_nats::NATS_SERVER, "nats://localhost:5222");
1041            jail.set_env(env_nats::auth::NATS_AUTH_USERNAME, "user");
1042            jail.set_env(env_nats::auth::NATS_AUTH_PASSWORD, "pass");
1043
1044            let opts = ClientOptions::builder().build();
1045            assert!(opts.is_ok());
1046            let opts = opts.unwrap();
1047
1048            assert_eq!(opts.server, "nats://localhost:5222");
1049            assert_eq!(
1050                opts.auth,
1051                NatsAuth::UserPass("user".to_string(), "pass".to_string())
1052            );
1053
1054            Ok(())
1055        });
1056
1057        Jail::expect_with(|jail| {
1058            jail.set_env(env_nats::NATS_SERVER, "nats://localhost:5222");
1059            jail.set_env(env_nats::auth::NATS_AUTH_USERNAME, "user");
1060            jail.set_env(env_nats::auth::NATS_AUTH_PASSWORD, "pass");
1061
1062            let opts = ClientOptions::builder()
1063                .server("nats://localhost:6222")
1064                .auth(NatsAuth::Token("token".to_string()))
1065                .build();
1066            assert!(opts.is_ok());
1067            let opts = opts.unwrap();
1068
1069            assert_eq!(opts.server, "nats://localhost:6222");
1070            assert_eq!(opts.auth, NatsAuth::Token("token".to_string()));
1071
1072            Ok(())
1073        });
1074    }
1075
1076    #[test]
1077    #[allow(clippy::result_large_err)]
1078    fn test_client_options_tls_url_validation() {
1079        // tls:// is accepted by the validator
1080        assert!(validate_nats_server("tls://nats:4222").is_ok());
1081        assert!(validate_nats_server("nats://nats:4222").is_ok());
1082        assert!(validate_nats_server("tcp://nats:4222").is_err());
1083        assert!(validate_nats_server("http://nats:4222").is_err());
1084
1085        // tls:// URL is preserved in options
1086        Jail::expect_with(|jail| {
1087            jail.set_env(env_nats::NATS_SERVER, "tls://nats:4222");
1088            let opts = ClientOptions::builder().build().unwrap();
1089            assert_eq!(opts.server, "tls://nats:4222");
1090            Ok(())
1091        });
1092    }
1093
1094    #[test]
1095    #[allow(clippy::result_large_err)]
1096    fn test_client_options_tls_ca_from_env() {
1097        Jail::expect_with(|jail| {
1098            jail.set_env(env_nats::NATS_SERVER, "tls://nats:4222");
1099            jail.set_env(env_nats::tls::NATS_TLS_CA_CERT_PATH, "/etc/certs/ca.pem");
1100            let opts = ClientOptions::builder().build().unwrap();
1101            assert_eq!(
1102                opts.tls_ca_cert_path,
1103                Some(PathBuf::from("/etc/certs/ca.pem"))
1104            );
1105            assert!(!opts.tls_insecure);
1106            Ok(())
1107        });
1108    }
1109
1110    #[test]
1111    #[allow(clippy::result_large_err)]
1112    fn test_client_options_tls_insecure_from_env() {
1113        Jail::expect_with(|jail| {
1114            jail.set_env(env_nats::NATS_SERVER, "tls://nats:4222");
1115            jail.set_env(env_nats::tls::NATS_TLS_INSECURE, "1");
1116            let opts = ClientOptions::builder().build().unwrap();
1117            assert!(opts.tls_insecure);
1118            assert!(opts.tls_ca_cert_path.is_none());
1119            Ok(())
1120        });
1121    }
1122
1123    #[test]
1124    #[allow(clippy::result_large_err)]
1125    fn test_client_options_no_tls_by_default() {
1126        Jail::expect_with(|_jail| {
1127            let opts = ClientOptions::builder().build().unwrap();
1128            assert!(opts.tls_ca_cert_path.is_none());
1129            assert!(!opts.tls_insecure);
1130            Ok(())
1131        });
1132    }
1133
1134    #[tokio::test]
1135    async fn test_connect_rejects_custom_tls_with_nats_url() {
1136        // Client is not Debug, so match rather than use unwrap_err/expect_err.
1137        fn assert_scheme_error(result: Result<Client>, case: &str) {
1138            match result {
1139                Ok(_) => panic!("{case}: expected an error, got a connection"),
1140                Err(e) => assert!(
1141                    e.to_string().contains("tls://"),
1142                    "{case}: unexpected error: {e}"
1143                ),
1144            }
1145        }
1146
1147        // CA cert set but server is nats:// (not tls://) → rejected before connecting.
1148        let opts = ClientOptions::builder()
1149            .server("nats://localhost:4222".to_string())
1150            .tls_ca_cert_path(Some(PathBuf::from("/etc/certs/ca.pem")))
1151            .build()
1152            .unwrap();
1153        assert_scheme_error(opts.connect().await, "nats:// + CA cert");
1154
1155        // Insecure mode set but server is nats:// → also rejected.
1156        let opts = ClientOptions::builder()
1157            .server("nats://localhost:4222".to_string())
1158            .tls_insecure(true)
1159            .build()
1160            .unwrap();
1161        assert_scheme_error(opts.connect().await, "nats:// + insecure");
1162    }
1163
1164    #[test]
1165    #[allow(clippy::result_large_err)]
1166    fn test_nats_mtls_client_cert_from_env() {
1167        Jail::expect_with(|jail| {
1168            jail.set_env(env_nats::NATS_SERVER, "tls://localhost:4222");
1169            jail.set_env(env_nats::tls::NATS_TLS_CA_CERT_PATH, "/etc/certs/ca.pem");
1170            jail.set_env(
1171                env_nats::tls::NATS_TLS_CLIENT_CERT_PATH,
1172                "/etc/certs/client.pem",
1173            );
1174            jail.set_env(
1175                env_nats::tls::NATS_TLS_CLIENT_KEY_PATH,
1176                "/etc/certs/client-key.pem",
1177            );
1178            let opts = ClientOptions::builder().build().unwrap();
1179            assert_eq!(
1180                opts.tls_client_cert_path,
1181                Some(PathBuf::from("/etc/certs/client.pem"))
1182            );
1183            assert_eq!(
1184                opts.tls_client_key_path,
1185                Some(PathBuf::from("/etc/certs/client-key.pem"))
1186            );
1187            Ok(())
1188        });
1189    }
1190
1191    #[tokio::test]
1192    async fn test_nats_mtls_client_cert_requires_ca() {
1193        // Client cert/key set, tls:// URL, but no CA and not insecure → rejected.
1194        let opts = ClientOptions::builder()
1195            .server("tls://localhost:4222".to_string())
1196            .tls_client_cert_path(Some(PathBuf::from("/tmp/client.pem")))
1197            .tls_client_key_path(Some(PathBuf::from("/tmp/client-key.pem")))
1198            .build()
1199            .unwrap();
1200        match opts.connect().await {
1201            Err(e) => assert!(
1202                e.to_string().contains("requires"),
1203                "expected CA-requirement error, got: {e}"
1204            ),
1205            Ok(_) => panic!("expected error when client cert set without CA"),
1206        }
1207    }
1208
1209    #[tokio::test]
1210    async fn test_nats_mtls_partial_client_identity_errors() {
1211        // Client cert without key → rejected before connecting.
1212        let opts = ClientOptions::builder()
1213            .server("tls://localhost:4222".to_string())
1214            .tls_ca_cert_path(Some(PathBuf::from("/tmp/ca.pem")))
1215            .tls_client_cert_path(Some(PathBuf::from("/tmp/client.pem")))
1216            .build()
1217            .unwrap();
1218        match opts.connect().await {
1219            Err(e) => assert!(
1220                e.to_string().contains("must be set together"),
1221                "expected both-or-neither error, got: {e}"
1222            ),
1223            Ok(_) => panic!("expected error when client cert set without key"),
1224        }
1225    }
1226
1227    // Integration test for object store data operations using bincode
1228    #[tokio::test]
1229    #[ignore] // Requires NATS server to be running
1230    async fn test_object_store_data_operations() {
1231        // Create test data
1232        let test_data = TestData {
1233            id: 42,
1234            name: "test_item".to_string(),
1235            values: vec![1.0, 2.5, 3.7, 4.2],
1236        };
1237
1238        // Set up client
1239        let client_options = ClientOptions::builder()
1240            .server("nats://localhost:4222")
1241            .build()
1242            .expect("Failed to build client options");
1243
1244        let client = client_options
1245            .connect()
1246            .await
1247            .expect("Failed to connect to NATS");
1248
1249        // Test URL (using .bin extension to indicate binary format)
1250        let url =
1251            Url::parse("nats://localhost/test-bucket/test-data.bin").expect("Failed to parse URL");
1252
1253        // Upload the data
1254        client
1255            .object_store_upload_data(&test_data, &url)
1256            .await
1257            .expect("Failed to upload data");
1258
1259        // Download the data
1260        let downloaded_data: TestData = client
1261            .object_store_download_data(&url)
1262            .await
1263            .expect("Failed to download data");
1264
1265        // Verify the data matches
1266        assert_eq!(test_data, downloaded_data);
1267
1268        // Clean up
1269        client
1270            .object_store_delete_bucket("test-bucket")
1271            .await
1272            .expect("Failed to delete bucket");
1273    }
1274
1275    // Integration test for broadcast pattern with purging
1276    #[tokio::test]
1277    #[ignore]
1278    async fn test_nats_queue_broadcast_with_purge() {
1279        use uuid::Uuid;
1280
1281        // Create unique stream name for this test
1282        let stream_name = format!("test-broadcast-{}", Uuid::new_v4());
1283        let nats_server = "nats://localhost:4222".to_string();
1284        let timeout = time::Duration::from_secs(0);
1285
1286        // Connect to NATS client first to delete stream if it exists
1287        let client_options = Client::builder()
1288            .server(nats_server.clone())
1289            .build()
1290            .expect("Failed to build client options");
1291
1292        let client = client_options
1293            .connect()
1294            .await
1295            .expect("Failed to connect to NATS");
1296
1297        // Delete the stream if it exists (to ensure clean start)
1298        let _ = client.jetstream().delete_stream(&stream_name).await;
1299
1300        // Create two consumers with different names for the same stream
1301        let consumer1_name = format!("consumer-{}", Uuid::new_v4());
1302        let consumer2_name = format!("consumer-{}", Uuid::new_v4());
1303
1304        let mut queue1 = NatsQueue::new_with_consumer(
1305            stream_name.clone(),
1306            nats_server.clone(),
1307            timeout,
1308            consumer1_name,
1309        );
1310
1311        // Connect queue1 first (it will create the stream)
1312        queue1.connect().await.expect("Failed to connect queue1");
1313
1314        // Send 4 messages using the EventPublisher trait
1315        let message_strings = [
1316            "message1".to_string(),
1317            "message2".to_string(),
1318            "message3".to_string(),
1319            "message4".to_string(),
1320        ];
1321
1322        // Publish messages using NatsQueue
1323        for (idx, msg) in message_strings.iter().enumerate() {
1324            queue1
1325                .publish_event("queue", msg)
1326                .await
1327                .unwrap_or_else(|_| panic!("Failed to publish message {}", idx + 1));
1328        }
1329
1330        // Convert messages to JSON-serialized Bytes for comparison
1331        let messages: Vec<Bytes> = message_strings
1332            .iter()
1333            .map(|s| Bytes::from(serde_json::to_vec(s).unwrap()))
1334            .collect();
1335
1336        // Give JetStream a moment to persist the messages
1337        tokio::time::sleep(time::Duration::from_millis(100)).await;
1338
1339        // Now create and connect queue2 and queue3 AFTER messages are published (to test persistence)
1340        let mut queue2 = NatsQueue::new_with_consumer(
1341            stream_name.clone(),
1342            nats_server.clone(),
1343            timeout,
1344            consumer2_name,
1345        );
1346
1347        // Create a third queue without consumer (publisher-only)
1348        let mut queue3 =
1349            NatsQueue::new_without_consumer(stream_name.clone(), nats_server.clone(), timeout);
1350
1351        // Connect queue2 and queue3 after messages are already published
1352        queue2.connect().await.expect("Failed to connect queue2");
1353        queue3.connect().await.expect("Failed to connect queue3");
1354
1355        // Purge the first two messages (sequence 1 and 2)
1356        // Note: JetStream sequences start at 1, and purge is exclusive of the sequence number
1357        queue1
1358            .purge_up_to_sequence(3)
1359            .await
1360            .expect("Failed to purge messages");
1361
1362        // Give JetStream a moment to process the purge
1363        tokio::time::sleep(time::Duration::from_millis(100)).await;
1364
1365        // Consumer 1 dequeues one message (message3)
1366        let msg3_consumer1 = queue1
1367            .dequeue_task(Some(time::Duration::from_millis(500)))
1368            .await
1369            .expect("Failed to dequeue from queue1");
1370        assert_eq!(
1371            msg3_consumer1,
1372            Some(messages[2].clone()),
1373            "Consumer 1 should get message3"
1374        );
1375
1376        // Give JetStream a moment to process acknowledgments
1377        tokio::time::sleep(time::Duration::from_millis(100)).await;
1378
1379        // Now run purge_acknowledged
1380        // At this point:
1381        // - Consumer 1 has ack'd message 3 (ack_floor = 3)
1382        // - Consumer 2 hasn't consumed anything yet (ack_floor = 0)
1383        // - Min ack_floor = 0, so nothing will be purged
1384        queue1
1385            .purge_acknowledged()
1386            .await
1387            .expect("Failed to purge acknowledged messages");
1388
1389        // Give JetStream a moment to process the purge
1390        tokio::time::sleep(time::Duration::from_millis(100)).await;
1391
1392        // Now collect remaining messages from both consumers
1393        let mut consumer1_remaining = Vec::new();
1394        let mut consumer2_remaining = Vec::new();
1395
1396        // Collect remaining messages from consumer 1
1397        while let Some(msg) = queue1
1398            .dequeue_task(None)
1399            .await
1400            .expect("Failed to dequeue from queue1")
1401        {
1402            consumer1_remaining.push(msg);
1403        }
1404
1405        // Collect remaining messages from consumer 2
1406        while let Some(msg) = queue2
1407            .dequeue_task(None)
1408            .await
1409            .expect("Failed to dequeue from queue2")
1410        {
1411            consumer2_remaining.push(msg);
1412        }
1413
1414        // Verify consumer 1 gets 1 remaining message (message4)
1415        assert_eq!(
1416            consumer1_remaining.len(),
1417            1,
1418            "Consumer 1 should have 1 remaining message"
1419        );
1420        assert_eq!(
1421            consumer1_remaining[0], messages[3],
1422            "Consumer 1 should get message4"
1423        );
1424
1425        // Verify consumer 2 gets 2 messages (message3 and message4)
1426        assert_eq!(
1427            consumer2_remaining.len(),
1428            2,
1429            "Consumer 2 should have 2 messages"
1430        );
1431        assert_eq!(
1432            consumer2_remaining[0], messages[2],
1433            "Consumer 2 should get message3"
1434        );
1435        assert_eq!(
1436            consumer2_remaining[1], messages[3],
1437            "Consumer 2 should get message4"
1438        );
1439
1440        // Test consumer count and shutdown behavior
1441        // First verify via consumer 1 that there are two consumers
1442        let consumer_count = queue1
1443            .count_consumers()
1444            .await
1445            .expect("Failed to count consumers");
1446        assert_eq!(consumer_count, 2, "Should have 2 consumers initially");
1447
1448        // Close consumer 1 and verify via consumer 2 that there are still two consumers
1449        queue1.close().await.expect("Failed to close queue1");
1450
1451        let consumer_count = queue2
1452            .count_consumers()
1453            .await
1454            .expect("Failed to count consumers");
1455        assert_eq!(
1456            consumer_count, 2,
1457            "Should still have 2 consumers after closing queue1"
1458        );
1459
1460        // Reconnect queue1 to be able to shutdown
1461        queue1.connect().await.expect("Failed to reconnect queue1");
1462
1463        // Shutdown consumer 1 and verify via consumer 2 that there is only one consumer left
1464        queue1
1465            .shutdown(None)
1466            .await
1467            .expect("Failed to shutdown queue1");
1468
1469        let consumer_count = queue2
1470            .count_consumers()
1471            .await
1472            .expect("Failed to count consumers");
1473        assert_eq!(
1474            consumer_count, 1,
1475            "Should have only 1 consumer after shutting down queue1"
1476        );
1477
1478        // Clean up by deleting the stream
1479        client
1480            .jetstream()
1481            .delete_stream(&stream_name)
1482            .await
1483            .expect("Failed to delete test stream");
1484    }
1485}