docling_rag/queue/
memory.rs

1//! In-process message queue built on a tokio broadcast channel.
2
3use super::{MessageQueue, QueueReceiver};
4use crate::Result;
5use async_trait::async_trait;
6use tokio::sync::broadcast;
7
8/// A fan-out, in-process queue. Subscribers created before a publish receive it.
9pub struct MemoryQueue {
10    tx: broadcast::Sender<Vec<u8>>,
11}
12
13impl MemoryQueue {
14    /// Create a queue with a default buffer capacity.
15    pub fn new() -> Self {
16        Self::with_capacity(1024)
17    }
18
19    /// Create a queue with a specific per-subscriber buffer capacity.
20    pub fn with_capacity(cap: usize) -> Self {
21        let (tx, _rx) = broadcast::channel(cap);
22        MemoryQueue { tx }
23    }
24}
25
26impl Default for MemoryQueue {
27    fn default() -> Self {
28        MemoryQueue::new()
29    }
30}
31
32#[async_trait]
33impl MessageQueue for MemoryQueue {
34    async fn publish(&self, payload: &[u8]) -> Result<()> {
35        // Err only means "no subscribers"; that is not a failure to publish.
36        let _ = self.tx.send(payload.to_vec());
37        Ok(())
38    }
39
40    async fn subscribe(&self) -> Result<Box<dyn QueueReceiver>> {
41        Ok(Box::new(MemoryReceiver {
42            rx: self.tx.subscribe(),
43        }))
44    }
45}
46
47struct MemoryReceiver {
48    rx: broadcast::Receiver<Vec<u8>>,
49}
50
51#[async_trait]
52impl QueueReceiver for MemoryReceiver {
53    async fn recv(&mut self) -> Option<Vec<u8>> {
54        loop {
55            match self.rx.recv().await {
56                Ok(v) => return Some(v),
57                // A slow consumer that lagged: skip dropped messages, keep going.
58                Err(broadcast::error::RecvError::Lagged(_)) => continue,
59                Err(broadcast::error::RecvError::Closed) => return None,
60            }
61        }
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68
69    #[tokio::test]
70    async fn publishes_to_subscriber() {
71        let q = MemoryQueue::new();
72        let mut sub = q.subscribe().await.unwrap();
73        q.publish(b"hello").await.unwrap();
74        assert_eq!(sub.recv().await, Some(b"hello".to_vec()));
75    }
76}
docling_rag/queue/memory.rs

docling_rag/queue/
memory.rs
