hanabi-core 0.1.0-alpha.2

mod command;
mod error;
mod message;
mod metrics;

pub use command::*;
pub use error::*;
pub use message::*;

use std::collections::HashMap;
use std::sync::atomic::Ordering;
use tokio::sync::mpsc::error::TrySendError;
use tokio::sync::{mpsc, watch};
use tracing::{error, info};

const PUBSUB_ENGINE_BACKLOG_SIZE: usize = 1000;
const PUBSUB_ENGINE_METRICS_INTERVAL: std::time::Duration = std::time::Duration::from_secs(5);

#[derive(Clone)]
pub struct PubSubEngine {
    command_tx: mpsc::Sender<PubSubCommand>,
    shutdown_tx: watch::Sender<()>,
}

impl Default for PubSubEngine {
    fn default() -> Self {
        Self::new()
    }
}

impl PubSubEngine {
    pub fn new() -> Self {
        let (command_tx, command_rx) = mpsc::channel(PUBSUB_ENGINE_BACKLOG_SIZE);
        let (shutdown_tx, shutdown_rx) = watch::channel(());

        let pubsub_metrics = metrics::PubSubMetrics::new();

        Self::start_engine(command_rx, shutdown_rx.clone(), pubsub_metrics.clone());
        Self::start_metrics_watcher(shutdown_rx, pubsub_metrics);

        Self {
            command_tx,
            shutdown_tx,
        }
    }

    /// Try to send a [PubSubCommand] into the PubSub Engine.
    ///
    /// This command will return a unit on success and either a [PubSubError::CommandQueueFull]
    /// or [PubSubError::CommandQueueClosed] in the cases of either the backlog is full
    /// or Engine has been shut down.
    pub fn send(&self, command: PubSubCommand) -> Result<(), PubSubError> {
        self.command_tx
            .try_send(command)
            .inspect_err(|error| error!(?error, "Failed to send command."))
            .map_err(|error| match error {
                TrySendError::Full(_) => PubSubError::CommandQueueFull,
                TrySendError::Closed(_) => PubSubError::CommandQueueClosed,
            })
    }

    /// Shuts down the PubSub engine gracefully.
    pub fn shutdown(self) {
        let _ = self.shutdown_tx.send(());
    }

    /// Starts the PubSub Engine in a specialized thread.
    fn start_engine(
        mut command_rx: mpsc::Receiver<PubSubCommand>,
        mut shutdown_rx: watch::Receiver<()>,
        pubsub_metrics: metrics::PubSubMetrics,
    ) {
        tokio::spawn(async move {
            let mut consumers: HashMap<String, Vec<mpsc::Sender<Message>>> = HashMap::new();
            loop {
                tokio::select! {
                    _ = shutdown_rx.changed() => {
                        info!("[PubSub] Engine shutting down.");
                        break;
                    }
                    Some(command) = command_rx.recv() => match command {
                        PubSubCommand::Publish { topic, message } => {
                            pubsub_metrics.messages_processed().fetch_add(1, Ordering::Relaxed);
                            if let Some(subscribers) = consumers.get_mut(&topic) {
                                subscribers.retain(|subscriber| {
                                    let result = subscriber.try_send(message.clone());
                                    match result {
                                        Ok(_) => true,
                                        Err(TrySendError::Full(_)) => {
                                            info!(?topic, "[PubSub] Subscriber backlog is full. The message will be dropped.");
                                            true
                                        }
                                        Err(TrySendError::Closed(_)) => {
                                            info!(?topic, "[PubSub] Subscriber closed. Removing from the consumer list.");
                                            false
                                        }
                                    }
                                });
                            }
                        }
                        PubSubCommand::Subscribe { topic, subscriber } => {
                            consumers.entry(topic).or_default().push(subscriber);
                        }
                        PubSubCommand::Unsubscribe { topic, subscriber } => {
                            if let Some(channels) = consumers.get_mut(&topic) {
                                if let Some(index) = channels.iter().position(|s| s.same_channel(&subscriber)) {
                                    channels.swap_remove(index);
                                }
                            }
                        }
                    }
                }
            }
        });
    }

    /// Starts the Pub/Sub engine metrics watcher task
    ///
    /// This function can gracefully shut down by sharing the engine's
    /// shutdown_rx channel.
    ///
    /// It will report the messages processed per second every [PUBSUB_ENGINE_METRICS_INTERVAL].
    fn start_metrics_watcher(
        mut shutdown_rx: watch::Receiver<()>,
        pubsub_metrics: metrics::PubSubMetrics,
    ) {
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(PUBSUB_ENGINE_METRICS_INTERVAL);
            loop {
                tokio::select! {
                    _ = shutdown_rx.changed() => {
                        info!("[PubSub] Metrics watcher shutting down.");
                        break;
                    }
                    _ = interval.tick() => {
                        let count = pubsub_metrics.messages_processed().swap(0, Ordering::Relaxed);
                        let rate = count as f64 / PUBSUB_ENGINE_METRICS_INTERVAL.as_secs_f64();
                        info!(messages_per_second = ?rate, "[PubSub] Metrics")
                    }
                }
            }
        });
    }
}