qm_nats/
lib.rs

1#![deny(missing_docs)]
2
3//! NATS JetStream integration for building distributed microservices.
4//!
5//! This crate provides utilities for connecting to NATS with JetStream support,
6//! enabling event-driven architectures with distributed locking and sequencing capabilities.
7//!
8//! ## Features
9//!
10//! - **Event Publishing**: Stream events to NATS JetStream with structured subject paths
11//! - **Distributed Locks**: Acquire and manage distributed locks across services
12//! - **Sequence Generation**: Generate unique, monotonically increasing sequences
13//! - **System Consumers**: Create durable pull consumers for event processing
14//! - **Configuration**: Environment-based configuration with sensible defaults
15//!
16//! ## Quick Start
17//!
18//! ```ignore
19//! use qm_nats::{Config, Nats};
20//!
21//! #[tokio::main]
22//! async fn main() -> anyhow::Result<()> {
23//!     let config = Config::new()?;
24//!     let nats = Nats::new(config).await?;
25//!
26//!     // Create a publisher
27//!     let publisher = nats.publisher().await?;
28//!     publisher.publish("subject.here", &"hello").await?;
29//!
30//!     // Or use distributed locks
31//!     let locks = nats.distributed_locks().await?;
32//!     let lock_manager = locks.sys_locks().await?;
33//!     let result = lock_manager.run_locked("my-resource", async {
34//!         // Critical section
35//!         Ok::<_, std::convert::Infalloid>(42)
36//!     }).await?;
37//!
38//!     Ok(())
39//! }
40//! ```
41//!
42//! ## Environment Variables
43//!
44//! | Variable | Description | Default |
45//! |----------|-------------|---------|
46//! | `NATS_HOST` | NATS server host | `127.0.0.1` |
47//! | `NATS_PORT` | NATS server port | `4222` |
48//! | `NATS_APP_NAME` | Application name | `edd-service-rs` |
49//! | `NATS_SYS_LOCKS` | Key-value bucket for locks | `SYS_LOCKS` |
50//! | `NATS_EVENTS_STREAM_NAME` | JetStream stream for events | `EVENTS` |
51//! | `NATS_EVENTS_STREAM_SUBJECT` | Subject pattern for events | `ev.>` |
52
53use std::{
54    error::Error,
55    sync::{
56        atomic::{AtomicU64, Ordering},
57        Arc,
58    },
59};
60
61use async_nats::{
62    jetstream::{
63        self,
64        consumer::PullConsumer,
65        context::{
66            CreateKeyValueErrorKind, CreateStreamError, GetStreamByNameErrorKind,
67            GetStreamErrorKind, KeyValueErrorKind,
68        },
69        kv::{self, Operation, Store},
70        stream::ConsumerError,
71        Context,
72    },
73    subject::ToSubject,
74    Client, ConnectError, ConnectErrorKind,
75};
76use futures::StreamExt;
77use tokio::task::JoinHandle;
78
79pub use async_nats;
80
81/// Subject module for event subject path generation.
82pub mod subject;
83
84/// Configuration for NATS JetStream connection.
85///
86/// Loads configuration from environment variables with sensible defaults.
87/// See module-level documentation for available environment variables.
88#[derive(Clone, serde::Deserialize)]
89pub struct Config {
90    app_name: Option<String>,
91    host: Option<String>,
92    port: Option<u16>,
93    #[serde(skip)]
94    address: Option<String>,
95    sys_locks: Option<String>,
96    events_stream_name: Option<String>,
97    events_stream_subject: Option<String>,
98}
99
100impl Config {
101    /// Creates a new Config from environment variables with default NATS_ prefix.
102    pub fn new() -> envy::Result<Self> {
103        ConfigBuilder::default().build()
104    }
105
106    /// Creates a new ConfigBuilder for custom configuration.
107    pub fn builder<'a>() -> ConfigBuilder<'a> {
108        ConfigBuilder::default()
109    }
110
111    /// Returns the NATS server address.
112    pub fn address(&self) -> &str {
113        self.address.as_deref().unwrap()
114    }
115
116    /// Returns the NATS server port.
117    pub fn port(&self) -> u16 {
118        self.port.unwrap_or(3000)
119    }
120
121    /// Returns the key-value bucket name for system locks.
122    pub fn sys_locks(&self) -> &str {
123        self.sys_locks.as_deref().unwrap_or("SYS_LOCKS")
124    }
125
126    /// Returns the JetStream stream name for events.
127    pub fn events_stream_name(&self) -> &str {
128        self.events_stream_name.as_deref().unwrap_or("EVENTS")
129    }
130
131    /// Returns the subject pattern for events stream.
132    pub fn events_stream_subject(&self) -> &str {
133        self.events_stream_subject.as_deref().unwrap_or("ev.>")
134    }
135}
136
137/// Builder for creating Config with custom settings.
138#[derive(Default)]
139pub struct ConfigBuilder<'a> {
140    prefix: Option<&'a str>,
141}
142
143impl<'a> ConfigBuilder<'a> {
144    /// Sets a custom environment variable prefix.
145    pub fn with_prefix(mut self, prefix: &'a str) -> Self {
146        self.prefix = Some(prefix);
147        self
148    }
149
150    /// Builds the Config from environment variables.
151    pub fn build(self) -> envy::Result<Config> {
152        let prefix = self.prefix.unwrap_or("NATS_");
153        let mut cfg: Config = envy::prefixed(prefix).from_env()?;
154        if cfg.app_name.is_none() {
155            cfg.app_name = Some("edd-service-rs".into());
156        }
157        let host = cfg.host.as_deref().unwrap_or("127.0.0.1");
158        let port = cfg.port.unwrap_or(4222);
159        cfg.address = Some(format!("nats://{}:{}", host, port));
160        Ok(cfg)
161    }
162}
163
164/// Internal state for the Nats client.
165pub struct Inner {
166    client: Client,
167    config: Config,
168}
169
170/// NATS JetStream client wrapper.
171///
172/// Provides high-level access to NATS JetStream features including
173/// event publishing, distributed locking, and sequence management.
174#[derive(Clone)]
175pub struct Nats {
176    inner: Arc<Inner>,
177}
178
179impl Nats {
180    /// Creates a new Nats client and connects to the NATS server.
181    pub async fn new(config: Config) -> Result<Self, ConnectError> {
182        let client = async_nats::ConnectOptions::new()
183            .max_reconnects(Some(1))
184            .connect(config.address())
185            .await?;
186        Ok(Self {
187            inner: Arc::new(Inner { client, config }),
188        })
189    }
190
191    /// Returns a reference to the underlying NATS client.
192    pub fn client(&self) -> &Client {
193        &self.inner.client
194    }
195
196    /// Returns a reference to the configuration.
197    pub fn config(&self) -> &Config {
198        &self.inner.config
199    }
200
201    /// Creates a new event publisher.
202    pub async fn publisher(&self) -> Result<Publisher, CreateStreamError> {
203        let ctx = jetstream::new(self.inner.client.clone());
204        let p = Publisher { ctx };
205        p.init(&self.inner.config).await?;
206        Ok(p)
207    }
208
209    /// Creates a durable pull consumer with the given name.
210    pub async fn sys_consumer(&self, name: String) -> Result<PullConsumer, ConsumerError> {
211        let ctx = jetstream::new(self.inner.client.clone());
212        ctx.create_consumer_on_stream(
213            jetstream::consumer::pull::Config {
214                durable_name: Some(name),
215                ..Default::default()
216            },
217            self.inner.config.events_stream_name(),
218        )
219        .await
220    }
221
222    /// Creates a durable pull consumer with a filter subject.
223    pub async fn sys_consumer_with_filter(
224        &self,
225        name: String,
226        filter_subject: String,
227    ) -> Result<PullConsumer, ConsumerError> {
228        let ctx = jetstream::new(self.inner.client.clone());
229        ctx.create_consumer_on_stream(
230            jetstream::consumer::pull::Config {
231                durable_name: Some(name),
232                filter_subject,
233                ..Default::default()
234            },
235            self.inner.config.events_stream_name(),
236        )
237        .await
238    }
239
240    /// Creates a durable pull consumer with multiple filter subjects.
241    pub async fn sys_consumer_with_filters(
242        &self,
243        name: String,
244        filter_subjects: Vec<String>,
245    ) -> Result<PullConsumer, ConsumerError> {
246        let ctx = jetstream::new(self.inner.client.clone());
247        ctx.create_consumer_on_stream(
248            jetstream::consumer::pull::Config {
249                durable_name: Some(name),
250                filter_subjects,
251                ..Default::default()
252            },
253            self.inner.config.events_stream_name(),
254        )
255        .await
256    }
257
258    /// Creates a temporary pull consumer with a filter subject.
259    pub async fn tmp_sys_consumer_with_filter(
260        &self,
261        filter_subject: String,
262    ) -> Result<PullConsumer, ConsumerError> {
263        let ctx = jetstream::new(self.inner.client.clone());
264        ctx.create_consumer_on_stream(
265            jetstream::consumer::pull::Config {
266                filter_subject,
267                deliver_policy: jetstream::consumer::DeliverPolicy::Last,
268                ..Default::default()
269            },
270            self.inner.config.events_stream_name(),
271        )
272        .await
273    }
274
275    /// Creates a distributed locks manager.
276    pub async fn distributed_locks(&self) -> Result<DistributedLocks, DistributedLocksError> {
277        let ctx = jetstream::new(self.inner.client.clone());
278        DistributedLocks::new(ctx, &self.inner.config).await
279    }
280
281    /// Creates a sequence manager.
282    pub fn sequence_manager(&self) -> SequenceManager {
283        let ctx = jetstream::new(self.inner.client.clone());
284        SequenceManager { ctx }
285    }
286}
287
288/// Trait for converting events to NATS subjects.
289///
290/// Implement this trait on your event types to enable automatic
291/// subject path generation for event publishing.
292pub trait EventToSubject<M> {
293    /// Converts the event to a NATS subject.
294    fn event_to_subject(&self) -> async_nats::Subject;
295}
296
297/// Event publisher for NATS JetStream.
298///
299/// Manages the events stream and provides methods to publish events
300/// with structured subject paths.
301pub struct Publisher {
302    ctx: Context,
303}
304
305impl Publisher {
306    async fn init(&self, config: &Config) -> Result<(), CreateStreamError> {
307        let stream_name = self
308            .ctx
309            .stream_by_subject(config.events_stream_subject().to_string())
310            .await;
311        if let Err(err) = stream_name {
312            if matches!(err.kind(), GetStreamByNameErrorKind::NotFound) {
313                self.ctx
314                    .create_stream(jetstream::stream::Config {
315                        name: config.events_stream_name().to_string(),
316                        subjects: vec![config.events_stream_subject().into()],
317                        allow_direct: true,
318                        deny_delete: true,
319                        deny_purge: true,
320                        ..Default::default()
321                    })
322                    .await?;
323            }
324        }
325        Ok(())
326    }
327
328    /// Publishes an event to the given subject.
329    pub async fn publish<S: ToSubject, P: ?Sized + serde::Serialize>(
330        &self,
331        subject: S,
332        payload: &P,
333    ) -> anyhow::Result<()> {
334        self.ctx
335            .publish(subject, serde_json::to_vec(payload)?.into())
336            .await?;
337        Ok(())
338    }
339
340    /// Publishes an event using the subject derived from the event type.
341    pub async fn publish_event<S, M, P>(&self, subject: &S, payload: &P) -> anyhow::Result<()>
342    where
343        S: ?Sized + EventToSubject<M>,
344        P: ?Sized + serde::Serialize,
345    {
346        self.ctx
347            .publish(
348                subject.event_to_subject(),
349                serde_json::to_vec(payload)?.into(),
350            )
351            .await?;
352        Ok(())
353    }
354
355    /// Publishes an event using the subject derived from the event type and additional headers.
356    pub async fn publish_event_with_headers<S, M, P>(
357        &self,
358        subject: &S,
359        payload: &P,
360        headers: async_nats::HeaderMap,
361    ) -> anyhow::Result<()>
362    where
363        S: ?Sized + EventToSubject<M>,
364        P: ?Sized + serde::Serialize,
365    {
366        self.ctx
367            .publish_with_headers(
368                subject.event_to_subject(),
369                headers,
370                serde_json::to_vec(payload)?.into(),
371            )
372            .await?;
373        Ok(())
374    }
375}
376
377impl AsRef<Context> for Publisher {
378    fn as_ref(&self) -> &Context {
379        &self.ctx
380    }
381}
382
383/// Error type for distributed lock operations.
384#[derive(thiserror::Error, Debug)]
385pub enum DistributedLocksError {
386    /// Error connecting to NATS.
387    #[error(transparent)]
388    Connect(#[from] async_nats::error::Error<ConnectErrorKind>),
389    /// Error creating a key-value store.
390    #[error(transparent)]
391    CreateKeyValue(#[from] async_nats::error::Error<CreateKeyValueErrorKind>),
392    /// Error accessing a key-value store.
393    #[error(transparent)]
394    KeyValue(#[from] async_nats::error::Error<KeyValueErrorKind>),
395}
396
397/// Distributed locks manager.
398///
399/// Manages a key-value store for distributed locking across services.
400#[derive(Clone)]
401pub struct DistributedLocks {
402    ctx: Context,
403    sys_locks: String,
404}
405
406impl DistributedLocks {
407    async fn new(ctx: Context, config: &Config) -> Result<Self, DistributedLocksError> {
408        let lm = DistributedLocks {
409            ctx,
410            sys_locks: config.sys_locks().to_string(),
411        };
412        if !lm.exists(config.sys_locks()).await? {
413            lm.create(config.sys_locks(), 5).await?;
414        }
415        Ok(lm)
416    }
417
418    async fn create<T: Into<String>>(
419        &self,
420        name: T,
421        max_age: u64,
422    ) -> Result<Store, DistributedLocksError> {
423        Ok(self
424            .ctx
425            .create_key_value(kv::Config {
426                bucket: name.into(),
427                max_age: std::time::Duration::from_secs(max_age),
428                history: 1,
429                ..Default::default()
430            })
431            .await?)
432    }
433
434    async fn exists<T: Into<String>>(&self, bucket: T) -> Result<bool, DistributedLocksError> {
435        if let Err(err) = self.ctx.get_key_value(bucket).await {
436            if err.kind() == KeyValueErrorKind::GetBucket {
437                if let Some(src) = err.source() {
438                    let err = src.downcast_ref::<async_nats::error::Error<GetStreamErrorKind>>();
439                    if let Some(err) = err {
440                        if let GetStreamErrorKind::JetStream(err) = err.kind() {
441                            if err.code() == 404 {
442                                return Ok(false);
443                            }
444                        }
445                    }
446                }
447            }
448            Err(err)?;
449        }
450        Ok(true)
451    }
452
453    /// Returns a lock manager for system locks.
454    pub async fn sys_locks(&self) -> anyhow::Result<LockManager> {
455        let kv = self.ctx.get_key_value(&self.sys_locks).await?;
456        Ok(LockManager { kv: Arc::new(kv) })
457    }
458}
459
460/// Error type for lock manager operations.
461#[derive(thiserror::Error, Debug)]
462pub enum LockManagerError {
463    /// Error creating a key-value store.
464    #[error(transparent)]
465    CreateKeyValue(#[from] async_nats::error::Error<CreateKeyValueErrorKind>),
466    /// Error accessing a key-value store.
467    #[error(transparent)]
468    KeyValue(#[from] async_nats::error::Error<KeyValueErrorKind>),
469    /// Error watching key-value store changes.
470    #[error(transparent)]
471    Watch(#[from] async_nats::error::Error<kv::WatchErrorKind>),
472    /// Unable to acquire lock after exhausting retries.
473    #[error("unable to lock resource after {0:?}")]
474    OutOfRetries(std::time::Duration),
475}
476
477/// Error type for sequence manager operations.
478#[derive(thiserror::Error, Debug)]
479pub enum SequenceManagerError {
480    /// Error connecting to NATS.
481    #[error(transparent)]
482    Connect(#[from] async_nats::error::Error<ConnectErrorKind>),
483    /// Error creating a key-value store.
484    #[error(transparent)]
485    CreateKeyValue(#[from] async_nats::error::Error<CreateKeyValueErrorKind>),
486    /// Error accessing a key-value store.
487    #[error(transparent)]
488    KeyValue(#[from] async_nats::error::Error<KeyValueErrorKind>),
489    /// Error putting a value to the key-value store.
490    #[error(transparent)]
491    Put(#[from] async_nats::error::Error<async_nats::jetstream::kv::PutErrorKind>),
492    /// Error reading an entry from the key-value store.
493    #[error(transparent)]
494    Entry(#[from] async_nats::error::Error<async_nats::jetstream::kv::EntryErrorKind>),
495}
496
497/// Sequence manager for generating unique, monotonically increasing IDs.
498///
499/// Uses NATS JetStream key-value stores to maintain sequence counters
500/// that can be used across distributed services.
501pub struct SequenceManager {
502    ctx: Context,
503}
504
505impl SequenceManager {
506    async fn create<T: Into<String>>(&self, name: T) -> Result<Store, SequenceManagerError> {
507        Ok(self
508            .ctx
509            .create_key_value(kv::Config {
510                bucket: name.into(),
511                ..Default::default()
512            })
513            .await?)
514    }
515
516    async fn exists<T: Into<String>>(&self, bucket: T) -> Result<bool, SequenceManagerError> {
517        if let Err(err) = self.ctx.get_key_value(bucket).await {
518            if err.kind() == KeyValueErrorKind::GetBucket {
519                if let Some(src) = err.source() {
520                    let err = src.downcast_ref::<async_nats::error::Error<GetStreamErrorKind>>();
521                    if let Some(err) = err {
522                        if let GetStreamErrorKind::JetStream(err) = err.kind() {
523                            if err.code() == 404 {
524                                return Ok(false);
525                            }
526                        }
527                    }
528                }
529            }
530            Err(err)?;
531        }
532        Ok(true)
533    }
534
535    async fn get<T: Into<String>>(&self, bucket: T) -> Result<Store, SequenceManagerError> {
536        Ok(self.ctx.get_key_value(bucket).await?)
537    }
538
539    /// Gets the next sequence number for the given prefix, creating the bucket if needed.
540    pub async fn next(&self, prefix: &str, id: i64) -> Result<i64, SequenceManagerError> {
541        let bucket = format!("sm-{prefix}");
542        if !self.exists(&bucket).await? {
543            let store = self.create(&bucket).await?;
544            let result = store.put("id", id.to_be_bytes().to_vec().into()).await?;
545            Ok(result as i64)
546        } else {
547            let store = self.get(&bucket).await?;
548            let e = store.entry("id").await?;
549            if let Some(e) = e {
550                Ok(e.revision as i64)
551            } else {
552                let result = store.put("id", id.to_be_bytes().to_vec().into()).await?;
553                Ok(result as i64)
554            }
555        }
556    }
557
558    /// Increments the sequence number for the given prefix.
559    pub async fn increment(&self, prefix: &str, id: i64) -> Result<i64, SequenceManagerError> {
560        let bucket = format!("sm-{prefix}");
561        let store = self.get(&bucket).await?;
562        let e = store.put("id", id.to_be_bytes().to_vec().into()).await?;
563        Ok(e as i64)
564    }
565}
566
567/// Lock manager for acquiring and managing distributed locks.
568///
569/// Provides automatic lock acquisition and release with retry logic.
570/// Locks are automatically refreshed and released when the critical
571/// section completes or the lock holder crashes.
572pub struct LockManager {
573    kv: Arc<Store>,
574}
575
576impl LockManager {
577    /// Runs the given future while holding a distributed lock.
578    ///
579    /// The lock is automatically acquired before the future runs and released
580    /// when the future completes or the holder crashes.
581    pub async fn run_locked<N, O, F, E>(&self, name: N, f: F) -> Result<O, E>
582    where
583        N: Into<String>,
584        F: std::future::Future<Output = Result<O, E>>,
585        E: From<LockManagerError>,
586    {
587        let lock = self.try_lock(name.into(), 3, 5).await?;
588        let result = f.await;
589        let w_kv = self.kv.clone();
590        tokio::spawn(async move {
591            if !lock.jh.is_finished() {
592                lock.jh.abort();
593                let result = lock.jh.await;
594                if let Err(err) = result {
595                    if !err.is_cancelled() {
596                        tracing::error!("{err:#?}");
597                    }
598                }
599            }
600            w_kv.delete_expect_revision(lock.name, Some(lock.revision.load(Ordering::SeqCst)))
601                .await
602                .ok();
603        });
604        result
605    }
606
607    async fn try_lock(
608        &self,
609        name: String,
610        timeout: u64,
611        retries: usize,
612    ) -> Result<Lock, LockManagerError> {
613        let now = std::time::Instant::now();
614        let max_retries = retries;
615        let mut tries = 0;
616        let revision = Arc::new(AtomicU64::new(0));
617        let kv = &self.kv;
618        loop {
619            if tries >= max_retries {
620                return Err(LockManagerError::OutOfRetries(now.elapsed()));
621            }
622            let v = kv.create(&name, "r".into()).await;
623            if let Err(err) = v {
624                if err.kind() == async_nats::jetstream::kv::CreateErrorKind::AlreadyExists {
625                    tracing::debug!("seems to be locked already, {tries} try to watch for changes");
626                    let mut w = kv.watch(&name).await?;
627                    let f = async {
628                        'inner: while let Some(m) = w.next().await {
629                            if let Ok(e) = m {
630                                if e.operation == Operation::Delete {
631                                    tracing::debug!("retry because prev lock was deleted");
632                                    break 'inner;
633                                }
634                            }
635                        }
636                    };
637                    let t = async {
638                        tokio::time::sleep(std::time::Duration::from_secs(timeout)).await;
639                    };
640                    let change = tokio::select! {
641                        _ = f => true,
642                        _ = t => false,
643                    };
644                    if !change {
645                        tries += 1;
646                    }
647                }
648            } else {
649                let r = v.unwrap();
650                revision.store(r, Ordering::SeqCst);
651                tracing::debug!("got lock: '{name}'");
652                break;
653            }
654        }
655        let w_kv = self.kv.clone();
656        let w_name = name.clone();
657        let w_revision = revision.clone();
658
659        let jh = tokio::spawn(async move {
660            let mut run = 0;
661            loop {
662                run += 1;
663                tokio::time::sleep(std::time::Duration::from_secs(2)).await;
664                tracing::debug!("refresh lock {w_name}");
665                let result = w_kv
666                    .update(&w_name, "u".into(), w_revision.load(Ordering::SeqCst))
667                    .await;
668                if let Err(err) = result {
669                    tracing::error!("{err:#?}");
670                    break;
671                } else {
672                    w_revision.store(result.unwrap(), Ordering::SeqCst);
673                }
674                if run >= 5 {
675                    tracing::debug!("release lock after timeout");
676                    break;
677                }
678            }
679            anyhow::Ok(())
680        });
681
682        Ok(Lock { name, revision, jh })
683    }
684}
685
686/// State of a distributed lock.
687#[derive(Debug, PartialEq, Eq)]
688pub enum LockState {
689    /// Lock is being acquired.
690    Registering,
691    /// Lock has been acquired.
692    Registered,
693}
694
695/// A distributed lock handle.
696///
697/// Represents an acquired lock. The lock is automatically released
698/// when the handle is dropped or the holder crashes.
699#[derive(Debug)]
700pub struct Lock {
701    name: String,
702    revision: Arc<AtomicU64>,
703    jh: JoinHandle<anyhow::Result<()>>,
704}
qm_nats/lib.rs

qm_nats/
lib.rs
