refactor: one channel per consumer

This commit is contained in:
izzy
2025-06-28 18:26:51 +01:00
parent 9d576d2430
commit 083229c7f5
6 changed files with 65 additions and 42 deletions

View File

@@ -108,5 +108,5 @@ pub async fn client_subscriber(
}
}
consumer.dispose().await;
consumer.dispose_channel().await;
}

View File

@@ -5,18 +5,21 @@ use std::{
use async_std::stream::StreamExt;
use lapin::{
Channel,
Channel, Connection,
options::BasicAckOptions,
types::{AMQPValue, FieldArray, FieldTable, LongLongInt},
};
use log::info;
use rand::Rng;
use revolt_config::config;
use revolt_config::{capture_internal_error, config};
use serde::de::DeserializeOwned;
use crate::event_stream::get_channel;
use crate::event_stream::{create_channel, get_connection};
pub struct Consumer {
channel: Arc<Channel>,
#[allow(dead_code)]
conn: Arc<Connection>,
channel: Channel,
tag: String,
topics: HashSet<String>,
topics_changed: bool,
@@ -27,8 +30,13 @@ pub struct Consumer {
impl Consumer {
/// Create a new event stream consumer
pub async fn new() -> Consumer {
let config = config().await;
let conn = get_connection().await;
let channel = create_channel(&conn, config.rabbit.event_stream).await;
Consumer {
channel: get_channel().await,
conn,
channel,
tag: rand::rng()
.sample_iter::<char, _>(&rand::distr::StandardUniform)
.take(32)
@@ -49,10 +57,13 @@ impl Consumer {
/// Get the current consumer
pub async fn ensure_consumer(&mut self) {
if self.topics_changed {
self.consumer = None;
info!("Topics changed, disposing the consumer.");
self.dispose_consumer().await;
self.topics_changed = false;
}
if self.consumer.is_none() {
info!("Creating a new consumer, tag={}", self.tag);
let config = config().await;
// Build arguments for consumer
@@ -88,27 +99,36 @@ impl Consumer {
}
}
/// Close the active consumer
pub async fn dispose(&mut self) {
// Close the channel -- don't do this actually
// capture_internal_error!(self.channel.close(0, "closing channel").await);
/// Close the active consumer if one exists
pub async fn dispose_consumer(&mut self) {
if let Some(consumer) = self.consumer.as_ref() {
if consumer.state().is_active() {
if let Err(err) = self
.channel
.basic_cancel(&self.tag, Default::default())
.await
{
eprintln!("Failed to close consumer! {:?}", err);
}
// Close the consumer
if let Err(err) = self
.channel
.basic_cancel(&self.tag, Default::default())
.await
{
eprintln!("Failed to close consumer! {:?}", err);
// is this necessary?
// else {
// Read the consumer to the end
// while let Some(delivery) = consumer.next().await {
// let delivery = delivery.expect("error in consumer");
// delivery.ack(BasicAckOptions::default()).await.expect("ack");
// }
// }
}
self.consumer = None;
}
// is this necessary?
// else {
// Read the consumer to the end
// while let Some(delivery) = consumer.next().await {
// let delivery = delivery.expect("error in consumer");
// delivery.ack(BasicAckOptions::default()).await.expect("ack");
// }
// }
}
/// Close the active channel
pub async fn dispose_channel(&mut self) {
// Close the channel -- don't do this actually
capture_internal_error!(self.channel.close(0, "closing channel").await);
}
/// Get the next item

View File

@@ -3,5 +3,5 @@ mod pool;
mod publish;
pub use consumer::Consumer;
pub use pool::get_channel;
pub use pool::{create_channel, get_connection};
pub use publish::publish_event;

View File

@@ -1,6 +1,6 @@
use async_std::sync::Mutex;
use lapin::{
Channel,
Connection,
options::QueueDeclareOptions,
types::{AMQPValue, FieldTable},
};
@@ -11,10 +11,10 @@ use std::sync::Arc;
use crate::create_client;
/// Get a handle to the event stream
pub async fn get_channel() -> Arc<Channel> {
pub async fn get_connection() -> Arc<Connection> {
let config = config().await;
static CONNECTIONS: Mutex<Vec<Arc<lapin::Channel>>> = Mutex::new(Vec::new());
static CONNECTIONS: Mutex<Vec<Arc<Connection>>> = Mutex::new(Vec::new());
let mut connections = CONNECTIONS.lock().await;
connections.retain(|item| {
@@ -39,22 +39,21 @@ pub async fn get_channel() -> Arc<Channel> {
.collect::<Vec<usize>>()
);
for channel in connections.iter() {
if Arc::strong_count(channel) < config.rabbit.event_stream.channels_per_conn {
return channel.clone();
for conn in connections.iter() {
if Arc::strong_count(conn) < config.rabbit.event_stream.channels_per_conn {
return conn.clone();
}
}
let conn = create_client().await;
let channel = Arc::new(create_channel(conn, config.rabbit.event_stream).await);
let conn = Arc::new(create_client().await);
connections.push(channel.clone());
channel
connections.push(conn.clone());
conn
}
/// Create a channel
async fn create_channel(
conn: lapin::Connection,
pub async fn create_channel(
conn: &lapin::Connection,
event_stream: RabbitEventStream,
) -> lapin::Channel {
let channel = conn.create_channel().await.unwrap();

View File

@@ -7,7 +7,7 @@ use lapin::{
use revolt_config::config;
use serde::Serialize;
use crate::event_stream::get_channel;
use crate::event_stream::{create_channel, get_connection};
/// Publish an event to the message broker
pub async fn publish_event<T: Serialize>(
@@ -22,7 +22,9 @@ pub async fn publish_event<T: Serialize>(
AMQPValue::LongString(channel.into()),
);
get_channel()
let conn = get_connection().await;
create_channel(&conn, config.rabbit.event_stream.clone())
.await
.basic_publish(
&config.rabbit.event_stream.exchange,

View File

@@ -40,7 +40,9 @@ impl TestHarness {
.collect();
static QUEUE_NAME: &str = "revolt.events";
let consumer = event_stream::get_channel()
let conn = event_stream::get_connection().await;
let consumer = event_stream::create_channel(&conn, config.rabbit.event_stream)
.await
.basic_consume(QUEUE_NAME, &tag, Default::default(), Default::default())
.await