celp_sdk/broker/

redis_broker.rs

use std::{os::unix::net::UnixStream, time::Duration};

use redis::{Client, Commands};
use thiserror::Error;

use crate::{
    cache::{Cache, CacheError},
    protobuf::{build_system_event_source, pack_system_event, system_event::Severity},
    warning,
};

use super::redis_pubsub::OwnedPubSub;
use super::{Message, Publish, Subscribe};

/*
 * Constants
 */

const BROKER_CACHE_KEY: &str = "celp:broker_cache";
const BROKER_ENDPOINT: &str = "/run/celp/redis.sock";
const BROKER_FALLBACK_ENDPOINT: &str = "redis://127.0.0.1/";

/*
 * Custom Errors
 */

#[derive(Error, Debug)]
pub enum BrokerError {
    #[error("cache error")]
    CacheError(#[from] CacheError),
    #[error("unable to encode protobuf")]
    EncodeError(#[from] prost::EncodeError),
    #[error(transparent)]
    RedisError(#[from] redis::RedisError),
}

/*
 * Trait implementations
 */

impl From<redis::Msg> for Message {
    fn from(redis_message: redis::Msg) -> Message {
        Message {
            data: redis_message.get_payload_bytes().to_vec(),
            topic: redis_message.get_channel_name().into(),
        }
    }
}

impl Publish for Broker {
    type Error = BrokerError;

    /// Function allowing to publish a protobuf message to the CELP broker
    /// # Arguments
    ///
    /// * `topic`   - The topic to publish the message to
    /// * `message` - The protobuf to be published
    ///
    /// # Examples
    /// ```rust,no_run
    /// use celp_sdk::{
    ///     protobuf::se::{AppInfo, app_info::SemVer},
    ///     broker::{Broker, Publish},
    /// };
    ///
    /// let app_info = AppInfo::new("test", SemVer::default());
    /// let mut broker = Broker::new().unwrap();
    /// if let Err(e) = broker.publish("app.info", &app_info) {
    ///     eprintln!("error publishing: {e:?}");
    /// }
    /// ```
    fn publish<M: prost::Message>(&mut self, topic: &str, message: &M) -> Result<(), Self::Error> {
        let buf = message.encode_to_vec();

        self.publish_raw(topic, &buf)
    }

    /// Function allowing to publish a protobuf message to the CELP broker while caching the latest version
    /// # Arguments
    ///
    /// * `topic`   - The topic to publish the message to
    /// * `message` - The protobuf to be published
    ///
    /// # Examples
    /// ```rust,no_run
    /// use celp_sdk::{
    ///     protobuf::se::{AppInfo, app_info::SemVer},
    ///     broker::{Broker, Publish},
    /// };
    ///
    /// let app_info = AppInfo::new("test", SemVer::default());
    /// let mut broker = Broker::new().unwrap();
    /// if let Err(e) = broker.publish_with_cache("app.info", &app_info) {
    ///     eprintln!("error publishing: {e:?}");
    /// }
    /// ```
    fn publish_with_cache<M: prost::Message>(
        &mut self,
        topic: &str,
        message: &M,
    ) -> Result<(), Self::Error> {
        let buf = message.encode_to_vec();

        self.cache.add_field(BROKER_CACHE_KEY, topic, &buf)?;
        self.publish_raw(topic, &buf)
    }

    /// Shorthand function for publishing a system event to the CELP broker while caching the latest version.
    /// The broker will perform the required wrapping of the system event
    /// # Arguments
    ///
    /// * `event_details`   - the system event details to publish
    /// * `severity`        - The severity with which to send the system event
    ///
    /// # Examples
    /// ```rust,no_run
    /// use celp_sdk::{
    ///     protobuf::{se::{AppInfo, app_info::SemVer}, system_event::Severity},
    ///     broker::{Broker, Publish},
    /// };
    ///
    /// let app_info = AppInfo::new("test", SemVer::default());
    /// let mut broker = Broker::new().unwrap();
    /// if let Err(e) = broker.publish_event(&app_info, Severity::Info) {
    ///     eprintln!("error publishing: {e:?}");
    /// }
    /// ```
    fn publish_event<E: prost::Message + prost::Name>(
        &mut self,
        event_details: &E,
        severity: Severity,
    ) -> Result<(), Self::Error> {
        let event = pack_system_event(event_details, severity)?;

        self.publish_with_cache(&event.source, &event)
    }

    /// Function allowing to publish a byte array to the CELP broker
    /// # Arguments
    ///
    /// * `topic`   - The topic to publish the message to
    /// * `message` - The protobuf to be published
    ///
    /// # Examples
    /// ```rust,no_run
    /// use prost::Message;
    /// use celp_sdk::{
    ///     protobuf::se::{AppInfo, app_info::SemVer},
    ///     broker::{Broker, Publish},
    /// };
    ///
    /// let app_info = AppInfo::new("test", SemVer::default());
    /// let app_info_data = app_info.encode_to_vec();
    /// let mut broker = Broker::new().unwrap();
    /// if let Err(e) = broker.publish_raw("app.info", &app_info_data) {
    ///     eprintln!("error publishing: {e:?}");
    /// }
    /// ```
    fn publish_raw(&mut self, topic: &str, data: &[u8]) -> Result<(), Self::Error> {
        Ok(self.connection.publish(topic, data)?)
    }
}

impl Subscribe for Broker {
    type Error = BrokerError;

    /// Function allowing subscription to the CELP broker for a specific topic
    /// # Arguments
    ///
    /// * `topic`   - The topic to subscribe to
    ///
    /// # Examples
    /// ```rust,no_run
    /// use celp_sdk::{
    ///     broker::{Broker, Subscribe},
    /// };
    ///
    /// let mut broker = Broker::new().unwrap();
    /// if let Err(e) = broker.subscribe("app.info") {
    ///     eprintln!("error subscribing: {e:?}");
    /// }
    /// ```
    fn subscribe(&mut self, topic: &str) -> Result<Option<Message>, Self::Error> {
        self.pubsub.subscribe(topic)?;

        let message = match self.cache.get_field_value(BROKER_CACHE_KEY, topic) {
            Ok(Some(val)) => Some(Message {
                data: val.into_bytes(),
                topic: topic.into(),
            }),
            Ok(_) => None,
            Err(e) => {
                warning!("Unable to retrieve cached data for {}: {e:?}", topic);
                None
            }
        };

        Ok(message)
    }

    /// Function allowing subscription to the CELP broker for a pattern
    /// # Arguments
    ///
    /// * `pattern` - The pattern to subscribe to
    ///
    /// # Examples
    /// ```rust,no_run
    /// use celp_sdk::{
    ///     broker::{Broker, Subscribe},
    /// };
    ///
    /// let mut broker = Broker::new().unwrap();
    /// if let Err(e) = broker.subscribe_pattern("app.info.*") {
    ///     eprintln!("error subscribing: {e:?}");
    /// }
    /// ```
    fn subscribe_pattern(&mut self, pattern: &str) -> Result<(), Self::Error> {
        Ok(self.pubsub.psubscribe(pattern)?)
    }

    /// Function allowing subscription to the CELP broker for a specific system event
    /// # Examples
    /// ```rust,no_run
    /// use celp_sdk::{
    ///     protobuf::se::AppInfo,
    ///     broker::{Broker, Subscribe},
    /// };
    ///
    /// let mut broker = Broker::new().unwrap();
    /// if let Err(e) = broker.subscribe_event::<AppInfo>() {
    ///     eprintln!("error subscribing: {e:?}");
    /// }
    /// ```
    fn subscribe_event<E: prost::Message + prost::Name>(
        &mut self,
    ) -> Result<Option<Message>, Self::Error> {
        self.subscribe(&build_system_event_source::<E>())
    }

    /// Function allowing unsubscription from the CELP broker for a specific topic
    /// # Arguments
    ///
    /// * `topic`   - The topic to unsubscribe for
    ///
    /// # Examples
    /// ```rust,no_run
    /// use celp_sdk::{
    ///     broker::{Broker, Subscribe},
    /// };
    ///
    /// let mut broker = Broker::new().unwrap();
    /// if let Err(e) = broker.unsubscribe("app.info") {
    ///     eprintln!("error unsubscribing: {e:?}");
    /// }
    /// ```
    fn unsubscribe(&mut self, topic: &str) -> Result<(), Self::Error> {
        Ok(self.pubsub.unsubscribe(topic)?)
    }

    /// Function allowing unsubscription from the CELP broker for a pattern
    /// # Arguments
    ///
    /// * `pattern` - The pattern to unsubscribe for
    ///
    /// # Examples
    /// ```rust,no_run
    /// use celp_sdk::{
    ///     broker::{Broker, Subscribe},
    /// };
    ///
    /// let mut broker = Broker::new().unwrap();
    /// if let Err(e) = broker.unsubscribe_pattern("app.info.*") {
    ///     eprintln!("error unsubscribing: {e:?}");
    /// }
    /// ```
    fn unsubscribe_pattern(&mut self, pattern: &str) -> Result<(), Self::Error> {
        Ok(self.pubsub.punsubscribe(pattern)?)
    }

    /// Function allowing unsubscription from the CELP broker for a specific system event
    /// # Examples
    /// ```rust,no_run
    /// use celp_sdk::{
    ///     protobuf::se::AppInfo,
    ///     broker::{Broker, Subscribe},
    /// };
    ///
    /// let mut broker = Broker::new().unwrap();
    /// if let Err(e) = broker.unsubscribe_event::<AppInfo>() {
    ///     eprintln!("error unsubscribing: {e:?}");
    /// }
    /// ```
    fn unsubscribe_event<E: prost::Message + prost::Name>(&mut self) -> Result<(), Self::Error> {
        self.unsubscribe(&build_system_event_source::<E>())
    }

    /// Function allowing to read messages from the broker
    /// # Examples
    /// ```rust,no_run
    /// use std::time::Duration;
    /// use celp_sdk::{
    ///     protobuf::{se::AppInfo, build_system_event_source, unpack_system_event_raw},
    ///     broker::{Broker, Subscribe, BrokerError},
    /// };
    ///
    /// let mut broker = Broker::new().unwrap()
    ///     .set_read_timeout(Duration::from_millis(100))
    ///     .unwrap();
    /// broker.subscribe_event::<AppInfo>().unwrap();
    ///
    /// let app_info_topic = build_system_event_source::<AppInfo>();
    ///
    /// loop {
    ///     match broker.get_message() {
    ///         Ok(message) => {
    ///             if message.topic == app_info_topic {
    ///                 let app_info = unpack_system_event_raw::<AppInfo>(&message.data).unwrap();
    ///                 // Handle event here
    ///             }
    ///         }
    ///         Err(BrokerError::RedisError(e)) if e.is_timeout() => continue,
    ///         Err(e) => eprintln!("Unable to receive messages from broker: {e:?}"),
    ///     }
    /// }
    /// ```
    fn get_message(&mut self) -> Result<super::interface::Message, Self::Error> {
        Ok(self.pubsub.get_message()?.into())
    }
}

/*
 * Broker implementation
 */

pub struct Broker {
    cache: Cache,
    connection: redis::Connection,
    pubsub: OwnedPubSub,
}

impl Broker {
    /// Create a new instance of broker
    /// # Examples
    /// ```rust,no_run
    /// use celp_sdk::broker::Broker;
    ///
    /// let broker = Broker::new().unwrap();
    /// ```
    pub fn new() -> Result<Self, BrokerError> {
        let client = Self::get_client()?;
        let connection = client.get_connection()?;
        let sub_connection = client.get_connection()?;

        Ok(Self {
            cache: Cache::new()?,
            connection,
            pubsub: OwnedPubSub::new(sub_connection),
        })
    }

    /// Set the read timeout, i.e. the amount of time the get_message call will block before returning when no
    /// messages can be read
    /// # Examples
    /// ```rust,no_run
    /// use std::time::Duration;
    /// use celp_sdk::broker::Broker;
    ///
    /// let broker = Broker::new().unwrap()
    ///     .set_read_timeout(Duration::from_millis(100))
    ///     .unwrap();
    /// ```
    pub fn set_read_timeout(self, dur: Duration) -> Result<Self, BrokerError> {
        self.pubsub.set_read_timeout(Some(dur))?;

        Ok(self)
    }

    fn get_client() -> Result<redis::Client, BrokerError> {
        let client = match UnixStream::connect(BROKER_ENDPOINT) {
            Ok(_) => Client::open(format!("unix:{}", BROKER_ENDPOINT))?,
            Err(e) => {
                warning!(
                    "Unable to open cache connection to {}: {e:?}. Trying fallback {}",
                    BROKER_ENDPOINT,
                    BROKER_FALLBACK_ENDPOINT
                );
                Client::open(BROKER_FALLBACK_ENDPOINT)?
            }
        };

        Ok(client)
    }
}