yeti_types/

pubsub.rs

//! PubSub — topic-based publish/subscribe system.
//!
//! Provides real-time messaging infrastructure for subscriptions.
//!
//! # Async surface
//!
//! Many `PubSubManager` methods are `async fn` despite the current
//! implementation being sync (DashMap + tokio broadcast channels). The
//! async keyword is part of the public API contract: callers `.await` it
//! everywhere, and we want the freedom to swap the in-process backend
//! for a remote one (NATS/Redis) without forcing a workspace-wide call-
//! site sweep. See `unused_async` note at module scope.
#![expect(
    clippy::unused_async,
    reason = "PubSubManager API is async-contract — callers .await; backend may swap to remote in future"
)]

use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering};

use tokio::sync::broadcast;

use crate::resource::{MessageType, SubscriptionMessage};

/// Default channel capacity for broadcast channels.
const DEFAULT_CHANNEL_CAPACITY: usize = 256;

/// Topic-based publish/subscribe manager.
///
/// Thread-safe manager for real-time subscriptions. Supports topic-based
/// subscriptions, automatic channel creation, and broadcast to all subscribers.
#[derive(Debug, Clone)]
pub struct PubSubManager {
    /// Topic -> broadcast sender
    topics: Arc<dashmap::DashMap<String, broadcast::Sender<SubscriptionMessage>>>,
    /// Default channel capacity for new topics
    capacity: usize,
    /// Counter for periodic topic cleanup
    cleanup_counter: Arc<AtomicU64>,
}

impl PubSubManager {
    /// Create a new `PubSubManager` with specified channel capacity.
    #[must_use]
    pub fn new(capacity: usize) -> Self {
        Self {
            topics: Arc::new(dashmap::DashMap::new()),
            capacity,
            cleanup_counter: Arc::new(AtomicU64::new(0)),
        }
    }

    /// Create a new `PubSubManager` with default capacity (256).
    #[must_use]
    pub fn with_default_capacity() -> Self {
        Self::new(DEFAULT_CHANNEL_CAPACITY)
    }

    /// Subscribe to a topic. Returns a receiver for messages.
    pub async fn subscribe(&self, topic: &str) -> broadcast::Receiver<SubscriptionMessage> {
        // Periodic cleanup every 100 subscriptions
        if self
            .cleanup_counter
            .fetch_add(1, AtomicOrdering::Relaxed)
            .is_multiple_of(100)
        {
            self.topics.retain(|_, sender| sender.receiver_count() > 0);
        }

        if let Some(sender) = self.topics.get(topic) {
            return sender.subscribe();
        }

        let entry = self
            .topics
            .entry(topic.to_owned())
            .or_insert_with(|| broadcast::channel(self.capacity).0);
        entry.value().subscribe()
    }

    /// Publish a message to a topic. Returns number of recipients.
    pub async fn publish(&self, topic: &str, message: SubscriptionMessage) -> usize {
        self.topics
            .get(topic)
            .map_or(0, |sender| sender.send(message).unwrap_or(0))
    }

    /// Notify subscribers of a record update (publishes to record and table topics).
    /// Short-circuits with zero allocation if no subscribers exist.
    pub async fn notify_update(&self, table: &str, id: &str, data: &serde_json::Value) {
        if self.topics.is_empty() {
            return;
        }
        let record_topic = format!("{table}/{id}");
        let has_record_sub = self.topics.contains_key(&record_topic);
        let has_table_sub = self.topics.contains_key(table);
        if !has_record_sub && !has_table_sub {
            return;
        }
        let message = SubscriptionMessage {
            message_type: MessageType::Update,
            data: data.clone(),
            id: Some(id.to_owned()),
            timestamp: chrono::Utc::now(),
        };
        if has_record_sub {
            self.publish(&record_topic, message.clone()).await;
        }
        if has_table_sub {
            self.publish(table, message).await;
        }
    }

    /// Notify subscribers of a record deletion.
    /// Short-circuits with zero allocation if no subscribers exist.
    pub async fn notify_delete(&self, table: &str, id: &str) {
        if self.topics.is_empty() {
            return;
        }
        let record_topic = format!("{table}/{id}");
        let has_record_sub = self.topics.contains_key(&record_topic);
        let has_table_sub = self.topics.contains_key(table);
        if !has_record_sub && !has_table_sub {
            return;
        }
        let message = SubscriptionMessage {
            message_type: MessageType::Delete,
            data: serde_json::json!({"id": id}),
            id: Some(id.to_owned()),
            timestamp: chrono::Utc::now(),
        };
        if has_record_sub {
            self.publish(&record_topic, message.clone()).await;
        }
        if has_table_sub {
            self.publish(table, message).await;
        }
    }

    /// Notify subscribers of a custom published message.
    /// Short-circuits if no subscribers exist for this topic.
    pub async fn notify_publish(&self, topic: &str, data: serde_json::Value) {
        if self.topics.is_empty() || !self.topics.contains_key(topic) {
            return;
        }
        let message = SubscriptionMessage {
            message_type: MessageType::Publish,
            data,
            id: None,
            timestamp: chrono::Utc::now(),
        };
        self.publish(topic, message).await;
    }

    /// Get the number of active topics.
    pub async fn topic_count(&self) -> usize {
        self.topics.len()
    }

    /// Whether any subscription topic is live. Cheap (`DashMap::is_empty`),
    /// sync — callers on a write hot path gate record decode/notify behind
    /// this so a write with no subscribers pays a single atomic-load check.
    #[must_use]
    pub fn has_topics(&self) -> bool {
        !self.topics.is_empty()
    }

    /// Get the number of subscribers for a topic.
    pub async fn subscriber_count(&self, topic: &str) -> usize {
        self.topics
            .get(topic)
            .map_or(0, |entry| entry.receiver_count())
    }

    /// Check if a topic exists.
    pub async fn has_topic(&self, topic: &str) -> bool {
        self.topics.contains_key(topic)
    }

    /// Remove a topic. Returns true if it existed.
    pub async fn remove_topic(&self, topic: &str) -> bool {
        self.topics.remove(topic).is_some()
    }

    /// Remove topics with zero active subscribers.
    pub async fn cleanup_empty_topics(&self) -> usize {
        let before = self.topics.len();
        self.topics.retain(|_, sender| sender.receiver_count() > 0);
        before - self.topics.len()
    }

    /// List all active topic names.
    pub async fn topics(&self) -> Vec<String> {
        self.topics.iter().map(|e| e.key().clone()).collect()
    }
}

impl Default for PubSubManager {
    fn default() -> Self {
        Self::with_default_capacity()
    }
}