Skip to main content

arminius_lib_nats/
client.rs

1// SPDX-FileCopyrightText: 2025 BWI GmbH
2//
3// SPDX-License-Identifier: Apache-2.0
4
5//! The NATS connection handle ([`NatsClient`]) and its shared alias
6//! ([`SharedNatsClient`]).
7
8mod maintenance;
9mod request;
10pub mod room;
11mod sync;
12
13use std::sync::Arc;
14
15use arminius_lib_message_bus::Event;
16use async_nats::jetstream::{self, kv};
17pub use request::NatsRequestBody;
18
19use crate::{NatsError, config::MessageBusConfig, resubscribe::ConnectionSignals};
20
21/// Live handle to a NATS client including a JetStream connection.
22///
23/// Implements [`Publisher`](arminius_lib_message_bus::Publisher),
24/// [`Subscriber`](arminius_lib_message_bus::Subscriber),
25/// [`Requester`](arminius_lib_message_bus::Requester), and
26/// [`Responder`](arminius_lib_message_bus::Responder).
27///
28/// Obtain one via [`NatsConfiguration::connect`](crate::NatsConfiguration).
29/// The cheap-to-clone [`SharedNatsClient`] alias
30/// (`Arc<NatsClient>`) is the type passed around in application code.
31#[derive(Clone, Debug)]
32pub struct NatsClient {
33	/// NATS connection.
34	// TODO: deprecate, kept public for callers that still need the raw
35	// `async_nats::Client`. New code should use the message-bus APIs.
36	pub client: async_nats::client::Client,
37	/// NATS Service for the component that owns this connection.
38	///
39	/// Having a NATS service allows us to use NATS to see which components are
40	/// online, what version they're using, and do additional introspection in
41	/// the long term (currently only pings work).
42	pub service: Arc<async_nats::service::Service>,
43	/// JetStream context, used by the
44	/// [`Publisher`](arminius_lib_message_bus::Publisher),
45	/// [`Subscriber`](arminius_lib_message_bus::Subscriber),
46	/// [`Requester`](arminius_lib_message_bus::Requester), and
47	/// [`Responder`](arminius_lib_message_bus::Responder) implementations.
48	pub(crate) jetstream: jetstream::Context,
49	/// Durable consumer name; instances sharing the same name load-balance
50	/// automatically.
51	pub(crate) service_name: String,
52	/// Shared reconnect signals, also held by the connection's event callback.
53	pub(crate) signals: Arc<ConnectionSignals>,
54	/// Message bus configuration to create different kinds of streams.
55	pub(crate) stream_config: MessageBusConfig,
56}
57
58/// A connection to NATS and JetStream.
59// TODO: make NatsClient a cloneable struct that hold an inner value to an
60// Arc<ClientInner> and ClientInner can be pub(crate) at best.
61pub type SharedNatsClient = Arc<NatsClient>;
62
63impl NatsClient {
64	/// Buffer size for the internal channel between a subscription's background
65	/// forward task and its consumer-facing stream. See
66	/// [`ConsumerConfig::buffer_size`](crate::config::ConsumerConfig::buffer_size).
67	pub(crate) fn internal_channel_size(&self) -> usize {
68		self.stream_config.consumer.buffer_size()
69	}
70
71	/// Given a JetStream context, returns a key-value store
72	pub async fn connect_to_store(
73		jetstream: &jetstream::context::Context,
74	) -> Result<kv::Store, NatsError> {
75		let store = jetstream
76			.create_key_value(kv::Config {
77				bucket: "arminius_store".to_owned(),
78				history: 1,
79				..Default::default()
80			})
81			.await?;
82		Ok(store)
83	}
84
85	/// Ensures the JetStream event stream for `E` exists and returns its
86	/// handle.
87	///
88	/// Shared by [`Publisher::publish`](arminius_lib_message_bus::Publisher)
89	/// (which discards the handle, just ensuring the stream exists) and
90	/// [`Subscriber::subscribe`](arminius_lib_message_bus::Subscriber) (which
91	/// uses the handle to create its consumer). The stream covers `E::SUBJECT`
92	/// plus any `COMPATIBLE_SUBJECTS`, using the connection's event retention
93	/// policy.
94	pub(crate) async fn ensure_event_stream<E: Event>(
95		&self,
96	) -> Result<jetstream::stream::Stream, jetstream::context::CreateStreamError> {
97		let config = self
98			.stream_config
99			.retention
100			.build_stream_config(stream_name(E::SUBJECT), event_subjects::<E>());
101		self.jetstream.get_or_create_stream(config).await
102	}
103
104	/// Ensures the `WorkQueue` JetStream stream for a command family exists and
105	/// returns its handle.
106	///
107	/// Mirrors [`NatsClient::ensure_event_stream`]. Shared by the responder
108	/// ([`Responder::respond`](arminius_lib_message_bus::Responder)) and the
109	/// requester ([`Requester::request`](arminius_lib_message_bus::Requester))
110	/// so both agree on a single stream per `SUBJECT_FILTER`. The responder
111	/// uses the handle to create its consumer; the requester just needs the
112	/// stream to exist before publishing (and discards the handle), so a
113	/// command sent before its responder is up is still durably captured.
114	///
115	/// The retention policy is always `WorkQueue`; only the limits come from
116	/// [`MessageBusConfig::command_retention`](crate::config::MessageBusConfig::command_retention).
117	pub(crate) async fn ensure_command_stream(
118		&self,
119		subject_filter: &str,
120	) -> Result<jetstream::stream::Stream, jetstream::context::CreateStreamError> {
121		let config = self
122			.stream_config
123			.command_retention
124			.build_stream_config(stream_name(subject_filter), vec![subject_filter.to_owned()]);
125		self.jetstream.get_or_create_stream(config).await
126	}
127}
128
129/// Derives a JetStream stream name from a (possibly wildcard) subject pattern.
130///
131/// Stream names cannot contain `.`, `*`, or `>`, so each is mapped to `_`. All
132/// instances of a per-instance event share one pattern (e.g.
133/// `rooms.*.events.v1`) and therefore one stable stream name
134/// (`rooms__events_v1`).
135pub fn stream_name(pattern: &str) -> String {
136	pattern.chars().map(|char| if matches!(char, '.' | '*' | '>') { '_' } else { char }).collect()
137}
138
139/// The full set of subjects an event occupies: its [`Event::SUBJECT`] plus any
140/// [`Event::COMPATIBLE_SUBJECTS`].
141///
142/// Used both to create the event stream and to filter the subscriber's
143/// consumer, so the two always agree.
144pub fn event_subjects<TEvent: Event>() -> Vec<String> {
145	std::iter::once(TEvent::SUBJECT)
146		.chain(TEvent::COMPATIBLE_SUBJECTS.iter().copied())
147		.map(ToOwned::to_owned)
148		.collect()
149}
150
151#[cfg(test)]
152mod tests;