1use 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#[derive(thiserror::Error, Debug)]
30pub enum NatsConfigurationError {
31 #[error("Could not start NATS client: {0}")]
33 StartClient(String),
34 #[error("NATS server version is not compatible: {0}")]
36 Version(String),
37 #[error(transparent)]
39 Connect(#[from] ConnectError),
40 #[error(transparent)]
42 KeyValue(#[from] CreateKeyValueError),
43 #[error("Other Error: {0}")]
45 Other(String),
46 #[error("Could not read NATS Credentials from environment: {0}")]
48 Credential(#[from] ArminiusConfigurationError),
49}
50
51#[derive(Debug, Clone, PartialEq, Eq)]
57pub enum StreamRetentionConfig {
58 Interest,
62
63 Limits {
67 max_age: Option<std::time::Duration>,
69 max_messages: Option<i64>,
71 max_bytes: Option<i64>,
74 },
75
76 WorkQueue {
80 max_age: Option<std::time::Duration>,
82 max_messages: Option<i64>,
84 max_bytes: Option<i64>,
87 },
88}
89
90impl Default for StreamRetentionConfig {
91 fn default() -> Self {
92 Self::Limits { max_age: None, max_messages: None, max_bytes: None }
95 }
96}
97
98impl StreamRetentionConfig {
99 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#[derive(Debug, Clone, Default, PartialEq, Eq)]
143pub struct CommandRetentionConfig {
144 pub max_age: Option<std::time::Duration>,
146 pub max_messages: Option<i64>,
148 pub max_bytes: Option<i64>,
150}
151
152impl CommandRetentionConfig {
153 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
169pub const DEFAULT_MAX_DELIVER: i64 = 5;
171
172pub const DEFAULT_MAX_ACK_PENDING: i64 = 256;
177
178pub const DEFAULT_STREAM_BUFFER_SIZE: usize = 1000;
182
183pub const DEFAULT_COMMAND_PREFETCH: NonZeroUsize = NonZeroUsize::MIN;
186
187#[derive(Debug, Clone)]
190pub struct ConsumerConfig {
191 pub ack_wait: std::time::Duration,
194 pub max_ack_pending: Option<i64>,
199 pub max_deliver: Option<i64>,
210 pub stream_buffer_size: Option<usize>,
213 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 #[must_use]
236 pub fn buffer_size(&self) -> usize {
237 self.stream_buffer_size.unwrap_or(DEFAULT_STREAM_BUFFER_SIZE)
238 }
239
240 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#[derive(Debug, Clone, Default)]
287pub struct MessageBusConfig {
288 pub retention: StreamRetentionConfig,
290 pub consumer: ConsumerConfig,
293 pub command_retention: CommandRetentionConfig,
295}
296
297#[derive(Clone, Debug)]
299pub struct NatsClientTls {
300 pub cert: PathBuf,
302 pub key: PathBuf,
304}
305
306#[derive(Clone, Debug)]
311pub struct NatsConfiguration {
312 pub urls: Vec<String>,
314 pub credentials: NatsCredentials,
317 pub client_tls: Option<NatsClientTls>,
319 pub stream_config: MessageBusConfig,
321}
322
323impl NatsConfiguration {
324 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 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 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 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 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 signals.connected.store(true, Ordering::SeqCst);
418
419 if !lazy {
421 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;