klieo-bus-memory 0.5.0

//! In-process `Pubsub` implementation.
//!
//! Backed by `tokio::sync::broadcast` per subject. Each `subscribe` call
//! returns a fresh receiver; messages published *after* subscribe are
//! delivered (broadcast semantics — no replay of pre-subscribe messages).
//! Messages are reconstructed per subscriber so each gets a fresh
//! [`klieo_core::Msg`] with its own no-op [`klieo_core::AckHandle`].
//!
//! The default per-subject channel capacity is 1024. Slow subscribers
//! may see `BusError::Retryable("subscriber lagged")` if they fall
//! behind.

use async_trait::async_trait;
use bytes::Bytes;
use klieo_core::bus::{AckHandle, AckHandleImpl, Headers, Msg, MsgStream, Pubsub};
use klieo_core::error::BusError;
use klieo_core::ids::DurableName;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{broadcast, Mutex};
use tokio_stream::wrappers::BroadcastStream;
use tokio_stream::StreamExt;

const DEFAULT_CAPACITY: usize = 1024;

struct State {
    /// Per-subject broadcast sender. Receivers are created via
    /// `Sender::subscribe`.
    subjects: HashMap<String, broadcast::Sender<(Bytes, Headers)>>,
    /// Per-subject channel capacity, applied lazily on first publish/subscribe.
    capacity: usize,
}

impl Default for State {
    fn default() -> Self {
        Self {
            subjects: HashMap::new(),
            capacity: DEFAULT_CAPACITY,
        }
    }
}

/// In-process `Pubsub` impl.
#[derive(Clone)]
pub struct MemoryPubsub {
    state: Arc<Mutex<State>>,
}

impl MemoryPubsub {
    /// Build an empty pubsub with the default 1024-message per-subject buffer.
    pub fn new() -> Self {
        Self::with_buffer_size(DEFAULT_CAPACITY)
    }

    /// Build an empty pubsub with a custom per-subject buffer size.
    /// Larger buffers tolerate slow subscribers; smaller buffers reject
    /// fast publishers earlier (slow subscribers see
    /// `BusError::Retryable("subscriber lagged")`).
    pub fn with_buffer_size(buffer: usize) -> Self {
        Self {
            state: Arc::new(Mutex::new(State {
                subjects: HashMap::new(),
                capacity: buffer.max(1),
            })),
        }
    }

    async fn sender(&self, subject: &str) -> broadcast::Sender<(Bytes, Headers)> {
        let mut g = self.state.lock().await;
        let cap = g.capacity;
        g.subjects
            .entry(subject.to_string())
            .or_insert_with(|| {
                let (tx, _rx) = broadcast::channel(cap);
                tx
            })
            .clone()
    }

    /// Drop the broadcast channel for `subject`. Used by
    /// [`crate::request_reply::MemoryRequestReply`] to clean up the
    /// short-lived `_reply.{id}` subjects it mints per RPC — without this
    /// the per-subject HashMap grows unbounded with one entry per
    /// request (W1.A1 / round-1 CRIT).
    ///
    /// Safe to call when no entry exists; behaves as a no-op.
    pub async fn remove_subject(&self, subject: &str) {
        self.state.lock().await.subjects.remove(subject);
    }

    /// Number of broadcast subjects currently retained. Test/observability
    /// helper; production callers should not rely on the exact value.
    pub async fn subject_count(&self) -> usize {
        self.state.lock().await.subjects.len()
    }
}

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

#[async_trait]
impl Pubsub for MemoryPubsub {
    async fn publish(
        &self,
        subject: &str,
        payload: Bytes,
        headers: Headers,
    ) -> Result<(), BusError> {
        let tx = self.sender(subject).await;
        // `send` returns Err only when there are no active receivers —
        // that is not an error in pub/sub semantics.
        let _ = tx.send((payload, headers));
        Ok(())
    }

    async fn subscribe(&self, subject: &str, _durable: DurableName) -> Result<MsgStream, BusError> {
        let tx = self.sender(subject).await;
        let rx = tx.subscribe();
        let subject = subject.to_string();
        let stream = BroadcastStream::new(rx).map(move |res| match res {
            Ok((payload, headers)) => Ok(Msg {
                subject: subject.clone(),
                payload,
                headers,
                ack: AckHandle::new(Box::new(NoopAck)),
            }),
            Err(_) => {
                tracing::warn!(
                    target: "klieo.bus.memory",
                    subject = %subject,
                    "subscriber lagged - increase MemoryPubsub buffer size or speed up consumer"
                );
                Err(BusError::Retryable("subscriber lagged".into()))
            }
        });
        Ok(Box::pin(stream))
    }
}

/// No-op ack handle. In-process pubsub has no redelivery semantics; the
/// trait surface is preserved for consistency.
pub(crate) struct NoopAck;

#[async_trait]
impl AckHandleImpl for NoopAck {
    async fn ack(self: Box<Self>) -> Result<(), BusError> {
        Ok(())
    }
    async fn nak(self: Box<Self>, _delay: Duration) -> Result<(), BusError> {
        Ok(())
    }
    async fn term(self: Box<Self>) -> Result<(), BusError> {
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use klieo_core::Pubsub;
    use tokio_stream::StreamExt;

    #[tokio::test]
    async fn publish_delivers_to_subscriber() {
        let bus = MemoryPubsub::new();
        let mut sub = bus
            .subscribe("subject.a", DurableName::new("d1"))
            .await
            .unwrap();
        bus.publish("subject.a", Bytes::from_static(b"hi"), Headers::new())
            .await
            .unwrap();
        let msg = sub.next().await.unwrap().unwrap();
        assert_eq!(msg.subject, "subject.a");
        assert_eq!(msg.payload, Bytes::from_static(b"hi"));
        // Memory impl: ack is a no-op but must succeed.
        msg.ack.ack().await.unwrap();
    }

    #[tokio::test]
    async fn two_subscribers_both_receive() {
        let bus = MemoryPubsub::new();
        let mut s1 = bus
            .subscribe("subject.b", DurableName::new("d1"))
            .await
            .unwrap();
        let mut s2 = bus
            .subscribe("subject.b", DurableName::new("d2"))
            .await
            .unwrap();
        bus.publish("subject.b", Bytes::from_static(b"x"), Headers::new())
            .await
            .unwrap();
        assert_eq!(
            s1.next().await.unwrap().unwrap().payload,
            Bytes::from_static(b"x")
        );
        assert_eq!(
            s2.next().await.unwrap().unwrap().payload,
            Bytes::from_static(b"x")
        );
    }

    #[tokio::test]
    async fn publish_with_no_subscribers_succeeds() {
        let bus = MemoryPubsub::new();
        bus.publish("nobody.home", Bytes::from_static(b"x"), Headers::new())
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn headers_preserved() {
        let bus = MemoryPubsub::new();
        let mut sub = bus
            .subscribe("subject.h", DurableName::new("d"))
            .await
            .unwrap();
        let mut h = Headers::new();
        h.insert("k".into(), "v".into());
        bus.publish("subject.h", Bytes::from_static(b"hi"), h)
            .await
            .unwrap();
        let msg = sub.next().await.unwrap().unwrap();
        assert_eq!(msg.headers.get("k").map(|s| s.as_str()), Some("v"));
    }

    #[tokio::test]
    async fn small_buffer_lags_on_burst() {
        let bus = MemoryPubsub::with_buffer_size(2);
        let mut sub = bus
            .subscribe("subject.lag", DurableName::new("d"))
            .await
            .unwrap();
        // Publish 4 with no consumption — the broadcast channel of size 2 will overflow.
        for i in 0..4u8 {
            bus.publish("subject.lag", Bytes::from(vec![i]), Headers::new())
                .await
                .unwrap();
        }
        // Drain until we observe a lag (broadcast may yield some messages
        // before flagging the lag). The lag must surface within the burst.
        let mut saw_lag = false;
        for _ in 0..6 {
            let next = sub.next().await.unwrap();
            if let Err(BusError::Retryable(ref m)) = next {
                if m.contains("lagged") {
                    saw_lag = true;
                    break;
                }
            }
        }
        assert!(saw_lag, "expected lagged err within the burst");
    }

    #[tokio::test]
    async fn large_buffer_does_not_lag() {
        let bus = MemoryPubsub::with_buffer_size(64);
        let mut sub = bus
            .subscribe("subject.nolag", DurableName::new("d"))
            .await
            .unwrap();
        for i in 0..10u8 {
            bus.publish("subject.nolag", Bytes::from(vec![i]), Headers::new())
                .await
                .unwrap();
        }
        for _ in 0..10 {
            let msg = sub.next().await.unwrap().expect("no lag");
            let _ = msg.payload;
        }
    }
}