klieo_bus_memory/
request_reply.rs

1//! In-process `RequestReply` implementation.
2//!
3//! Builds on top of [`crate::pubsub::MemoryPubsub`]: each `request`
4//! generates a unique reply subject, subscribes to it, publishes the
5//! request with a `reply_to` header carrying the reply subject, then
6//! awaits exactly one message on the reply subject (bounded by the
7//! caller's timeout).
8
9use crate::pubsub::MemoryPubsub;
10use async_trait::async_trait;
11use bytes::Bytes;
12use klieo_core::bus::{Headers, RequestReply};
13use klieo_core::error::BusError;
14use klieo_core::ids::DurableName;
15use klieo_core::Pubsub;
16use std::sync::atomic::{AtomicU64, Ordering};
17use std::sync::Arc;
18use std::time::Duration;
19use tokio_stream::StreamExt;
20
21/// RAII guard that removes a reply subject from the underlying pubsub
22/// when the outer `request` future completes — successfully, on timeout,
23/// or via cancellation (caller drops the future). Without this,
24/// [`MemoryPubsub`] retains a `broadcast::Sender` per ever-emitted reply
25/// id, leaking memory at request-rate (W1.A1 / round-1 CRIT).
26struct ReplySubjectGuard {
27    pubsub: Arc<MemoryPubsub>,
28    subject: String,
29}
30
31impl Drop for ReplySubjectGuard {
32    fn drop(&mut self) {
33        // `Drop` cannot await; spawn the cleanup so it runs on the
34        // current Tokio runtime regardless of how the parent future
35        // unwound. The cleanup task is O(1) — one HashMap::remove under
36        // the same mutex `subscribe` takes.
37        let pubsub = Arc::clone(&self.pubsub);
38        let subject = std::mem::take(&mut self.subject);
39        tokio::spawn(async move {
40            pubsub.remove_subject(&subject).await;
41        });
42    }
43}
44
45/// In-process `RequestReply` impl.
46pub struct MemoryRequestReply {
47    pubsub: Arc<MemoryPubsub>,
48    next_reply_id: AtomicU64,
49}
50
51impl MemoryRequestReply {
52    /// Build atop the supplied pubsub.
53    pub fn new(pubsub: Arc<MemoryPubsub>) -> Self {
54        Self {
55            pubsub,
56            next_reply_id: AtomicU64::new(0),
57        }
58    }
59}
60
61#[async_trait]
62impl RequestReply for MemoryRequestReply {
63    async fn request(
64        &self,
65        subject: &str,
66        payload: Bytes,
67        timeout: Duration,
68    ) -> Result<Bytes, BusError> {
69        let id = self.next_reply_id.fetch_add(1, Ordering::Relaxed);
70        let reply_subject = format!("_reply.{id}");
71
72        let mut sub = self
73            .pubsub
74            .subscribe(&reply_subject, DurableName::new("rr"))
75            .await?;
76
77        // Hold the guard until the future completes (any path).
78        let _guard = ReplySubjectGuard {
79            pubsub: Arc::clone(&self.pubsub),
80            subject: reply_subject.clone(),
81        };
82
83        let mut headers = Headers::new();
84        headers.insert("reply_to".into(), reply_subject);
85        self.pubsub.publish(subject, payload, headers).await?;
86
87        match tokio::time::timeout(timeout, sub.next()).await {
88            Ok(Some(Ok(msg))) => Ok(msg.payload),
89            Ok(Some(Err(e))) => Err(e),
90            Ok(None) => Err(BusError::NotFound("reply stream closed".into())),
91            Err(_) => Err(BusError::Timeout),
92        }
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99    use crate::pubsub::MemoryPubsub;
100    use klieo_core::ids::DurableName;
101    use klieo_core::Pubsub;
102    use std::sync::Arc;
103    use std::time::Duration;
104    use tokio_stream::StreamExt;
105
106    /// Polling interval used when waiting for a subscription to register
107    /// before sending a request in tests.
108    const IN_PROCESS_POLL_INTERVAL: Duration = Duration::from_millis(10);
109    /// Polling interval used when waiting for reply-subject cleanup to
110    /// propagate in tests.
111    const IN_PROCESS_REPLY_POLL_INTERVAL: Duration = Duration::from_millis(5);
112
113    #[tokio::test]
114    async fn request_returns_reply() {
115        let pubsub = Arc::new(MemoryPubsub::new());
116        let rr = MemoryRequestReply::new(pubsub.clone());
117
118        // Spawn a tiny "echo handler" on the request subject.
119        let pubsub_h = pubsub.clone();
120        tokio::spawn(async move {
121            let mut sub = pubsub_h
122                .subscribe("svc.echo", DurableName::new("h1"))
123                .await
124                .unwrap();
125            let msg = sub.next().await.unwrap().unwrap();
126            let reply_to = msg.headers.get("reply_to").unwrap().clone();
127            pubsub_h
128                .publish(&reply_to, msg.payload, Headers::new())
129                .await
130                .unwrap();
131        });
132
133        // Give the spawned subscribe a moment to register.
134        tokio::time::sleep(IN_PROCESS_POLL_INTERVAL).await;
135
136        let reply = rr
137            .request(
138                "svc.echo",
139                Bytes::from_static(b"ping"),
140                Duration::from_millis(500),
141            )
142            .await
143            .unwrap();
144        assert_eq!(reply, Bytes::from_static(b"ping"));
145    }
146
147    #[tokio::test]
148    async fn request_times_out_when_no_handler() {
149        let pubsub = Arc::new(MemoryPubsub::new());
150        let rr = MemoryRequestReply::new(pubsub);
151        let err = rr
152            .request(
153                "svc.nobody",
154                Bytes::from_static(b"x"),
155                Duration::from_millis(50),
156            )
157            .await
158            .unwrap_err();
159        assert!(matches!(err, klieo_core::BusError::Timeout));
160    }
161
162    /// W1.A1 / round-1 CRIT: every `request()` mints a fresh
163    /// `_reply.{id}` subject. Without the RAII guard those subjects
164    /// stay in `MemoryPubsub.state.subjects` forever, leaking one entry
165    /// per RPC. After this fix, subject count returns to the pre-burst
166    /// baseline once cleanup tasks drain.
167    #[tokio::test]
168    async fn reply_subjects_are_cleaned_up_after_timeout() {
169        let pubsub = Arc::new(MemoryPubsub::new());
170        let rr = MemoryRequestReply::new(pubsub.clone());
171
172        let baseline = pubsub.subject_count().await;
173
174        for _ in 0..50 {
175            let _ = rr
176                .request(
177                    "svc.nobody",
178                    Bytes::from_static(b"x"),
179                    Duration::from_millis(5),
180                )
181                .await;
182        }
183
184        // Cleanup tasks are spawned in Drop — give them a moment to
185        // acquire the lock and remove their entry.
186        for _ in 0..20 {
187            tokio::task::yield_now().await;
188            tokio::time::sleep(IN_PROCESS_REPLY_POLL_INTERVAL).await;
189            if pubsub.subject_count().await <= baseline + 1 {
190                break;
191            }
192        }
193
194        let after = pubsub.subject_count().await;
195        // +1 tolerance because the request subject ("svc.nobody") gets
196        // created lazily by the publish call and is intentionally kept.
197        assert!(
198            after <= baseline + 1,
199            "expected reply subjects to be cleaned up (baseline={baseline}, after={after})"
200        );
201    }
202
203    /// Cancellation path: the caller's future is dropped before the
204    /// reply arrives. RAII guard must still fire.
205    #[tokio::test]
206    async fn reply_subject_cleaned_on_cancellation() {
207        let pubsub = Arc::new(MemoryPubsub::new());
208        let rr = Arc::new(MemoryRequestReply::new(pubsub.clone()));
209        let baseline = pubsub.subject_count().await;
210
211        // Drive 30 requests through cancellation by wrapping each in a
212        // tight `tokio::time::timeout` and dropping immediately.
213        for _ in 0..30 {
214            let rr = Arc::clone(&rr);
215            let fut = rr.request(
216                "svc.nobody",
217                Bytes::from_static(b"x"),
218                Duration::from_secs(60),
219            );
220            let _ = tokio::time::timeout(Duration::from_millis(1), fut).await;
221        }
222
223        for _ in 0..30 {
224            tokio::task::yield_now().await;
225            tokio::time::sleep(IN_PROCESS_REPLY_POLL_INTERVAL).await;
226            if pubsub.subject_count().await <= baseline + 1 {
227                break;
228            }
229        }
230        let after = pubsub.subject_count().await;
231        assert!(
232            after <= baseline + 1,
233            "cancel path leaked reply subjects (baseline={baseline}, after={after})"
234        );
235    }
236}
klieo_bus_memory/request_reply.rs

klieo_bus_memory/
request_reply.rs
