arminius_lib_nats/

consumergroup.rs

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

//! Consumer groups for NATS
//!
//! Allows creating Kafka consumer-group style sets of consumers for a JetStream
//! stream. This class assumes that it has a monopoly on consuming the stream,
//! i.e. that no other interface is used to consume it. Given a subject `subject`
//! a group of [`ConsumerClient`]s will consume all messages on the subject
//! `subject.[part].>` for every `part < partition_parts`. `partition_parts` is
//! defined in the [`ConsumerClientConfig`]. When a set of ConsumerClients are
//! consuming a stream, it is guaranteed that messages belonging to the same
//! part will be treated in order. But message consumption is not guaranteed to
//! be globally ordered, as more consumers running in parallel can outpace one
//! another. A group of ConsumerClients will try their best to have the same
//! parts allocated to the same ConsumerClient over time. However, when a new
//! ConsumerClient goes online or offline, all ConsumerClients will briefly stop
//! reading the event stream, and repartition the stream among them before they
//! start reading messages again, with no regard to the previous distribution of
//! parts in the partition.
//!
//! ## Example
//! This example connects to a stream, that may or may not exist, and prints
//! all messages received. The cool part is, you can run multiple instances of
//! this program, and they'll distribute the work among themselves!
//! ```no_run
//! # use async_nats::jetstream;
//! # use async_nats::jetstream::stream;
//! # use async_nats::jetstream::stream::SubjectTransform;
//! # use arminius_lib_nats::consumergroup;
//! # use futures::StreamExt;
//! # use std::process;
//!
//! #[tokio::main]
//! async fn main() -> Result<(), async_nats::Error> {
//! 	let client = async_nats::connect("localhost:4242").await?;
//! 	let jetstream = jetstream::new(client.clone());
//! 	let stream = jetstream
//! 		.get_or_create_stream(stream::Config {
//! 			name: "eventstream".to_owned(),
//! 			subjects: vec!["events.>".to_owned()],
//! 			retention: stream::RetentionPolicy::WorkQueue,
//! 			subject_transform: Some(SubjectTransform {
//! 				source: "events.*.>".to_owned(),
//! 				destination: "events.{{partition(100, 1)}}.>".to_owned(),
//! 			}),
//! 			..Default::default()
//! 		})
//! 		.await?;
//!
//! 	let mut client = consumergroup::ConsumerClient::new(
//! 		consumergroup::ConsumerClientConfig {
//! 			consumer_name: format!("worker_{}", process::id()),
//! 			subject: "events".to_owned(),
//! 			num_parts: 100,
//! 			..Default::default()
//! 		},
//! 		client,
//! 		stream,
//! 	)
//! 	.await?;
//!
//! 	loop {
//! 		if let Some(message) = client.get_message().await? {
//! 			println!("Received message {:?}", message);
//! 		}
//! 	}
//! 	Ok(())
//! }
//! ```

mod tests;

use std::time::{Duration, Instant};

use async_nats::{
	Message, Subscriber,
	jetstream::{
		consumer::{self, Consumer},
		stream::{self, Stream},
	},
};
use futures::{FutureExt as _, StreamExt as _, TryStreamExt as _};
use rand::RngExt as _;
use tokio::{runtime::Handle, task::block_in_place};

use crate::client::room::ROOMSERVER_STREAM_PARTITION_PARTS;

/// Configuration for a [ConsumerClient].
#[derive(Debug, Clone)]
pub struct ConsumerClientConfig {
	/// Name of the consumer. Must be unique over the stream.
	///
	/// Defaults to a random number between `0` and `10000`.
	pub consumer_name: String,
	/// Subject to subscribe to.
	///
	/// Defaults to `"consumer_subject"`.
	pub subject: String,
	/// The number of parts the stream is partitioned into. This number should
	/// be consistent with how the stream was partitioned, e.g with its
	/// `subject_transform` option.
	///
	/// Defaults to `100`.
	pub num_parts: usize,
	/// How long the consumer should stay alive before JetStream kills it when
	/// inactive.
	///
	/// Defaults to `30 seconds`.
	pub inactive_threshold: Duration,
	/// How many times the client should try updating the consumer's
	/// `filter_subjects` before giving up. See module's documentation.
	///
	/// Defaults to `5`.
	pub repartition_attempts: usize,
	/// Time to wait between attempts to update the `filter_subjects`. See
	/// module's documentation.
	///
	/// Defaults to `1 second`.
	pub repartition_delay: Duration,
	/// Timeout after which message fetching fails. If this timeout expires we
	/// return `None`.
	///
	/// Defaults to `30 seconds`.
	pub fetch_timeout: Duration,
}

impl Default for ConsumerClientConfig {
	fn default() -> Self {
		ConsumerClientConfig {
			consumer_name: rand::rng().random_range(0..10000).to_string(),
			subject: "consumer_subject".to_owned(),
			num_parts: 100,
			inactive_threshold: Duration::from_secs(30),
			repartition_attempts: 30,
			repartition_delay: Duration::from_secs(1),
			fetch_timeout: Duration::from_secs(60),
		}
	}
}

impl ConsumerClientConfig {
	/// Creates a consumer config that subscribes to the room subject
	#[must_use]
	pub fn for_rooms(name: String) -> Self {
		Self {
			consumer_name: name,
			subject: "rooms".to_owned(),
			num_parts: ROOMSERVER_STREAM_PARTITION_PARTS,
			// A non-responsive consumer blocks a large part of the roomserver
			// stream (depending on how many consumers we have). So I prefer having this number
			// low, until we have performance reasons to increase it
			inactive_threshold: Duration::from_secs(10),
			..Default::default()
		}
	}
}

/// Updates a consumer to give it a subset of the parts in the stream partition,
/// compared to all other consumers currently registered in the stream. If the
/// consumer no longer exists, possibly as a result of the `inactive_threshold`,
/// the consumer is recreated.
#[tracing::instrument(skip(config, stream))]
async fn update_or_recreate_consumer(
	config: &ConsumerClientConfig,
	num_partitions: usize,
	stream: &Stream<stream::Info>,
) -> Result<Consumer<consumer::pull::Config>, async_nats::Error> {
	for _ in 0..config.repartition_attempts {
		// In case our consumer died, we simply recreate it
		let consumers: Vec<String> = stream.consumer_names().try_collect().await?;

		tracing::trace!(
			"{} trying to get filters with list of consumers: {:?}",
			config.consumer_name,
			consumers
		);

		let Some(worker_index) =
			consumers.iter().position(|consumer| consumer == &config.consumer_name)
		else {
			tracing::trace!("Recreating decayed consumer {}", config.consumer_name);
			_ = stream
				.create_consumer(consumer::pull::Config {
					name: Some(config.consumer_name.clone()),
					filter_subjects: vec![format!("placeholder_filter.{}", config.consumer_name)],
					inactive_threshold: config.inactive_threshold,
					max_ack_pending: 1,
					..Default::default()
				})
				.await?;
			continue;
		};

		let num_workers = consumers.len();

		let filter_subjects: Vec<String> = (0..num_partitions)
			.skip(worker_index)
			.step_by(num_workers)
			.map(|i| format!("{}.{}.>", config.subject, i))
			.collect();

		if let Ok(consumer) = stream
			.update_consumer(consumer::pull::Config {
				name: Some(config.consumer_name.clone()),
				// If we don't put anything, we subscribe to everything
				filter_subjects: filter_subjects.clone(),
				inactive_threshold: config.inactive_threshold,
				max_ack_pending: 1,
				..Default::default()
			})
			.await
		{
			tracing::trace!("{} successfully updated filter_subjects", config.consumer_name);

			return Ok(consumer);
		}

		tokio::time::sleep(config.repartition_delay).await;
	}

	// If we can't update our consumer in a distributed manner, we don't claim any.
	// This way, we don't disrupt any other consumers, since only one consumer can access any
	// stream.
	//
	// This also means that the streams that this subscriber should have been subscribed to will be
	// without a subscriber until it's recreated.
	Ok(stream
		.update_consumer(consumer::pull::Config {
			name: Some(config.consumer_name.clone()),
			filter_subjects: vec![format!("placeholder_filter.{}", config.consumer_name)],
			inactive_threshold: config.inactive_threshold,
			max_ack_pending: 1,
			..Default::default()
		})
		.await?)
}

/// Handle controlling a single consumer.
#[derive(Debug)]
pub struct ConsumerClient {
	/// Configuration used for setting up this consumer. See
	/// [ConsumerClientConfig] for more info.
	config: ConsumerClientConfig,
	/// The stream over which the consumer was created.
	stream: Stream<stream::Info>,
	/// The actual `JetStream` consumer.
	consumer: Consumer<consumer::pull::Config>,
	/// The list of **other** consumers we currently know about.
	consumer_list: Vec<String>,
	/// When we last updated the list of consumers on the stream.
	consumer_list_last_update: Instant,
	/// A subscription to the subject in which new workers are announced.
	new_worker_subscription: Subscriber,
	/// A subscription to the subject in which deleted consumers are announced.
	consumer_deleted_subscription: Subscriber,
}

/// Tries deleting the [ConsumerClient] from the NATS JetStream.
impl Drop for ConsumerClient {
	fn drop(&mut self) {
		// TODO: Replace with AsyncDrop when stabilized
		// Tracking issue: https://github.com/rust-lang/rust/issues/126482
		block_in_place(move || {
			Handle::current().block_on(async move {
				if let Err(err) = self.stream.delete_consumer(&self.config.consumer_name).await {
					tracing::warn!("Failed to delete consumer: {err}");
				}
			});
		});

		tracing::info!("Finished deleting consumer");
	}
}

impl ConsumerClient {
	/// Creates a new [ConsumerClient] over a stream.
	pub async fn new(
		config: ConsumerClientConfig,
		client: async_nats::Client,
		stream: Stream<stream::Info>,
	) -> Result<Self, async_nats::Error> {
		let stream_name = &stream.cached_info().config.name;
		let new_worker_subscription = client
			.subscribe(format!("$JS.EVENT.ADVISORY.CONSUMER.CREATED.{stream_name}.*"))
			.await?;
		let consumer_deleted_subscription = client
			.subscribe(format!("$JS.EVENT.ADVISORY.CONSUMER.DELETED.{stream_name}.*"))
			.await?;

		let consumer = stream
			.create_consumer_strict(consumer::pull::Config {
				name: Some(config.consumer_name.clone()),
				filter_subjects: vec![format!("placeholder_filter.{}", config.consumer_name)],
				inactive_threshold: config.inactive_threshold,
				max_ack_pending: 1,
				..Default::default()
			})
			.await?;

		Ok(ConsumerClient {
			config,
			stream,
			consumer,
			consumer_list: Vec::new(),
			consumer_list_last_update: Instant::now() - Duration::from_secs(90),
			new_worker_subscription,
			consumer_deleted_subscription,
		})
	}

	/// Updates the list of **other** consumers we know are connected on the
	/// stream.
	async fn update_consumer_list(&mut self) -> Result<(), async_nats::Error> {
		self.consumer_list = self.stream.consumer_names().try_collect().await?;
		self.consumer_list_last_update = Instant::now();

		Ok(())
	}

	/// Updates the list of subjects we are consuming.
	#[tracing::instrument(skip(self))]
	async fn update_filter_subjects(&mut self) -> Result<(), async_nats::Error> {
		tracing::trace!("{}: Rebuilding consumer", self.config.consumer_name);
		match self
			.stream
			.update_consumer(consumer::pull::Config {
				name: Some(self.config.consumer_name.clone()),
				filter_subjects: vec![format!("placeholder_filter.{}", self.config.consumer_name)],
				inactive_threshold: self.config.inactive_threshold,
				max_ack_pending: 1,
				..Default::default()
			})
			.await
		{
			Ok(_) => Ok(()),
			Err(err) => match err.kind() {
				// We want to ignore this error
				stream::ConsumerUpdateErrorKind::DoesNotExist => Ok(()),
				_ => Err(err),
			},
		}?;

		self.consumer =
			update_or_recreate_consumer(&self.config, self.config.num_parts, &self.stream).await?;

		self.update_consumer_list().await?;

		Ok(())
	}

	/// Handles a notification for a new consumer arriving,
	/// possibly triggering a re-distribution of the stream's partition.
	#[tracing::instrument(skip(self))]
	async fn handle_new_worker_message(
		&mut self,
		msg: Option<Message>,
	) -> Result<(), async_nats::Error> {
		if let Some(msg) = msg {
			let pending_worker = msg.subject.rsplit(".").next().unwrap_or("UNKNOWN");
			if !self.consumer_list.iter().any(|w| w == pending_worker) {
				tracing::trace!("{}: New worker {}!", self.config.consumer_name, pending_worker);
				self.update_filter_subjects().await?;
			}
		}
		Ok(())
	}

	/// Handles a notification for a consumer on the stream going offline,
	/// possibly triggering a re-distribution of the stream partition.
	#[tracing::instrument(skip(self))]
	async fn handle_worker_deleted_message(
		&mut self,
		msg: Option<Message>,
	) -> Result<(), async_nats::Error> {
		if let Some(msg) = msg {
			// We can just ignore unknown workers going offline
			let offline_worker = msg.subject.rsplit(".").next().unwrap_or("UNKNOWN");
			tracing::trace!("{}: {} went offline!", self.config.consumer_name, offline_worker);
			if self.consumer_list.iter().any(|w| w == offline_worker) {
				self.update_filter_subjects().await?;
			}
		}
		Ok(())
	}

	/// Makes sure that the distribution of rooms between [ConsumerClient]s is
	/// correct, and then tries to fetch the next message
	///
	/// ## Returns
	///
	/// Returns the Message. If it returns `None`, this means
	/// it has reached `fetch_timeout` without a message being received
	#[tracing::instrument(skip(self))]
	pub async fn get_message(
		&mut self,
	) -> Result<Option<async_nats::jetstream::Message>, async_nats::Error> {
		// Update distribution of rooms until it stabilizes
		loop {
			if self.consumer_list_last_update < Instant::now() - Duration::from_secs(90) {
				tracing::trace!("Rebuilding consumer due to timeout");
				self.update_filter_subjects().await?;
				continue;
			}

			if let Some(msg) = self.new_worker_subscription.next().now_or_never() {
				self.handle_new_worker_message(msg).await?;
				continue;
			}
			if let Some(msg) = self.consumer_deleted_subscription.next().now_or_never() {
				self.handle_worker_deleted_message(msg).await?;
				continue;
			}

			break;
		}

		// And then we try to fetch a message
		let consumer_clone = self.consumer.clone();
		consumer_clone
			.fetch()
			.max_messages(1)
			.expires(self.config.fetch_timeout)
			.messages()
			.await?
			.try_next()
			.await
	}
}