klieo-bus-memory 0.3.0

//! In-process `RequestReply` implementation.
//!
//! Builds on top of [`crate::pubsub::MemoryPubsub`]: each `request`
//! generates a unique reply subject, subscribes to it, publishes the
//! request with a `reply_to` header carrying the reply subject, then
//! awaits exactly one message on the reply subject (bounded by the
//! caller's timeout).

use crate::pubsub::MemoryPubsub;
use async_trait::async_trait;
use bytes::Bytes;
use klieo_core::bus::{Headers, RequestReply};
use klieo_core::error::BusError;
use klieo_core::ids::DurableName;
use klieo_core::Pubsub;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio_stream::StreamExt;

/// RAII guard that removes a reply subject from the underlying pubsub
/// when the outer `request` future completes — successfully, on timeout,
/// or via cancellation (caller drops the future). Without this,
/// [`MemoryPubsub`] retains a `broadcast::Sender` per ever-emitted reply
/// id, leaking memory at request-rate (W1.A1 / round-1 CRIT).
struct ReplySubjectGuard {
    pubsub: Arc<MemoryPubsub>,
    subject: String,
}

impl Drop for ReplySubjectGuard {
    fn drop(&mut self) {
        // `Drop` cannot await; spawn the cleanup so it runs on the
        // current Tokio runtime regardless of how the parent future
        // unwound. The cleanup task is O(1) — one HashMap::remove under
        // the same mutex `subscribe` takes.
        let pubsub = Arc::clone(&self.pubsub);
        let subject = std::mem::take(&mut self.subject);
        tokio::spawn(async move {
            pubsub.remove_subject(&subject).await;
        });
    }
}

/// In-process `RequestReply` impl.
pub struct MemoryRequestReply {
    pubsub: Arc<MemoryPubsub>,
    next_reply_id: AtomicU64,
}

impl MemoryRequestReply {
    /// Build atop the supplied pubsub.
    pub fn new(pubsub: Arc<MemoryPubsub>) -> Self {
        Self {
            pubsub,
            next_reply_id: AtomicU64::new(0),
        }
    }
}

#[async_trait]
impl RequestReply for MemoryRequestReply {
    async fn request(
        &self,
        subject: &str,
        payload: Bytes,
        timeout: Duration,
    ) -> Result<Bytes, BusError> {
        let id = self.next_reply_id.fetch_add(1, Ordering::Relaxed);
        let reply_subject = format!("_reply.{id}");

        let mut sub = self
            .pubsub
            .subscribe(&reply_subject, DurableName::new("rr"))
            .await?;

        // Hold the guard until the future completes (any path).
        let _guard = ReplySubjectGuard {
            pubsub: Arc::clone(&self.pubsub),
            subject: reply_subject.clone(),
        };

        let mut headers = Headers::new();
        headers.insert("reply_to".into(), reply_subject);
        self.pubsub.publish(subject, payload, headers).await?;

        match tokio::time::timeout(timeout, sub.next()).await {
            Ok(Some(Ok(msg))) => Ok(msg.payload),
            Ok(Some(Err(e))) => Err(e),
            Ok(None) => Err(BusError::NotFound("reply stream closed".into())),
            Err(_) => Err(BusError::Timeout),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::pubsub::MemoryPubsub;
    use klieo_core::ids::DurableName;
    use klieo_core::Pubsub;
    use std::sync::Arc;
    use std::time::Duration;
    use tokio_stream::StreamExt;

    #[tokio::test]
    async fn request_returns_reply() {
        let pubsub = Arc::new(MemoryPubsub::new());
        let rr = MemoryRequestReply::new(pubsub.clone());

        // Spawn a tiny "echo handler" on the request subject.
        let pubsub_h = pubsub.clone();
        tokio::spawn(async move {
            let mut sub = pubsub_h
                .subscribe("svc.echo", DurableName::new("h1"))
                .await
                .unwrap();
            let msg = sub.next().await.unwrap().unwrap();
            let reply_to = msg.headers.get("reply_to").unwrap().clone();
            pubsub_h
                .publish(&reply_to, msg.payload, Headers::new())
                .await
                .unwrap();
        });

        // Give the spawned subscribe a moment to register.
        tokio::time::sleep(Duration::from_millis(10)).await;

        let reply = rr
            .request(
                "svc.echo",
                Bytes::from_static(b"ping"),
                Duration::from_millis(500),
            )
            .await
            .unwrap();
        assert_eq!(reply, Bytes::from_static(b"ping"));
    }

    #[tokio::test]
    async fn request_times_out_when_no_handler() {
        let pubsub = Arc::new(MemoryPubsub::new());
        let rr = MemoryRequestReply::new(pubsub);
        let err = rr
            .request(
                "svc.nobody",
                Bytes::from_static(b"x"),
                Duration::from_millis(50),
            )
            .await
            .unwrap_err();
        assert!(matches!(err, klieo_core::BusError::Timeout));
    }

    /// W1.A1 / round-1 CRIT: every `request()` mints a fresh
    /// `_reply.{id}` subject. Without the RAII guard those subjects
    /// stay in `MemoryPubsub.state.subjects` forever, leaking one entry
    /// per RPC. After this fix, subject count returns to the pre-burst
    /// baseline once cleanup tasks drain.
    #[tokio::test]
    async fn reply_subjects_are_cleaned_up_after_timeout() {
        let pubsub = Arc::new(MemoryPubsub::new());
        let rr = MemoryRequestReply::new(pubsub.clone());

        let baseline = pubsub.subject_count().await;

        for _ in 0..50 {
            let _ = rr
                .request(
                    "svc.nobody",
                    Bytes::from_static(b"x"),
                    Duration::from_millis(5),
                )
                .await;
        }

        // Cleanup tasks are spawned in Drop — give them a moment to
        // acquire the lock and remove their entry.
        for _ in 0..20 {
            tokio::task::yield_now().await;
            tokio::time::sleep(Duration::from_millis(5)).await;
            if pubsub.subject_count().await <= baseline + 1 {
                break;
            }
        }

        let after = pubsub.subject_count().await;
        // +1 tolerance because the request subject ("svc.nobody") gets
        // created lazily by the publish call and is intentionally kept.
        assert!(
            after <= baseline + 1,
            "expected reply subjects to be cleaned up (baseline={baseline}, after={after})"
        );
    }

    /// Cancellation path: the caller's future is dropped before the
    /// reply arrives. RAII guard must still fire.
    #[tokio::test]
    async fn reply_subject_cleaned_on_cancellation() {
        let pubsub = Arc::new(MemoryPubsub::new());
        let rr = Arc::new(MemoryRequestReply::new(pubsub.clone()));
        let baseline = pubsub.subject_count().await;

        // Drive 30 requests through cancellation by wrapping each in a
        // tight `tokio::time::timeout` and dropping immediately.
        for _ in 0..30 {
            let rr = Arc::clone(&rr);
            let fut = rr.request(
                "svc.nobody",
                Bytes::from_static(b"x"),
                Duration::from_secs(60),
            );
            let _ = tokio::time::timeout(Duration::from_millis(1), fut).await;
        }

        for _ in 0..30 {
            tokio::task::yield_now().await;
            tokio::time::sleep(Duration::from_millis(5)).await;
            if pubsub.subject_count().await <= baseline + 1 {
                break;
            }
        }
        let after = pubsub.subject_count().await;
        assert!(
            after <= baseline + 1,
            "cancel path leaked reply subjects (baseline={baseline}, after={after})"
        );
    }
}