resilient_rabbitmq_client/

client_handler.rs

use amqprs::callbacks::ConnectionCallback;
use amqprs::channel::BasicConsumeArguments;
use amqprs::channel::Channel;
use amqprs::connection::{Connection, OpenConnectionArguments};
use amqprs::consumer::AsyncConsumer;
use amqprs::error::Error;
use amqprs::{AmqpChannelId, Close};
use async_trait::async_trait;
use futures::lock::Mutex;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::mpsc::{self, Receiver, Sender};
use tokio::time::{Duration, sleep};
#[allow(unused)]
use tracing::{debug, error, info, trace, warn};

use crate::consumer::ChannelConsumerHandler;
use crate::publisher::DataToPublish;
use crate::publisher::PublishChannel;

pub type ConnectionArguments = OpenConnectionArguments;

/// This struct is holding the state of the RabbitMQ client. The whole resiliency revolves around it.
/// It will handle connection errors and backup the list of consumers (basic.consume) you declare and recreate them when something happens.
/// If you need to publish to RabbitMQ as well, a PublishChannel will be extracted from it and regenerated if needed.
pub struct RabbitMqClientHandler<BasicConsumerType: AsyncConsumer + Clone + Send + 'static> {
    // Globally useful part
    need_to_reload_sender: Sender<NeedToReload>,
    need_to_reload_receiver: Receiver<NeedToReload>,
    open_connection_arguments: OpenConnectionArguments,
    new_valid_channel_sender: Option<Sender<()>>,

    // Consumer specific part
    with_consumer: bool,
    consumer_connection: Connection,
    expected_consumers: Vec<(BasicConsumerType, BasicConsumeArguments)>,
    consume_channels: HashMap<AmqpChannelId, (Channel, BasicConsumerType, BasicConsumeArguments)>,

    // Publisher specific part
    with_publisher: bool,
    publisher_connection: Option<Connection>,
    publish_channel: Option<Arc<Mutex<Channel>>>,
    data_to_publish_sender: Option<Sender<DataToPublish>>
}

impl<BasicConsumerType: AsyncConsumer + Clone + Send + 'static>
    RabbitMqClientHandler<BasicConsumerType>
{
    pub async fn new(
        open_connection_arguments: OpenConnectionArguments,
    ) -> RabbitMqClientHandler<BasicConsumerType> {
        loop {
            match Connection::open(&open_connection_arguments).await {
                Ok(consumer_connection) => {
                    let (need_to_reload_sender, need_to_reload_receiver) =
                        mpsc::channel::<NeedToReload>(100);

                    let connection_callbacks =
                        ConnectionHandler::from(need_to_reload_sender.clone());

                    match consumer_connection.register_callback(connection_callbacks).await {
                        Ok(_) => {
                            debug!(
                                connection_name = consumer_connection.connection_name(),
                                "RabbitMqClientHandler creation successful",
                            );

                            return RabbitMqClientHandler {
                                // Global
                                need_to_reload_sender,
                                need_to_reload_receiver,
                                open_connection_arguments,
                                new_valid_channel_sender: None,
                                // Consume
                                with_consumer: true,
                                consumer_connection,
                                expected_consumers: Vec::new(),
                                consume_channels: HashMap::new(),
                                // Publish
                                with_publisher: false,
                                publisher_connection: None,
                                publish_channel: None,
                                data_to_publish_sender: None,
                            };
                        }
                        Err(error_content) => {
                            warn!("Failure to register callbacks : {:?}", error_content);
                        }
                    }
                }
                Err(error_content) => {
                    warn!(
                        "Failure to create RabbitMqClientHandler : {:?}",
                        error_content
                    );
                }
            }
            sleep(Duration::from_secs(2)).await;
        }
    }

    pub async fn try_reconnect(&mut self) -> Result<(), String> {
        match Connection::open(&self.open_connection_arguments).await {
            Ok(new_consumer_connection) => {
                let connection_callbacks =
                    ConnectionHandler::from(self.need_to_reload_sender.clone());

                match new_consumer_connection
                    .register_callback(connection_callbacks)
                    .await
                {
                    Ok(_) => {
                        self.consumer_connection = new_consumer_connection;

                        debug!(
                            connection_name = self.consumer_connection.connection_name(),
                            "RabbitMqClientHandler creation successful"
                        );

                        Ok(())
                    }
                    Err(error_content) => Err(format!(
                        "Failure to register callbacks : {:?}",
                        error_content
                    )),
                }
            }
            Err(error_content) => Err(format!(
                "Failure to create RabbitMqClientHandler : {:?}",
                error_content
            )),
        }
    }

    pub async fn create_publisher(&mut self) -> Result<PublishChannel, String> {
        debug!(
            connection_name = self.consumer_connection.connection_name(),
            "Beginning to create a publisher based on RabbitMqClientHandler",
        );

        match self.consumer_connection.open_channel(None).await {
            Ok(publish_channel) => {
                let (data_to_publish_sender, data_to_publish_receiver) =
                    mpsc::channel::<DataToPublish>(100);
                let (new_valid_channel_sender, new_valid_channel_receiver) =
                    mpsc::channel::<()>(10);

                let publish_channel_handle = Arc::new(Mutex::new(publish_channel));

                debug!("Publisher created and turned into Arc<Mutex>");
                self.with_publisher = true;
                self.publish_channel = Some(publish_channel_handle.clone());
                self.data_to_publish_sender = Some(data_to_publish_sender);
                self.new_valid_channel_sender = Some(new_valid_channel_sender);

                Ok(PublishChannel::from(
                    publish_channel_handle,
                    new_valid_channel_receiver,
                    self.need_to_reload_sender.clone(),
                    data_to_publish_receiver,
                ))
            }
            Err(error_content) => Err(format!("Failed to create publisher : {:?}", error_content)),
        }
    }

    pub async fn new_consumer(
        &mut self,
        basic_consumer: BasicConsumerType,
        basic_consumer_args: BasicConsumeArguments,
    ) {
        self.expected_consumers
            .push((basic_consumer.clone(), basic_consumer_args.clone()));

        loop {
            match self
                .create_consumer(basic_consumer.clone(), basic_consumer_args.clone())
                .await
            {
                Ok(_) => {
                    break;
                }
                Err(error_content) => {
                    warn!("Failed to regenerate consumer : {:?}", error_content);
                    sleep(Duration::from_secs(2)).await;
                }
            }
        }
    }

    pub async fn create_consumer(
        &mut self,
        basic_consumer: BasicConsumerType,
        basic_consumer_args: BasicConsumeArguments,
    ) -> Result<(), Error> {
        debug!(
            connection_name = self.consumer_connection.connection_name(),
            "create_consumer : trying to create channel on connection"
        );
        if !self.consumer_connection.is_open() {
            return Err(Error::ConnectionCloseError(
                "Connection is closed.".to_string(),
            ));
        }

        match self.consumer_connection.open_channel(None).await {
            Ok(new_channel) => {
                debug!(channel_id = new_channel.channel_id(), "new channel created");

                match new_channel
                    .basic_consume(basic_consumer.clone(), basic_consumer_args.clone())
                    .await
                {
                    Ok(consumer_tag) => {
                        debug!(
                            channel_id = new_channel.channel_id(),
                            consumer_tag, "Consumer created"
                        );

                        match new_channel
                            .register_callback(ChannelConsumerHandler::from(
                                basic_consumer.clone(),
                                basic_consumer_args.clone(),
                                self.need_to_reload_sender.clone(),
                            ))
                            .await
                        {
                            Ok(_) => {
                                info!(
                                    channel_id = new_channel.channel_id(),
                                    "Consumer created with callbacks : {:?}", basic_consumer_args
                                );

                                self.consume_channels.insert(
                                    new_channel.channel_id(),
                                    (new_channel, basic_consumer, basic_consumer_args),
                                );

                                Ok(())
                            }
                            Err(error_content) => Err(error_content),
                        }
                    }
                    Err(error_content) => Err(error_content),
                }
            }
            Err(error_content) => Err(error_content),
        }
    }

    pub fn get_publish_sender_clone(&self) -> Result<Sender<DataToPublish>, String> {
        match &self.data_to_publish_sender {
            Some(sender) => { Ok(sender.clone()) }
            None => { Err("Called get_publish_sender_clone() on this RabbitMqClientHandler but no publish channel is initialized".to_string()) }
        }
    }

    pub async fn keep_alive(&mut self) {
        while let Some(need_to_reload_details) = self.need_to_reload_receiver.recv().await {
            match need_to_reload_details {
                NeedToReload::All => {
                    debug!("Need to reload whole connection and channels");
                    debug!(
                        "Old channels list : {}",
                        display_channels_list(&self.consume_channels)
                    );
                    self.consume_channels.clear();
                    'outer: loop {
                        match self.try_reconnect().await {
                            Ok(_) => {
                                for (basic_consumer, basic_consumer_args) in
                                    self.expected_consumers.clone()
                                {
                                    loop {
                                        match self
                                            .create_consumer(
                                                basic_consumer.clone(),
                                                basic_consumer_args.clone(),
                                            )
                                            .await
                                        {
                                            Ok(_) => {
                                                break;
                                            }
                                            Err(Error::InternalChannelError(error_details)) => {
                                                warn!(
                                                    "Failure to create consumer {:?} : {}",
                                                    basic_consumer_args, error_details
                                                );
                                                sleep(Duration::from_secs(2)).await;
                                                continue;
                                            }
                                            Err(error_content) => {
                                                debug!(
                                                    "Failure to create consumer : {:?}",
                                                    error_content
                                                );
                                                continue 'outer;
                                            }
                                        }
                                    }
                                }

                                info!(
                                    "Reload connection/channels/consumer ok: {}",
                                    self.consumer_connection.connection_name()
                                );
                                if let Err(error_detail) = self
                                    .need_to_reload_sender
                                    .send(NeedToReload::PublishChannel)
                                    .await
                                {
                                    error!(
                                        "Failed to send NeedToReload signal over channel : {:?}",
                                        error_detail
                                    )
                                }

                                break;
                            }
                            Err(error_content) => {
                                warn!(
                                    "Failed to reload connection/channel/consumer : {:?}",
                                    error_content
                                );
                                sleep(Duration::from_secs(2)).await;
                            }
                        }
                    }
                }
                NeedToReload::ConsumerChannel(channel_id) => {
                    debug!("Need to reload channel {}", channel_id);
                    debug!(
                        "Old channels list : {}",
                        display_channels_list(&self.consume_channels)
                    );

                    let (basic_consumer, basic_consumer_args) = match self
                        .consume_channels
                        .get(&channel_id)
                    {
                        Some((
                            _problematic_channel,
                            old_basic_consumer,
                            old_basic_consumer_args,
                        )) => (old_basic_consumer.clone(), old_basic_consumer_args.clone()),
                        None => {
                            error!(
                                "NeedToReload::ConsumerChannel({}) received but missing from cache. Reload all.",
                                channel_id
                            );
                            if let Err(error_detail) =
                                self.need_to_reload_sender.send(NeedToReload::All).await
                            {
                                error!(
                                    "Failed to send NeedToReload signal over channel : {:?}",
                                    error_detail
                                );
                            }
                            continue;
                        }
                    };

                    self.consume_channels.remove(&channel_id);

                    loop {
                        match &mut self
                            .create_consumer(basic_consumer.clone(), basic_consumer_args.clone())
                            .await
                        {
                            Ok(_) => {
                                info!("Reload consumer ok: {:?}", basic_consumer_args);
                                break;
                            }
                            Err(Error::InternalChannelError(error_details)) => {
                                warn!(
                                    "Failure to create consumer {:?} : {}",
                                    basic_consumer_args, error_details
                                );
                                sleep(Duration::from_secs(2)).await;
                                continue;
                            }
                            Err(_) => {
                                debug!("Need to reload whole connection an channels");
                                debug!(
                                    "Old channels list : {}",
                                    display_channels_list(&self.consume_channels)
                                );
                                self.consume_channels.clear();
                                'outer: loop {
                                    match self.try_reconnect().await {
                                        Ok(_) => {
                                            for (basic_consumer, basic_consumer_args) in
                                                self.expected_consumers.clone()
                                            {
                                                loop {
                                                    match self
                                                        .create_consumer(
                                                            basic_consumer.clone(),
                                                            basic_consumer_args.clone(),
                                                        )
                                                        .await
                                                    {
                                                        Ok(_) => {
                                                            break;
                                                        }
                                                        Err(Error::InternalChannelError(
                                                            error_details,
                                                        )) => {
                                                            warn!(
                                                                "Failure to create consumer {:?} : {}",
                                                                basic_consumer_args, error_details
                                                            );
                                                            sleep(Duration::from_secs(2)).await;
                                                            continue;
                                                        }
                                                        Err(error_content) => {
                                                            debug!(
                                                                "Failure to create consumer : {:?}",
                                                                error_content
                                                            );
                                                            continue 'outer;
                                                        }
                                                    }
                                                }
                                            }

                                            info!(
                                                "Reload connection/channels/consumer ok: {}",
                                                self.consumer_connection.connection_name()
                                            );
                                            break;
                                        }
                                        Err(error_content) => {
                                            warn!(
                                                "Échec reload connection/channel/consumer : {:?}",
                                                error_content
                                            );
                                            sleep(Duration::from_secs(2)).await;
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
                NeedToReload::PublishChannel => {
                    if self.with_publisher {
                        loop {
                            match self.consumer_connection.open_channel(None).await {
                                Ok(new_publish_channel) => {
                                    info!(
                                        "Reload publisher ok on channel {}",
                                        new_publish_channel.channel_id()
                                    );

                                    let publish_channel_clone =
                                        self.publish_channel.clone().unwrap();
                                    let mut publish_channel_handle =
                                        publish_channel_clone.lock().await;
                                    *publish_channel_handle = new_publish_channel;

                                    self.new_valid_channel_sender
                                        .clone()
                                        .unwrap()
                                        .send(())
                                        .await
                                        .unwrap();

                                    break;
                                }
                                Err(Error::InternalChannelError(error_details)) => {
                                    warn!("Failure to create PublishChannel : {}", error_details);
                                    sleep(Duration::from_secs(2)).await;
                                    continue;
                                }
                                Err(error_content) => {
                                    debug!(
                                        "Failure to create PublishChannel : {:?}. Need to reload everything.",
                                        error_content
                                    );
                                    if let Err(error_detail) =
                                        self.need_to_reload_sender.send(NeedToReload::All).await
                                    {
                                        error!(
                                            "Failed to send NeedToReload signal over channel : {:?}",
                                            error_detail
                                        )
                                    }
                                }
                            }
                        }
                    }
                }
            }

            debug!(
                "New channels list : {}",
                display_channels_list(&self.consume_channels)
            );
        }
    }
}

/// The type which implements the ConnectionCallback trait from amqprs
#[derive(Clone)]
pub struct ConnectionHandler {
    need_to_reload_sender: Sender<NeedToReload>,
}

impl ConnectionHandler {
    pub fn from(need_to_reload_sender: Sender<NeedToReload>) -> ConnectionHandler {
        ConnectionHandler {
            need_to_reload_sender,
        }
    }
}

#[async_trait]
impl ConnectionCallback for ConnectionHandler {
    async fn close(&mut self, connection: &Connection, close: Close) -> Result<(), Error> {
        warn!(
            "End of connection received from broker : connection {}, cause: {}",
            connection, close
        );

        loop {
            match self.need_to_reload_sender.send(NeedToReload::All).await {
                Ok(_) => {
                    sleep(Duration::from_secs(1)).await;
                    break;
                }
                Err(error_content) => {
                    warn!(
                        "Failed to send need_new_connection MPSC signal : {:?}",
                        error_content
                    );
                    sleep(Duration::from_secs(5)).await;
                }
            }
        }

        Ok(())
    }

    async fn blocked(&mut self, connection: &Connection, reason: String) {
        warn!(
            "Blocked connection received from broker : connection {}, cause: {}",
            connection, reason
        );

        loop {
            match self.need_to_reload_sender.send(NeedToReload::All).await {
                Ok(_) => {
                    sleep(Duration::from_secs(1)).await;
                    break;
                }
                Err(error_content) => {
                    warn!(
                        "Failed to send need_new_connection MPSC signal : {:?}",
                        error_content
                    );
                    sleep(Duration::from_secs(5)).await;
                }
            }
        }
    }

    async fn unblocked(&mut self, connection: &Connection) {
        warn!(
            "Unblocked connection received from broker : connection {}",
            connection
        );
    }

    async fn secret_updated(&mut self, connection: &Connection) {
        warn!("secret_updated : connection {}", connection);

        loop {
            match self.need_to_reload_sender.send(NeedToReload::All).await {
                Ok(_) => {
                    sleep(Duration::from_secs(1)).await;
                    break;
                }
                Err(error_content) => {
                    warn!(
                        "Failed to send need_new_connection MPSC signal : {:?}",
                        error_content
                    );
                    sleep(Duration::from_secs(5)).await;
                }
            }
        }
    }
}

/// Convenient enum type for this crate's internals. This enum is being sent by the ChannelConsumerHandler instances to the ConnectionHandler whenever something bad is detected.
#[derive(Debug, Clone)]
pub enum NeedToReload {
    All,
    ConsumerChannel(AmqpChannelId),
    PublishChannel,
}

/// Convenient method for debugging purposes. When used in a DEBUG log level context, this crate will display channels lists at specific points (mainly when an item needs to be reloaded).
fn display_channels_list<BasicConsumerType: AsyncConsumer + Clone + Send + 'static>(
    channels: &HashMap<AmqpChannelId, (Channel, BasicConsumerType, BasicConsumeArguments)>,
) -> String {
    let mut list = String::new();
    for channel_id in channels.keys() {
        list.push_str(&format!("{} ", channel_id));
    }
    list
}