use std::collections::HashMap;
use tracing::{debug, info};
use crate::backends::rabbitmq::client::RabbitMqClient;
use crate::backends::rabbitmq::consumer_group::{ConsumerGroup, ConsumerGroupConfig};
use crate::backends::rabbitmq::topology::RabbitMqTopologyDeclarer;
use crate::error::{Result, ShoveError};
use crate::handler::MessageHandler;
use crate::topic::Topic;
use crate::topology::TopologyDeclarer;
pub struct ConsumerGroupRegistry {
groups: HashMap<String, ConsumerGroup>,
client: RabbitMqClient,
}
impl ConsumerGroupRegistry {
pub fn new(client: RabbitMqClient) -> Self {
Self {
groups: HashMap::new(),
client,
}
}
pub async fn register<T, H>(
&mut self,
config: ConsumerGroupConfig,
handler_factory: impl Fn() -> H + Send + Sync + 'static,
) -> Result<()>
where
T: Topic + 'static,
H: MessageHandler<T> + Clone + 'static,
{
let topology = T::topology();
let name = topology.queue().to_string();
if self.groups.contains_key(&name) {
return Err(ShoveError::Topology(format!(
"consumer group '{name}' is already registered"
)));
}
let channel = self.client.create_channel().await?;
let declarer = RabbitMqTopologyDeclarer::new(channel);
declarer.declare(topology).await?;
info!(group = %name, "registering consumer group");
let group_token = self.client.shutdown_token().child_token();
let group = ConsumerGroup::new::<T, H>(
name.clone(),
name.clone(),
config,
self.client.clone(),
group_token,
handler_factory,
);
self.groups.insert(name, group);
Ok(())
}
pub fn start_all(&mut self) {
info!(count = self.groups.len(), "starting all consumer groups");
for group in self.groups.values_mut() {
group.start();
}
}
pub fn groups(&self) -> &HashMap<String, ConsumerGroup> {
&self.groups
}
pub fn groups_mut(&mut self) -> &mut HashMap<String, ConsumerGroup> {
&mut self.groups
}
pub async fn shutdown_all(&mut self) {
info!(
count = self.groups.len(),
"shutting down all consumer groups"
);
for group in self.groups.values_mut() {
group.shutdown().await;
}
debug!("all consumer groups shut down");
}
}