PiCloud/crates/orchestrator-core/src/realtime.rs

//! In-process `RealtimeBroadcaster` — the SSE fan-out registry (v1.1.6).
//!
//! Sibling of [`crate::inbox::InboxRegistry`], but multi-receiver and
//! repeated-event: a `Mutex<HashMap<(AppId, topic), broadcast::Sender>>`
//! over `tokio::sync::broadcast` instead of a oneshot map. The publish
//! side ([`PubsubServiceImpl`]) and the SSE subscribe side both hold one
//! shared `Arc<InProcessBroadcaster>`.
//!
//! Delivery is best-effort: each channel has a bounded buffer
//! (`PICLOUD_REALTIME_BROADCAST_CAPACITY`, default 64); a slow consumer
//! that falls behind sees the oldest events dropped (standard
//! `broadcast` lag semantics — the receiver gets `RecvError::Lagged`).
//! SSE's transport-layer auto-reconnect is the recovery path; there's no
//! server-side replay in v1.1.6.
//!
//! Channels are created lazily on first subscribe. A periodic GC task
//! ([`spawn_realtime_gc`]) drops senders whose receiver count has fallen
//! to zero so one-shot subscribers don't grow the map unboundedly.
//!
//! Cluster mode (v1.3+) swaps this for a Postgres `LISTEN/NOTIFY`-backed
//! resolver behind the same `RealtimeBroadcaster` trait.

use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use async_trait::async_trait;
use picloud_shared::{AppId, BroadcasterError, GroupId, RealtimeBroadcaster, RealtimeEvent};
use tokio::sync::broadcast;

/// Default per-channel broadcast buffer depth.
pub const DEFAULT_BROADCAST_CAPACITY: usize = 64;
const ENV_CAPACITY: &str = "PICLOUD_REALTIME_BROADCAST_CAPACITY";

/// Default cap on the number of live channels PER map (per-app and per-group).
/// A subscribe that would create a NEW channel past this is refused, so an
/// attacker naming unbounded distinct topics can't grow the maps until OOM.
pub const DEFAULT_MAX_CHANNELS: usize = 100_000;
const ENV_MAX_CHANNELS: &str = "PICLOUD_REALTIME_MAX_CHANNELS";

/// Default GC sweep interval for empty channels.
pub const DEFAULT_GC_INTERVAL_SECS: u64 = 60;

pub struct InProcessBroadcaster {
    inner: Mutex<HashMap<(AppId, String), broadcast::Sender<RealtimeEvent>>>,
    /// §11.6 D2: group SHARED topic channels, keyed by the OWNING group so every
    /// subtree app subscribing to the same shared topic shares one channel.
    groups: Mutex<HashMap<(GroupId, String), broadcast::Sender<RealtimeEvent>>>,
    capacity: usize,
    /// Max live channels per map (see [`DEFAULT_MAX_CHANNELS`]).
    max_channels: usize,
}

impl InProcessBroadcaster {
    #[must_use]
    pub fn new(capacity: usize) -> Self {
        Self {
            inner: Mutex::new(HashMap::new()),
            groups: Mutex::new(HashMap::new()),
            capacity: capacity.max(1),
            max_channels: DEFAULT_MAX_CHANNELS,
        }
    }

    /// Override the per-map channel cap (0 → treated as 1; never unbounded).
    #[must_use]
    pub fn with_max_channels(mut self, max_channels: usize) -> Self {
        self.max_channels = max_channels.max(1);
        self
    }

    /// Build from `PICLOUD_REALTIME_BROADCAST_CAPACITY` (default 64) and
    /// `PICLOUD_REALTIME_MAX_CHANNELS` (default 100_000).
    #[must_use]
    pub fn from_env() -> Self {
        let capacity = match std::env::var(ENV_CAPACITY) {
            Err(_) => DEFAULT_BROADCAST_CAPACITY,
            Ok(v) => match v.parse::<usize>() {
                Ok(n) if n > 0 => n,
                Ok(_) => {
                    tracing::warn!(env = ENV_CAPACITY, value = %v, "must be > 0; using default");
                    DEFAULT_BROADCAST_CAPACITY
                }
                Err(e) => {
                    tracing::warn!(env = ENV_CAPACITY, value = %v, error = %e, "invalid; using default");
                    DEFAULT_BROADCAST_CAPACITY
                }
            },
        };
        let max_channels = match std::env::var(ENV_MAX_CHANNELS) {
            Err(_) => DEFAULT_MAX_CHANNELS,
            Ok(v) => match v.parse::<usize>() {
                Ok(n) if n > 0 => n,
                Ok(_) => {
                    tracing::warn!(env = ENV_MAX_CHANNELS, value = %v, "must be > 0; using default");
                    DEFAULT_MAX_CHANNELS
                }
                Err(e) => {
                    tracing::warn!(env = ENV_MAX_CHANNELS, value = %v, error = %e, "invalid; using default");
                    DEFAULT_MAX_CHANNELS
                }
            },
        };
        Self::new(capacity).with_max_channels(max_channels)
    }

    /// Number of live channels in both maps (test/observability helper).
    #[must_use]
    pub fn channel_count(&self) -> usize {
        let app = self.inner.lock().map(|g| g.len()).unwrap_or(0);
        let grp = self.groups.lock().map(|g| g.len()).unwrap_or(0);
        app + grp
    }

    /// Drop senders with zero receivers across both maps. Returns how many were
    /// removed. Called periodically by [`spawn_realtime_gc`].
    pub fn gc(&self) -> usize {
        let mut removed = 0;
        if let Ok(mut g) = self.inner.lock() {
            let before = g.len();
            g.retain(|_, tx| tx.receiver_count() > 0);
            removed += before - g.len();
        }
        if let Ok(mut g) = self.groups.lock() {
            let before = g.len();
            g.retain(|_, tx| tx.receiver_count() > 0);
            removed += before - g.len();
        }
        removed
    }
}

#[async_trait]
impl RealtimeBroadcaster for InProcessBroadcaster {
    async fn subscribe(
        &self,
        app_id: AppId,
        topic: &str,
    ) -> Result<broadcast::Receiver<RealtimeEvent>, BroadcasterError> {
        let mut g = self
            .inner
            .lock()
            .map_err(|_| BroadcasterError::Unavailable("broadcaster map poisoned".into()))?;
        let key = (app_id, topic.to_string());
        // Refuse to open a NEW channel past the cap; subscribing to an existing
        // one (another receiver) is always allowed.
        if !g.contains_key(&key) && g.len() >= self.max_channels {
            return Err(BroadcasterError::Unavailable(
                "realtime channel limit reached".into(),
            ));
        }
        let tx = g
            .entry(key)
            .or_insert_with(|| broadcast::channel(self.capacity).0);
        Ok(tx.subscribe())
    }

    async fn publish(&self, app_id: AppId, topic: &str, event: RealtimeEvent) {
        let Ok(g) = self.inner.lock() else {
            return;
        };
        // Only fan out to an existing channel: a topic with no live
        // subscribers has no sender (publish never creates one). `send`
        // returns Err iff every receiver has dropped — a benign no-op.
        if let Some(tx) = g.get(&(app_id, topic.to_string())) {
            let _ = tx.send(event);
        }
    }

    async fn drop_topic(&self, app_id: AppId, topic: &str) {
        if let Ok(mut g) = self.inner.lock() {
            // Removing the sender closes the channel; existing receivers
            // observe `RecvError::Closed` and disconnect cleanly.
            g.remove(&(app_id, topic.to_string()));
        }
    }

    async fn subscribe_group(
        &self,
        group_id: GroupId,
        topic: &str,
    ) -> Result<broadcast::Receiver<RealtimeEvent>, BroadcasterError> {
        let mut g = self
            .groups
            .lock()
            .map_err(|_| BroadcasterError::Unavailable("group broadcaster map poisoned".into()))?;
        let key = (group_id, topic.to_string());
        if !g.contains_key(&key) && g.len() >= self.max_channels {
            return Err(BroadcasterError::Unavailable(
                "realtime channel limit reached".into(),
            ));
        }
        let tx = g
            .entry(key)
            .or_insert_with(|| broadcast::channel(self.capacity).0);
        Ok(tx.subscribe())
    }

    async fn publish_group(&self, group_id: GroupId, topic: &str, event: RealtimeEvent) {
        let Ok(g) = self.groups.lock() else {
            return;
        };
        if let Some(tx) = g.get(&(group_id, topic.to_string())) {
            let _ = tx.send(event);
        }
    }

    async fn drop_group_topic(&self, group_id: GroupId, topic: &str) {
        if let Ok(mut g) = self.groups.lock() {
            g.remove(&(group_id, topic.to_string()));
        }
    }
}

/// Spawn the background GC sweep that drops empty channels every
/// `interval_secs` (default [`DEFAULT_GC_INTERVAL_SECS`]). Spawned at
/// startup alongside the other housekeeping tasks.
pub fn spawn_realtime_gc(broadcaster: Arc<InProcessBroadcaster>, interval_secs: u64) {
    let period = Duration::from_secs(interval_secs.max(1));
    tokio::spawn(async move {
        let mut ticker = tokio::time::interval(period);
        ticker.tick().await; // skip the immediate first fire
        loop {
            ticker.tick().await;
            let removed = broadcaster.gc();
            if removed > 0 {
                tracing::debug!(removed, "realtime broadcaster GC dropped empty channels");
            }
        }
    });
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::Utc;
    use serde_json::json;

    fn event(topic: &str, n: i64) -> RealtimeEvent {
        RealtimeEvent {
            topic: topic.to_string(),
            message: json!({ "n": n }),
            published_at: Utc::now(),
        }
    }

    #[tokio::test]
    async fn multiple_subscribers_each_receive_each_event() {
        let b = InProcessBroadcaster::new(16);
        let app = AppId::new();
        let mut rx1 = b.subscribe(app, "chat").await.unwrap();
        let mut rx2 = b.subscribe(app, "chat").await.unwrap();

        b.publish(app, "chat", event("chat", 1)).await;
        b.publish(app, "chat", event("chat", 2)).await;

        for rx in [&mut rx1, &mut rx2] {
            assert_eq!(rx.recv().await.unwrap().message, json!({ "n": 1 }));
            assert_eq!(rx.recv().await.unwrap().message, json!({ "n": 2 }));
        }
    }

    #[tokio::test]
    async fn dropped_subscriber_does_not_leak_after_gc() {
        let b = InProcessBroadcaster::new(16);
        let app = AppId::new();
        let rx = b.subscribe(app, "t").await.unwrap();
        assert_eq!(b.channel_count(), 1);
        drop(rx);
        // GC reclaims the now-empty channel.
        assert_eq!(b.gc(), 1);
        assert_eq!(b.channel_count(), 0);
    }

    #[tokio::test]
    async fn drop_topic_disconnects_existing_subscribers() {
        let b = InProcessBroadcaster::new(16);
        let app = AppId::new();
        let mut rx = b.subscribe(app, "t").await.unwrap();
        b.drop_topic(app, "t").await;
        // Sender gone → receiver observes a closed channel.
        assert!(rx.recv().await.is_err());
        assert_eq!(b.channel_count(), 0);
    }

    #[tokio::test]
    async fn slow_consumer_loses_oldest_events() {
        // Capacity 2: a consumer that never drains sees the oldest
        // events dropped (broadcast Lagged semantics).
        let b = InProcessBroadcaster::new(2);
        let app = AppId::new();
        let mut rx = b.subscribe(app, "t").await.unwrap();
        for i in 0..5 {
            b.publish(app, "t", event("t", i)).await;
        }
        // First recv reports the lag rather than event 0.
        let first = rx.recv().await;
        assert!(
            matches!(first, Err(broadcast::error::RecvError::Lagged(_))),
            "expected Lagged, got {first:?}"
        );
        // Subsequent recvs return the most recent buffered events.
        let next = rx.recv().await.unwrap();
        assert_eq!(next.message, json!({ "n": 3 }));
    }

    #[tokio::test]
    async fn cross_app_isolation() {
        let b = InProcessBroadcaster::new(16);
        let app_a = AppId::new();
        let app_b = AppId::new();
        let mut rx_a = b.subscribe(app_a, "shared").await.unwrap();
        let mut rx_b = b.subscribe(app_b, "shared").await.unwrap();

        b.publish(app_a, "shared", event("shared", 1)).await;
        // App B's subscriber must not see app A's publish.
        assert_eq!(rx_a.recv().await.unwrap().message, json!({ "n": 1 }));
        assert!(rx_b.try_recv().is_err());
    }

    #[tokio::test]
    async fn publish_with_no_subscribers_is_noop() {
        let b = InProcessBroadcaster::new(16);
        let app = AppId::new();
        // No subscriber → no sender created → no panic, nothing fanned out.
        b.publish(app, "ghost", event("ghost", 1)).await;
        assert_eq!(b.channel_count(), 0);
    }

    #[tokio::test]
    async fn group_channel_fans_out_and_isolates_by_group() {
        // §11.6 D2: a group shared-topic channel is keyed by the OWNING group —
        // a publish to group A never reaches a subscriber on group B, and the
        // group map is separate from the per-app map.
        let b = InProcessBroadcaster::new(16);
        let g_a = GroupId::new();
        let g_b = GroupId::new();
        let mut rx_a1 = b.subscribe_group(g_a, "events").await.unwrap();
        let mut rx_a2 = b.subscribe_group(g_a, "events").await.unwrap();
        let mut rx_b = b.subscribe_group(g_b, "events").await.unwrap();
        assert_eq!(b.channel_count(), 2, "one channel per (group, topic)");

        b.publish_group(g_a, "events", event("events", 7)).await;
        // Both of group A's subscribers see it...
        assert_eq!(rx_a1.recv().await.unwrap().message, json!({ "n": 7 }));
        assert_eq!(rx_a2.recv().await.unwrap().message, json!({ "n": 7 }));
        // ...group B's does not (owning-group isolation).
        assert!(rx_b.try_recv().is_err());

        // GC reclaims both maps.
        drop((rx_a1, rx_a2, rx_b));
        assert_eq!(b.gc(), 2);
        assert_eq!(b.channel_count(), 0);
    }

    #[tokio::test]
    async fn subscribe_refuses_new_channel_past_the_cap() {
        // Bound the maps: a subscribe that would open a NEW channel past the cap
        // is refused, but subscribing to an already-open channel still works.
        let b = InProcessBroadcaster::new(16).with_max_channels(2);
        let app = AppId::new();
        let _rx1 = b.subscribe(app, "a").await.unwrap();
        let _rx2 = b.subscribe(app, "b").await.unwrap();
        // Third distinct topic → over cap → refused.
        assert!(matches!(
            b.subscribe(app, "c").await,
            Err(BroadcasterError::Unavailable(_))
        ));
        // An additional receiver on an existing channel is still allowed.
        assert!(b.subscribe(app, "a").await.is_ok());
    }

    #[tokio::test]
    async fn drop_group_topic_disconnects_subscribers() {
        let b = InProcessBroadcaster::new(16);
        let g = GroupId::new();
        let mut rx = b.subscribe_group(g, "t").await.unwrap();
        b.drop_group_topic(g, "t").await;
        assert!(rx.recv().await.is_err());
        assert_eq!(b.channel_count(), 0);
    }
}