klieo_bus_memory/

pubsub.rs

//! In-process `Pubsub` implementation.
//!
//! Backed by `tokio::sync::broadcast` per subscription pattern. 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`].
//!
//! # NATS wildcard parity
//!
//! Subscribe patterns mirror NATS subject semantics (wildcard matching
//! is faithful; delivery guarantees are NOT — see the divergence section
//! below):
//!
//! - Tokens are dot-separated.
//! - `*` matches exactly one token at its position.
//! - `>` matches one or more remaining tokens; valid only as the LAST
//!   token. A pattern that uses `>` anywhere else matches nothing.
//! - Concrete (wildcard-free) patterns match only their literal subject.
//! - Empty tokens (`a..b`) are rejected (no match).
//!
//! On publish, every live subscription whose pattern matches the
//! published subject receives the message.
//!
//! The default per-subscription channel capacity is 1024. Slow
//! subscribers may see `BusError::Retryable("subscriber lagged")` if
//! they fall behind.
//!
//! # Delivery-guarantee divergence from NATS — READ BEFORE RELYING ON IT
//!
//! This in-process bus does NOT honour two parts of the
//! [`klieo_core::Pubsub`] contract that the NATS impl does. It is a
//! fan-out simulator, not a faithful JetStream substitute for these:
//!
//! - **`durable` is ignored.** Every `subscribe` gets its own broadcast
//!   receiver, so multiple subscribers with the SAME durable name each
//!   receive EVERY message (fan-out). They do NOT form a competing-consumer
//!   group that load-balances messages the way NATS durables do. Code that
//!   relies on competing-consumer semantics behaves differently here vs NATS.
//! - **`nak`/`term` are no-ops; there is no redelivery.** A `nak` does not
//!   re-queue the message. Retry logic built on redelivery passes its tests
//!   on this bus but will only actually retry on NATS.
//!
//! Use [`klieo-bus-nats`](../../klieo_bus_nats/index.html) when those
//! guarantees matter.

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;

type PatternSender = broadcast::Sender<(String, Bytes, Headers)>;

struct State {
    /// Subscribers whose pattern is fully concrete (no `*` / `>`).
    /// O(1) lookup on publish — the common case for per-id subjects
    /// like `klieo.a2a.task.t-1`, `klieo.mcp.progress.tok-1`, and
    /// `_reply.{id}`. The broadcast payload carries the concrete
    /// published subject alongside the body/headers so the delivery
    /// shape matches the wildcard path.
    concretes: HashMap<String, PatternSender>,
    /// Subscribers whose pattern contains `*` or `>`. Linear-scanned
    /// on publish; expected to stay tiny (one per cancel-fanout
    /// subscriber per replica).
    wildcards: Vec<(String, PatternSender)>,
    /// Per-channel capacity, applied lazily on first subscribe.
    capacity: usize,
}

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

/// `true` when the pattern contains a NATS wildcard token (`*` or
/// `>`). Subscribers with such patterns go into `State.wildcards`;
/// everything else goes into `State.concretes` for O(1) publish
/// lookup.
fn is_wildcard_pattern(pattern: &str) -> bool {
    pattern.contains('*') || pattern.contains('>')
}

/// 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-channel buffer.
    pub fn new() -> Self {
        Self::with_buffer_size(DEFAULT_CAPACITY)
    }

    /// Build an empty pubsub with a custom per-channel 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 {
                concretes: HashMap::new(),
                wildcards: Vec::new(),
                capacity: buffer.max(1),
            })),
        }
    }

    /// Get-or-create the broadcast sender for a subscription pattern.
    /// Used by `subscribe` only — `publish` MUST NOT create new entries
    /// so the per-pattern stores do not grow with publish-only
    /// subjects. Routes concrete patterns into `State.concretes`
    /// (O(1) publish lookup) and wildcard patterns into
    /// `State.wildcards` (linear scan only when needed).
    async fn subscribe_sender(&self, pattern: &str) -> PatternSender {
        let mut g = self.state.lock().await;
        let cap = g.capacity;
        if is_wildcard_pattern(pattern) {
            if let Some((_, tx)) = g.wildcards.iter().find(|(p, _)| p == pattern) {
                return tx.clone();
            }
            let (tx, _rx) = broadcast::channel(cap);
            g.wildcards.push((pattern.to_string(), tx.clone()));
            tx
        } else {
            g.concretes
                .entry(pattern.to_string())
                .or_insert_with(|| {
                    let (tx, _rx) = broadcast::channel(cap);
                    tx
                })
                .clone()
        }
    }

    /// Drop the broadcast channel for `pattern`. Used by
    /// [`crate::request_reply::MemoryRequestReply`] to clean up the
    /// short-lived `_reply.{id}` subjects it mints per RPC — without
    /// this the per-pattern stores grow unbounded with one entry
    /// per request.
    ///
    /// Safe to call when no entry exists; behaves as a no-op. Removes
    /// from whichever store (concrete or wildcard) holds the entry.
    pub async fn remove_subject(&self, pattern: &str) {
        let mut g = self.state.lock().await;
        if g.concretes.remove(pattern).is_some() {
            return;
        }
        g.wildcards.retain(|(p, _)| p != pattern);
    }

    /// Number of active subscription patterns currently retained
    /// across both concrete and wildcard stores. Test/observability
    /// helper; production callers should not rely on the exact value.
    pub async fn subject_count(&self) -> usize {
        let g = self.state.lock().await;
        g.concretes.len() + g.wildcards.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> {
        // Snapshot the matching senders under the lock, then send
        // outside the lock so a slow subscriber cannot block publish.
        // Concrete subscribers resolve via O(1) HashMap lookup;
        // wildcards (`*` / `>`) require a linear scan but the vec
        // stays tiny (one per cancel-fanout subscriber per replica).
        let matched: Vec<PatternSender> = {
            let g = self.state.lock().await;
            let mut out: Vec<PatternSender> = Vec::new();
            if let Some(tx) = g.concretes.get(subject) {
                out.push(tx.clone());
            }
            for (pattern, tx) in &g.wildcards {
                if subject_matches(pattern, subject) {
                    out.push(tx.clone());
                }
            }
            out
        };
        for tx in matched {
            // `send` returns Err only when there are no active receivers
            // — not an error in pub/sub semantics.
            let _ = tx.send((subject.to_string(), payload.clone(), headers.clone()));
        }
        Ok(())
    }

    async fn subscribe(&self, subject: &str, _durable: DurableName) -> Result<MsgStream, BusError> {
        let tx = self.subscribe_sender(subject).await;
        let rx = tx.subscribe();
        let pattern = subject.to_string();
        let stream = BroadcastStream::new(rx).map(move |res| match res {
            Ok((concrete_subject, payload, headers)) => Ok(Msg {
                // Mirror NATS: deliver the concrete published subject,
                // not the subscription filter. Wildcard consumers need
                // this to recover which subject actually fired (e.g.
                // strip prefix to extract the task id).
                subject: concrete_subject,
                payload,
                headers,
                ack: AckHandle::new(Box::new(NoopAck)),
            }),
            Err(_) => {
                tracing::warn!(
                    target: "klieo.bus.memory",
                    pattern = %pattern,
                    "subscriber lagged - increase MemoryPubsub buffer size or speed up consumer"
                );
                Err(BusError::Retryable("subscriber lagged".into()))
            }
        });
        Ok(Box::pin(stream))
    }
}

/// NATS-compatible subject matcher. Returns `true` when `concrete` (a
/// wildcard-free published subject) matches the subscription `pattern`
/// per NATS rules:
///
/// - `*` matches exactly one token at its position.
/// - `>` matches one or more remaining tokens; only valid as the LAST
///   token. Malformed (`>` in non-tail position) → `false`.
/// - Empty tokens (`a..b`, leading/trailing `.`) → `false`.
/// - Concrete patterns match only their literal subject.
fn subject_matches(pattern: &str, concrete: &str) -> bool {
    if pattern.is_empty() || concrete.is_empty() {
        return false;
    }
    let mut p = pattern.split('.');
    let mut c = concrete.split('.');
    loop {
        match (p.next(), c.next()) {
            // Empty token (leading/trailing/double `.`) on either
            // side is invalid and never matches.
            (Some(""), _) | (_, Some("")) => return false,
            // `>` is greedy: matches the current concrete token plus
            // any remaining ones, but ONLY when it is the LAST token
            // in the pattern. `>` mid-pattern is malformed.
            (Some(">"), Some(_)) => return p.next().is_none(),
            // `>` requires at least one trailing concrete token —
            // `a.>` does not match `a`.
            (Some(">"), None) => return false,
            // `*` matches exactly one concrete token.
            (Some("*"), Some(_)) => continue,
            (Some("*"), None) => return false,
            // Literal token must match exactly.
            (Some(pt), Some(ct)) if pt == ct => continue,
            (Some(_), Some(_)) => return false,
            // Both ran out at the same time — full match.
            (None, None) => return true,
            // Length mismatch — pattern still has tokens or concrete
            // still has tokens with no `>` to consume them.
            (Some(_), None) | (None, Some(_)) => return false,
        }
    }
}

/// 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;
        }
    }

    #[test]
    fn concrete_matches_self() {
        assert!(subject_matches("a.b.c", "a.b.c"));
        assert!(!subject_matches("a.b.c", "a.b"));
        assert!(!subject_matches("a.b.c", "a.b.c.d"));
        assert!(!subject_matches("a.b.c", "a.x.c"));
    }

    #[test]
    fn single_token_wildcard() {
        assert!(subject_matches("a.*.c", "a.x.c"));
        assert!(subject_matches("a.*.c", "a.42.c"));
        assert!(!subject_matches("a.*.c", "a.b.c.d"));
        assert!(!subject_matches("a.*.c", "a.c"));
        assert!(!subject_matches("*", "a.b"));
        assert!(subject_matches("*", "a"));
    }

    #[test]
    fn greedy_wildcard_matches_remaining_tokens() {
        assert!(subject_matches("a.>", "a.b"));
        assert!(subject_matches("a.>", "a.b.c.d"));
        assert!(!subject_matches("a.>", "a"));
        assert!(!subject_matches("a.>", "b.x"));
        // `a.b.>` requires at least one token after `a.b`.
        assert!(subject_matches("a.b.>", "a.b.x"));
        assert!(!subject_matches("a.b.>", "a.b"));
    }

    #[test]
    fn greedy_only_valid_as_last_token() {
        // `>` mid-pattern is malformed → never matches.
        assert!(!subject_matches("a.>.c", "a.x.c"));
        assert!(!subject_matches("a.>.c", "a.b.c"));
        assert!(!subject_matches("a.>.c", "anything"));
        assert!(!subject_matches(">.a", "x.a"));
    }

    #[test]
    fn bare_greedy_matches_anything_with_one_plus_token() {
        assert!(subject_matches(">", "a"));
        assert!(subject_matches(">", "a.b.c"));
        assert!(!subject_matches(">", ""));
    }

    #[test]
    fn prefix_does_not_partially_match() {
        assert!(!subject_matches("a.b", "a"));
        assert!(!subject_matches("a.b", "a.b.c"));
    }

    #[test]
    fn empty_token_rejected() {
        assert!(!subject_matches("a..b", "a..b"));
        assert!(!subject_matches("a..b", "a.x.b"));
        assert!(!subject_matches("a.b", ""));
        assert!(!subject_matches("", "a.b"));
        assert!(!subject_matches(".a", "x.a"));
        assert!(!subject_matches("a.", "a.x"));
    }

    #[tokio::test]
    async fn greedy_wildcard_subscriber_receives_matching_publishes() {
        let bus = MemoryPubsub::new();
        let mut sub = bus
            .subscribe("klieo.a2a.cancel.>", DurableName::new("d"))
            .await
            .unwrap();
        bus.publish(
            "klieo.a2a.cancel.t-1",
            Bytes::from_static(b"signal"),
            Headers::new(),
        )
        .await
        .unwrap();
        let msg = sub.next().await.unwrap().unwrap();
        assert_eq!(msg.subject, "klieo.a2a.cancel.t-1");
        assert_eq!(msg.payload, Bytes::from_static(b"signal"));
    }
}