From 083229c7f519a79d5e4003f6f69b69f9dde34b16 Mon Sep 17 00:00:00 2001 From: izzy Date: Sat, 28 Jun 2025 18:26:51 +0100 Subject: [PATCH] refactor: one channel per consumer --- crates/bonfire/src/client/subscriber.rs | 2 +- .../core/broker/src/event_stream/consumer.rs | 70 ++++++++++++------- crates/core/broker/src/event_stream/mod.rs | 2 +- crates/core/broker/src/event_stream/pool.rs | 23 +++--- .../core/broker/src/event_stream/publish.rs | 6 +- crates/delta/src/util/test.rs | 4 +- 6 files changed, 65 insertions(+), 42 deletions(-) diff --git a/crates/bonfire/src/client/subscriber.rs b/crates/bonfire/src/client/subscriber.rs index c0f13e3e..f3de104c 100644 --- a/crates/bonfire/src/client/subscriber.rs +++ b/crates/bonfire/src/client/subscriber.rs @@ -108,5 +108,5 @@ pub async fn client_subscriber( } } - consumer.dispose().await; + consumer.dispose_channel().await; } diff --git a/crates/core/broker/src/event_stream/consumer.rs b/crates/core/broker/src/event_stream/consumer.rs index 50126658..59d52ff1 100644 --- a/crates/core/broker/src/event_stream/consumer.rs +++ b/crates/core/broker/src/event_stream/consumer.rs @@ -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, + #[allow(dead_code)] + conn: Arc, + channel: Channel, tag: String, topics: HashSet, 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::(&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 diff --git a/crates/core/broker/src/event_stream/mod.rs b/crates/core/broker/src/event_stream/mod.rs index 7b9d4b08..e29cd0ea 100644 --- a/crates/core/broker/src/event_stream/mod.rs +++ b/crates/core/broker/src/event_stream/mod.rs @@ -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; diff --git a/crates/core/broker/src/event_stream/pool.rs b/crates/core/broker/src/event_stream/pool.rs index ff706c19..6870d7b6 100644 --- a/crates/core/broker/src/event_stream/pool.rs +++ b/crates/core/broker/src/event_stream/pool.rs @@ -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 { +pub async fn get_connection() -> Arc { let config = config().await; - static CONNECTIONS: Mutex>> = Mutex::new(Vec::new()); + static CONNECTIONS: Mutex>> = Mutex::new(Vec::new()); let mut connections = CONNECTIONS.lock().await; connections.retain(|item| { @@ -39,22 +39,21 @@ pub async fn get_channel() -> Arc { .collect::>() ); - 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(); diff --git a/crates/core/broker/src/event_stream/publish.rs b/crates/core/broker/src/event_stream/publish.rs index 4be20efa..668d6b40 100644 --- a/crates/core/broker/src/event_stream/publish.rs +++ b/crates/core/broker/src/event_stream/publish.rs @@ -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( @@ -22,7 +22,9 @@ pub async fn publish_event( 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, diff --git a/crates/delta/src/util/test.rs b/crates/delta/src/util/test.rs index de47ecb6..88a18f8e 100644 --- a/crates/delta/src/util/test.rs +++ b/crates/delta/src/util/test.rs @@ -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