Skip to main content

arminius_lib_nats/
config.rs

1// SPDX-FileCopyrightText: 2025 BWI GmbH
2//
3// SPDX-License-Identifier: Apache-2.0
4
5//! Configuration for connecting to NATS.
6use std::{
7	num::NonZeroUsize,
8	path::PathBuf,
9	sync::{
10		Arc,
11		atomic::{AtomicBool, Ordering},
12	},
13};
14
15use arminius_lib_config::ArminiusConfigurationError;
16use async_nats::{
17	ConnectError, ConnectOptions,
18	jetstream::{
19		self, consumer,
20		context::CreateKeyValueError,
21		stream::{self as nats_stream, DiscardPolicy, RetentionPolicy},
22	},
23	service::ServiceExt as _,
24};
25
26use crate::{NatsClient, NatsCredentials, SharedNatsClient, resubscribe::ConnectionSignals};
27
28/// Errors that can occur while building or connecting the NATS client.
29#[derive(thiserror::Error, Debug)]
30pub enum NatsConfigurationError {
31	/// The NATS client could not be started.
32	#[error("Could not start NATS client: {0}")]
33	StartClient(String),
34	/// The server's version is incompatible with this client.
35	#[error("NATS server version is not compatible: {0}")]
36	Version(String),
37	/// Connecting to the NATS server failed.
38	#[error(transparent)]
39	Connect(#[from] ConnectError),
40	/// Creating or accessing a JetStream key-value store failed.
41	#[error(transparent)]
42	KeyValue(#[from] CreateKeyValueError),
43	/// A failure with no more specific variant, e.g. a missing `NATS_URLS`.
44	#[error("Other Error: {0}")]
45	Other(String),
46	/// The NATS credentials could not be read from the environment.
47	#[error("Could not read NATS Credentials from environment: {0}")]
48	Credential(#[from] ArminiusConfigurationError),
49}
50
51/// Controls how long messages are retained in a JetStream stream.
52///
53/// The retention policy is immutable after stream creation — unlike most other
54/// stream settings, it cannot be changed via `update_stream` and requires
55/// deleting and recreating the stream.
56#[derive(Debug, Clone, PartialEq, Eq)]
57pub enum StreamRetentionConfig {
58	/// Keep messages only while at least one consumer group is interested.
59	/// Messages are removed once every subscribed consumer has acknowledged
60	/// them. No event replay for consumers that join later.
61	Interest,
62
63	/// Keep messages until explicit size/age/count limits are reached, then
64	/// discard the oldest. `None` on any limit means unlimited.
65	/// Enables event replay for consumers that join after the fact.
66	Limits {
67		/// Drop messages older than this age (`None` = unlimited).
68		max_age: Option<std::time::Duration>,
69		/// Drop oldest messages when count exceeds this (`None` = unlimited).
70		max_messages: Option<i64>,
71		/// Drop oldest messages when total bytes exceeds this (`None` =
72		/// unlimited).
73		max_bytes: Option<i64>,
74	},
75
76	/// Delete messages as soon as any consumer acknowledges them.
77	/// Suitable for work-queue patterns where each message should be
78	/// processed exactly once by any one worker.
79	WorkQueue {
80		/// Drop messages older than this age (`None` = unlimited).
81		max_age: Option<std::time::Duration>,
82		/// Drop oldest messages when count exceeds this (`None` = unlimited).
83		max_messages: Option<i64>,
84		/// Drop oldest messages when total bytes exceeds this (`None` =
85		/// unlimited).
86		max_bytes: Option<i64>,
87	},
88}
89
90impl Default for StreamRetentionConfig {
91	fn default() -> Self {
92		// Unlimited Limits: event replay is available for consumers that join
93		// later.
94		Self::Limits { max_age: None, max_messages: None, max_bytes: None }
95	}
96}
97
98impl StreamRetentionConfig {
99	/// Builds a [`nats_stream::Config`] for `get_or_create_stream` calls.
100	pub(crate) fn build_stream_config(
101		&self,
102		name: String,
103		subjects: Vec<String>,
104	) -> nats_stream::Config {
105		const UNLIMITED: i64 = -1;
106		let (retention, max_age, max_messages, max_bytes) = match self {
107			Self::Interest => {
108				(RetentionPolicy::Interest, std::time::Duration::ZERO, UNLIMITED, UNLIMITED)
109			}
110			Self::Limits { max_age, max_messages, max_bytes } => (
111				RetentionPolicy::Limits,
112				max_age.unwrap_or(std::time::Duration::ZERO),
113				max_messages.unwrap_or(UNLIMITED),
114				max_bytes.unwrap_or(UNLIMITED),
115			),
116			Self::WorkQueue { max_age, max_messages, max_bytes } => (
117				RetentionPolicy::WorkQueue,
118				max_age.unwrap_or(std::time::Duration::ZERO),
119				max_messages.unwrap_or(UNLIMITED),
120				max_bytes.unwrap_or(UNLIMITED),
121			),
122		};
123		nats_stream::Config {
124			name,
125			subjects,
126			discard: DiscardPolicy::Old,
127			retention,
128			max_age,
129			max_messages,
130			max_bytes,
131			..Default::default()
132		}
133	}
134}
135
136/// Retention limits for command request-reply streams.
137///
138/// The retention *policy* is fixed to `WorkQueue` and cannot be configured:
139/// command dispatch relies on WorkQueue semantics (each command delivered to
140/// exactly one responder and removed on ack). Only these limits are tunable;
141/// `None` means unlimited.
142#[derive(Debug, Clone, Default, PartialEq, Eq)]
143pub struct CommandRetentionConfig {
144	/// Drop messages older than this age (`None` = unlimited).
145	pub max_age: Option<std::time::Duration>,
146	/// Drop oldest messages when count exceeds this (`None` = unlimited).
147	pub max_messages: Option<i64>,
148	/// Drop oldest messages when total bytes exceeds this (`None` = unlimited).
149	pub max_bytes: Option<i64>,
150}
151
152impl CommandRetentionConfig {
153	/// Builds a `WorkQueue` [`nats_stream::Config`] for a command stream from
154	/// these limits. The policy is always `WorkQueue` by construction.
155	pub(crate) fn build_stream_config(
156		&self,
157		name: String,
158		subjects: Vec<String>,
159	) -> nats_stream::Config {
160		StreamRetentionConfig::WorkQueue {
161			max_age: self.max_age,
162			max_messages: self.max_messages,
163			max_bytes: self.max_bytes,
164		}
165		.build_stream_config(name, subjects)
166	}
167}
168
169/// How many times a message is redelivered before the broker gives up on it.
170pub const DEFAULT_MAX_DELIVER: i64 = 5;
171
172/// How many messages the broker will have unacknowledged with one consumer.
173///
174/// Kept within the client's internal buffer so the broker stops delivering
175/// before that buffer overflows.
176pub const DEFAULT_MAX_ACK_PENDING: i64 = 256;
177
178/// Capacity of the channel between a subscription's background forward task and
179/// its consumer-facing stream, when
180/// [`ConsumerConfig::stream_buffer_size`] is `None`.
181pub const DEFAULT_STREAM_BUFFER_SIZE: usize = 1000;
182
183/// How many commands one responder pulls from the broker ahead of handling
184/// them.
185pub const DEFAULT_COMMAND_PREFETCH: NonZeroUsize = NonZeroUsize::MIN;
186
187/// Controls JetStream consumer behaviour: acknowledgment window, backpressure,
188/// poison-pill protection, and internal buffer sizes.
189#[derive(Debug, Clone)]
190pub struct ConsumerConfig {
191	/// How long the server waits for an ack before redelivering the message.
192	/// Default: 30 s.
193	pub ack_wait: std::time::Duration,
194	/// Maximum in-flight (unacked) messages. `None` uses the NATS default
195	/// (1 000). Set lower to apply backpressure on slow consumers.
196	///
197	/// Defaults to [`DEFAULT_MAX_ACK_PENDING`].
198	pub max_ack_pending: Option<i64>,
199	/// Maximum delivery attempts per message. `None` means unlimited retries.
200	/// Set to a small number to drop poison-pill messages automatically.
201	///
202	/// Defaults to [`DEFAULT_MAX_DELIVER`], leaving it unlimited would turn a
203	/// message no consumer can handle into an endless redelivery loop.
204	///
205	/// Overridden to 1 for a command declaring
206	/// [`Delivery::AtMostOnce`](arminius_lib_message_bus::Delivery::AtMostOnce):
207	/// such a command is acknowledged before it is handled, so any redelivery
208	/// would mean the handler ran twice, whatever this is set to.
209	pub max_deliver: Option<i64>,
210	/// Capacity of the internal mpsc channel used by the auto-reconnect task.
211	/// `None` falls back to 1 000.
212	pub stream_buffer_size: Option<usize>,
213	/// How many commands one responder pulls ahead of handling them. Keep it
214	/// small so idle instances get the queued ones. Defaults to
215	/// [`DEFAULT_COMMAND_PREFETCH`].
216	pub command_prefetch: NonZeroUsize,
217}
218
219impl Default for ConsumerConfig {
220	fn default() -> Self {
221		Self {
222			ack_wait: std::time::Duration::from_secs(30),
223			max_ack_pending: Some(DEFAULT_MAX_ACK_PENDING),
224			max_deliver: Some(DEFAULT_MAX_DELIVER),
225			stream_buffer_size: None,
226			command_prefetch: DEFAULT_COMMAND_PREFETCH,
227		}
228	}
229}
230
231impl ConsumerConfig {
232	/// Capacity of the channel between a subscription's background forward task
233	/// and its consumer-facing stream, falling back to
234	/// [`DEFAULT_STREAM_BUFFER_SIZE`].
235	#[must_use]
236	pub fn buffer_size(&self) -> usize {
237		self.stream_buffer_size.unwrap_or(DEFAULT_STREAM_BUFFER_SIZE)
238	}
239
240	/// Builds a durable JetStream pull-consumer config for `durable_name`,
241	/// filtering on `filter_subjects`, with this struct's tuning applied (ack
242	/// window, backpressure, poison-pill limit). `None` limits keep the
243	/// `async-nats` defaults.
244	///
245	/// All subjects go in `filter_subjects` (requires NATS 2.10+, which the
246	/// connection mandates). `deliver_policy` is `All`; `ack_policy` stays at
247	/// the `async-nats` default (`Explicit`, as required for pull consumers).
248	///
249	/// The in-flight limit is capped at [`buffer_size`](Self::buffer_size).
250	/// The two are separate fields, so shrinking the buffer alone would
251	/// otherwise let the broker deliver more than the client can hold, and the
252	/// surplus would be dropped and redelivered purely because of
253	/// backpressure.
254	pub(crate) fn pull_config(
255		&self,
256		durable_name: String,
257		filter_subjects: Vec<String>,
258	) -> consumer::pull::Config {
259		let mut config = consumer::pull::Config {
260			durable_name: Some(durable_name),
261			filter_subjects,
262			deliver_policy: consumer::DeliverPolicy::All,
263			ack_wait: self.ack_wait,
264			..Default::default()
265		};
266		if let Some(max_ack_pending) = self.max_ack_pending {
267			config.max_ack_pending =
268				max_ack_pending.min(i64::try_from(self.buffer_size()).unwrap_or(i64::MAX));
269		}
270		if let Some(max_deliver) = self.max_deliver {
271			config.max_deliver = max_deliver;
272		}
273		config
274	}
275}
276
277/// Message-bus tuning carried by a live connection.
278///
279/// Groups the stream/consumer settings the
280/// [`Publisher`](arminius_lib_message_bus::Publisher),
281/// [`Subscriber`](arminius_lib_message_bus::Subscriber), and
282/// [`Responder`](arminius_lib_message_bus::Responder) implementations need, so
283/// the connection holds one config value rather than several. Connection-only
284/// inputs (URLs, credentials, TLS) stay on [`NatsConfiguration`] and are not
285/// retained after connect.
286#[derive(Debug, Clone, Default)]
287pub struct MessageBusConfig {
288	/// Stream retention policy applied when an event stream is first created.
289	pub retention: StreamRetentionConfig,
290	/// Consumer behaviour (ack timeout, backpressure, poison-pill limits) for
291	/// the subscriber and responder.
292	pub consumer: ConsumerConfig,
293	/// Tunable limits for command request-reply (`WorkQueue`) streams.
294	pub command_retention: CommandRetentionConfig,
295}
296
297/// Options for client TLS authentication with NATS
298#[derive(Clone, Debug)]
299pub struct NatsClientTls {
300	/// Path to a TLS certificate
301	pub cert: PathBuf,
302	/// Path to a TLS key
303	pub key: PathBuf,
304}
305
306/// Configuration for connecting to NATS.
307///
308/// Construct one with [`from_env`](NatsConfiguration::from_env), then establish
309/// a connection with [`connect`](crate::NatsClient).
310#[derive(Clone, Debug)]
311pub struct NatsConfiguration {
312	/// List of URLs to connect to NATS over
313	pub urls: Vec<String>,
314	/// The credentials for authentication against NATS
315	// TODO this should also be an Option<NatsCredentials>
316	pub credentials: NatsCredentials,
317	/// Optional client TLS config
318	pub client_tls: Option<NatsClientTls>,
319	/// Configuration for NATS streams
320	pub stream_config: MessageBusConfig,
321}
322
323impl NatsConfiguration {
324	/// Creates a new NATS configuration from environment variables
325	pub fn from_env() -> Result<Self, NatsConfigurationError> {
326		let urls = Self::nats_url_from_env()?;
327		let credentials = NatsCredentials::from_env()?;
328
329		let client_tls = if let Ok(cert) = std::env::var("ARMINIUS_NATS_CERT")
330			&& let Ok(key) = std::env::var("ARMINIUS_NATS_KEY")
331		{
332			let cert = PathBuf::from(cert);
333			let key = PathBuf::from(key);
334			Some(NatsClientTls { cert, key })
335		} else {
336			None
337		};
338		Ok(NatsConfiguration {
339			urls,
340			credentials,
341			client_tls,
342			stream_config: MessageBusConfig::default(),
343		})
344	}
345
346	/// Creates a new NATS configuration from environment variables, using the
347	/// given Credentials This is primarily for unit tests, where separate
348	/// accounts are created for test isolation.
349	pub fn with_credentials(credentials: NatsCredentials) -> Result<Self, NatsConfigurationError> {
350		let env_config = NatsConfiguration::from_env()?;
351		Ok(NatsConfiguration { credentials, ..env_config })
352	}
353
354	/// reads NATS_URLS from env
355	pub(crate) fn nats_url_from_env() -> Result<Vec<String>, NatsConfigurationError> {
356		let nats_urls = std::env::var("NATS_URLS").map_err(|_| {
357			NatsConfigurationError::Other("Missing `NATS_URLS` environment variable".into())
358		})?;
359		Ok(nats_urls.split(" ").map(ToOwned::to_owned).collect())
360	}
361
362	/// Return a NATS connection, given the above parameters
363	pub async fn connect(
364		&self,
365		service_name: String,
366		version: String,
367		description: String,
368	) -> Result<SharedNatsClient, NatsConfigurationError> {
369		let signals = Arc::new(ConnectionSignals {
370			connected: Arc::new(AtomicBool::new(false)),
371			reconnect_notify: Arc::new(tokio::sync::Notify::new()),
372		});
373		let signals_cb = signals.clone();
374
375		// FIXME(#1377): Normalize the content of the variable.
376		let lazy = std::env::var("ARMINIUS_LAZY").ok().filter(|var| !var.is_empty()).is_some();
377		let mut options: ConnectOptions = ConnectOptions::from(self.credentials.clone())
378			.event_callback(move |event| {
379				let signals = signals_cb.clone();
380				async move {
381					match event {
382						async_nats::Event::Disconnected => {
383							signals.connected.store(false, Ordering::SeqCst);
384							tracing::error!("disconnected from NATS server");
385						}
386						async_nats::Event::Connected => {
387							signals.connected.store(true, Ordering::SeqCst);
388							signals.reconnect_notify.notify_waiters();
389							tracing::info!("reconnected to NATS server");
390						}
391						_ => {}
392					}
393				}
394			});
395
396		if lazy {
397			options = options.retry_on_initial_connect();
398		}
399
400		if let Some(client_tls) = &self.client_tls {
401			options =
402				options.add_client_certificate(client_tls.cert.clone(), client_tls.key.clone());
403		}
404
405		let client = async_nats::connect_with_options(self.urls.clone(), options)
406			.await
407			.inspect_err(|err| {
408				tracing::error!(
409					"connecting to NATS server failed. credentials={credentials:?}, cause={err:?}",
410					credentials = &self.credentials
411				);
412			})?;
413
414		// We're connected at this point. The `Connected` event fires on
415		// reconnects, not on the initial connect, so we set the flag
416		// explicitly here.
417		signals.connected.store(true, Ordering::SeqCst);
418
419		// The version check returns "not compatible" when in initial retry mode
420		if !lazy {
421			// Version 2.10 introduces subject-transforms on streams
422			if !client.is_server_compatible(2, 10, 0) {
423				return Err(NatsConfigurationError::Version(
424					"Arminius requires NATS >= 2.10".into(),
425				));
426			}
427			tracing::info!(nats_server_version=%client.server_info().version, "Connected to nats version");
428		}
429
430		let service = client
431			.service_builder()
432			.description(description)
433			.start(service_name.clone(), version)
434			.await
435			.map_err(|err| {
436				tracing::error!(error = err.to_string(), "Failed to start NATS service");
437				NatsConfigurationError::StartClient(err.to_string())
438			})?;
439		Ok(Arc::new(NatsClient {
440			jetstream: jetstream::new(client.clone()),
441			client,
442			service: Arc::new(service),
443			service_name,
444			signals,
445			stream_config: self.stream_config.clone(),
446		}))
447	}
448}
449
450#[cfg(test)]
451mod tests;