arminius_lib_nats/

resubscribe.rs

// SPDX-FileCopyrightText: 2026 BWI GmbH
//
// SPDX-License-Identifier: Apache-2.0

//! Functionality to detect NATS disconnects and seamless reconnects

use std::{
	ops::ControlFlow,
	pin::Pin,
	sync::{
		Arc,
		atomic::{AtomicBool, Ordering},
	},
	time::Duration,
};

use arminius_lib_message_bus::Error;
use backon::{BackoffBuilder as _, ExponentialBuilder};
use futures::{Stream, StreamExt as _};

/// Maximum backoff delay between re-subscription attempts (also the fallback
/// delay when the backoff iterator is exhausted).
const RECONNECT_MAX_DELAY: Duration = Duration::from_secs(32);

/// Boxed stream of the messages a subscription delivers; re-created on every
/// re-subscribe.
pub type MessageStream<T> = Pin<Box<dyn Stream<Item = T> + Send>>;

/// Reconnect signals shared between the NATS event callback and the
/// subscriber's background task.
#[derive(Debug)]
pub struct ConnectionSignals {
	/// `true` when the connection is alive; `false` after a `Disconnected`
	/// event.
	pub connected: Arc<AtomicBool>,
	/// Fired on every `Connected` event so waiting tasks can attempt to
	/// re-subscribe.
	pub reconnect_notify: Arc<tokio::sync::Notify>,
}

/// Outcome returned by [`forward_stream`].
pub enum ForwardOutcome {
	/// The upstream stream ended or a disconnection was detected; re-subscribe.
	Resubscribe,
	/// The receiver end of the channel was dropped; the background task should
	/// stop.
	ReceiverDropped,
}

/// Forwards items from `stream` into `tx`, watching `signals.connected` for
/// disconnections.
///
/// Returns [`ForwardOutcome::Resubscribe`] when re-subscription is needed and
/// [`ForwardOutcome::ReceiverDropped`] when the channel receiver has been
/// closed.
pub async fn forward_stream<T: Send + 'static>(
	signals: &ConnectionSignals,
	subject: &'static str,
	mut stream: MessageStream<T>,
	tx: &tokio::sync::mpsc::Sender<T>,
) -> ForwardOutcome {
	while let Some(item) = stream.next().await {
		if !signals.connected.load(Ordering::SeqCst) {
			tracing::warn!(subject, "Detected disconnection, will re-subscribe after reconnection");
			return ForwardOutcome::Resubscribe;
		}
		// The only point of the try_send is to detect backpressure
		match tx.try_send(item) {
			Ok(()) => {}
			Err(tokio::sync::mpsc::error::TrySendError::Full(item)) => {
				// TODO: We should have OTel metrics here.
				tracing::warn!(subject, "Internal channel full, waiting to send");
				if let Err(error) = tx.send(item).await {
					tracing::error!(subject, %error, "Receiver dropped while waiting to send");
					return ForwardOutcome::ReceiverDropped;
				}
			}
			Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
				tracing::warn!(subject, "Receiver closed while forwarding");
				return ForwardOutcome::ReceiverDropped;
			}
		}
	}

	tracing::warn!(subject, "Stream closed, will re-subscribe");
	ForwardOutcome::Resubscribe
}

/// Waits until the connection is usable again.
///
/// Returns immediately while connected. Otherwise it parks until the NATS event
/// callback signals a reconnect, or until the receiver is dropped — the latter
/// yields [`ControlFlow::Break`], telling the caller to stop.
async fn await_connection<T>(
	signals: &ConnectionSignals,
	tx: &tokio::sync::mpsc::Sender<T>,
) -> ControlFlow<()> {
	if signals.connected.load(Ordering::SeqCst) {
		return ControlFlow::Continue(());
	}

	tokio::select! {
		() = signals.reconnect_notify.notified() => ControlFlow::Continue(()),
		() = tx.closed() => ControlFlow::Break(()),
	}
}

/// Re-creates the message stream once the connection is back, retrying failed
/// attempts with exponential backoff.
///
/// Returns `None` when the receiver was dropped while waiting or retrying: with
/// nobody left to receive the messages, the background task should stop.
async fn resubscribe<T, F, Fut>(
	signals: &ConnectionSignals,
	subject: &'static str,
	make_stream: &F,
	tx: &tokio::sync::mpsc::Sender<T>,
) -> Option<MessageStream<T>>
where
	F: Fn() -> Fut + Sync,
	Fut: Future<Output = Result<MessageStream<T>, Error>>,
{
	// Built per call, so the delay only grows across consecutive failed
	// attempts and starts over once a subscription succeeds.
	let mut backoff = ExponentialBuilder::default()
		.with_max_delay(RECONNECT_MAX_DELAY)
		.without_max_times()
		.build();

	loop {
		if await_connection(signals, tx).await.is_break() {
			return None;
		}

		match make_stream().await {
			Ok(stream) => {
				tracing::debug!(subject, "Re-subscribed after reconnection");
				return Some(stream);
			}
			Err(error) => {
				let delay = backoff.next().unwrap_or(RECONNECT_MAX_DELAY);
				tracing::error!(
					subject,
					%error,
					delay_secs = delay.as_secs(),
					"Failed to re-subscribe, retrying with backoff"
				);
				tokio::time::sleep(delay).await;
			}
		}
	}
}

/// Body of the background task behind [`resubscribing_stream`]: forwards
/// `initial` into `tx`, then alternates between re-subscribing and forwarding
/// until the receiver is dropped.
async fn forward_and_resubscribe<T, F, Fut>(
	signals: &ConnectionSignals,
	subject: &'static str,
	initial: MessageStream<T>,
	make_stream: F,
	tx: tokio::sync::mpsc::Sender<T>,
) where
	T: Send + 'static,
	F: Fn() -> Fut + Sync,
	Fut: Future<Output = Result<MessageStream<T>, Error>>,
{
	let mut stream = initial;

	loop {
		if matches!(
			forward_stream(signals, subject, stream, &tx).await,
			ForwardOutcome::ReceiverDropped
		) {
			break;
		}

		match resubscribe(signals, subject, &make_stream, &tx).await {
			Some(next) => stream = next,
			None => break,
		}
	}

	tracing::debug!(subject, "Receiver dropped, stopping subscription task");
}

/// Spawns the background re-subscribe/forward task for a durable consumer and
/// returns the consumer-facing stream.
///
/// `initial` is the already-created message stream (so setup errors surface
/// synchronously at the call site); `make_stream` re-creates it after each
/// disconnect. Items are forwarded through an internal channel holding
/// `channel_size` items, and the task re-subscribes with exponential backoff
/// until the receiver is dropped.
///
/// Shared by [`Subscriber::subscribe`](arminius_lib_message_bus::Subscriber)
/// and [`Responder::respond`](arminius_lib_message_bus::Responder), which
/// differ only in the per-message stream `make_stream` produces.
pub fn resubscribing_stream<T, F, Fut>(
	signals: Arc<ConnectionSignals>,
	channel_size: usize,
	subject: &'static str,
	initial: MessageStream<T>,
	make_stream: F,
) -> impl Stream<Item = T> + Send + 'static
where
	T: Send + 'static,
	F: Fn() -> Fut + Send + Sync + 'static,
	Fut: Future<Output = Result<MessageStream<T>, Error>> + Send,
{
	let (tx, rx) = tokio::sync::mpsc::channel(channel_size);

	tokio::spawn(async move {
		forward_and_resubscribe(&signals, subject, initial, make_stream, tx).await;
	});

	futures::stream::unfold(rx, |mut rx| async move { rx.recv().await.map(|item| (item, rx)) })
}

#[cfg(test)]
mod tests;