1use 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 pub fn builder() -> ClientOptionsBuilder {
69 ClientOptionsBuilder::default()
70 }
71
72 pub fn client(&self) -> &client::Client {
74 &self.client
75 }
76
77 pub fn jetstream(&self) -> &jetstream::Context {
79 &self.js_ctx
80 }
81
82 pub fn addr(&self) -> String {
84 let info = self.client.server_info();
85 format!("{}:{}", info.host, info.port)
86 }
87
88 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 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 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 self.client
127 .publish_with_reply(subject, reply_subject, "".into())
128 .await?;
129
130 Ok(subscription)
131 }
132
133 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 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 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 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 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 pub async fn object_store_upload_data<T>(&self, data: &T, nats_url: &Url) -> anyhow::Result<()>
230 where
231 T: Serialize,
232 {
233 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 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 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 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 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#[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 #[builder(default = "default_nats_tls_ca_cert_path()")]
299 tls_ca_cert_path: Option<PathBuf>,
300
301 #[builder(default = "default_nats_tls_client_cert_path()")]
304 tls_client_cert_path: Option<PathBuf>,
305
306 #[builder(default = "default_nats_tls_client_key_path()")]
308 tls_client_key_path: Option<PathBuf>,
309
310 #[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
355const NATS_WORKER_THREADS: usize = 4;
357
358impl ClientOptions {
359 pub fn builder() -> ClientOptionsBuilder {
361 ClientOptionsBuilder::default()
362 }
363
364 pub async fn connect(self) -> Result<Client> {
366 self.validate()?;
367
368 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 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 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 if custom_tls || tls_url {
428 let _ = rustls::crypto::ring::default_provider().install_default();
429 }
430
431 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 options = options.require_tls(true);
447 }
448
449 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
536pub 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
551pub struct NatsQueue {
553 stream_name: String,
555 nats_server: String,
557 dequeue_timeout: time::Duration,
559 client: Option<Client>,
561 subject: String,
563 subscriber: Option<jetstream::consumer::PullConsumer>,
565 consumer_name: Option<String>,
567 message_stream: Option<jetstream::consumer::pull::Stream>,
569}
570
571impl NatsQueue {
572 pub fn new(stream_name: String, nats_server: String, dequeue_timeout: time::Duration) -> Self {
574 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 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 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 pub async fn connect(&mut self) -> Result<()> {
637 self.connect_with_reset(false).await
638 }
639
640 pub async fn connect_with_reset(&mut self, reset_stream: bool) -> Result<()> {
642 if self.client.is_none() {
643 let client_options = Client::builder().server(self.nats_server.clone()).build()?;
645
646 let client = client_options.connect().await?;
647
648 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 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 {
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 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), ..Default::default()
692 };
693
694 let subscriber = stream.create_consumer(consumer_config).await?;
695
696 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 pub async fn ensure_connection(&mut self) -> Result<()> {
711 if self.client.is_none() {
712 self.connect().await?;
713 }
714 Ok(())
715 }
716
717 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 pub async fn shutdown(&mut self, consumer_name: Option<String>) -> Result<()> {
731 let target_consumer = consumer_name.as_ref().or(self.consumer_name.as_ref());
733
734 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 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 if consumer_name.is_none() {
769 self.close().await
770 } else {
771 Ok(())
772 }
773 }
774
775 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 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 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 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 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 Err(_) => Ok(None),
839 }
840 }
841
842 pub async fn get_queue_size(&mut self) -> Result<u64> {
844 self.ensure_connection().await?;
845
846 if let Some(client) = &self.client {
847 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 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 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 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 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 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 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 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 if min_ack_sequence < u64::MAX && min_ack_sequence > 0 {
951 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 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
1007pub 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 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 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 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 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 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 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 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 #[tokio::test]
1229 #[ignore] async fn test_object_store_data_operations() {
1231 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 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 let url =
1251 Url::parse("nats://localhost/test-bucket/test-data.bin").expect("Failed to parse URL");
1252
1253 client
1255 .object_store_upload_data(&test_data, &url)
1256 .await
1257 .expect("Failed to upload data");
1258
1259 let downloaded_data: TestData = client
1261 .object_store_download_data(&url)
1262 .await
1263 .expect("Failed to download data");
1264
1265 assert_eq!(test_data, downloaded_data);
1267
1268 client
1270 .object_store_delete_bucket("test-bucket")
1271 .await
1272 .expect("Failed to delete bucket");
1273 }
1274
1275 #[tokio::test]
1277 #[ignore]
1278 async fn test_nats_queue_broadcast_with_purge() {
1279 use uuid::Uuid;
1280
1281 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 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 let _ = client.jetstream().delete_stream(&stream_name).await;
1299
1300 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 queue1.connect().await.expect("Failed to connect queue1");
1313
1314 let message_strings = [
1316 "message1".to_string(),
1317 "message2".to_string(),
1318 "message3".to_string(),
1319 "message4".to_string(),
1320 ];
1321
1322 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 let messages: Vec<Bytes> = message_strings
1332 .iter()
1333 .map(|s| Bytes::from(serde_json::to_vec(s).unwrap()))
1334 .collect();
1335
1336 tokio::time::sleep(time::Duration::from_millis(100)).await;
1338
1339 let mut queue2 = NatsQueue::new_with_consumer(
1341 stream_name.clone(),
1342 nats_server.clone(),
1343 timeout,
1344 consumer2_name,
1345 );
1346
1347 let mut queue3 =
1349 NatsQueue::new_without_consumer(stream_name.clone(), nats_server.clone(), timeout);
1350
1351 queue2.connect().await.expect("Failed to connect queue2");
1353 queue3.connect().await.expect("Failed to connect queue3");
1354
1355 queue1
1358 .purge_up_to_sequence(3)
1359 .await
1360 .expect("Failed to purge messages");
1361
1362 tokio::time::sleep(time::Duration::from_millis(100)).await;
1364
1365 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 tokio::time::sleep(time::Duration::from_millis(100)).await;
1378
1379 queue1
1385 .purge_acknowledged()
1386 .await
1387 .expect("Failed to purge acknowledged messages");
1388
1389 tokio::time::sleep(time::Duration::from_millis(100)).await;
1391
1392 let mut consumer1_remaining = Vec::new();
1394 let mut consumer2_remaining = Vec::new();
1395
1396 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 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 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 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 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 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 queue1.connect().await.expect("Failed to reconnect queue1");
1462
1463 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 client
1480 .jetstream()
1481 .delete_stream(&stream_name)
1482 .await
1483 .expect("Failed to delete test stream");
1484 }
1485}