Skip to main content

ruststream_lapin/testing/
publisher.rs

1//! The in-process publisher, with buffered transactions.
2
3use std::sync::{Arc, Mutex};
4
5use bytes::Bytes;
6use ruststream::{Headers, OutgoingMessage, Publisher, TransactionalPublisher};
7
8use super::broker::TestBrokerState;
9use crate::error::AmqpError;
10
11type Buffered = (String, Bytes, Headers);
12
13/// Publisher into the in-process router.
14///
15/// Mirrors [`ConfirmsPublisher`](crate::ConfirmsPublisher) transaction semantics: publishes
16/// buffer between `begin_transaction` and `commit`, and `abort` discards them. Clones share the
17/// transaction buffer.
18#[derive(Debug, Clone)]
19pub struct LapinTestPublisher {
20    state: Arc<TestBrokerState>,
21    txn: Arc<Mutex<Option<Vec<Buffered>>>>,
22}
23
24impl LapinTestPublisher {
25    pub(crate) fn new(state: Arc<TestBrokerState>) -> Self {
26        Self {
27            state,
28            txn: Arc::new(Mutex::new(None)),
29        }
30    }
31
32    fn route(&self, queue: &str, payload: &Bytes, headers: &Headers) {
33        self.state
34            .router
35            .publish(queue, payload, headers, self.state.coordinator().as_ref());
36    }
37}
38
39impl Publisher for LapinTestPublisher {
40    type Error = AmqpError;
41
42    /// Routes `msg` to subscribers of the queue named by `msg.name()`.
43    ///
44    /// # Errors
45    ///
46    /// Returns [`AmqpError::InvalidOptions`] when the routing key is empty.
47    async fn publish(&self, msg: OutgoingMessage<'_>) -> Result<(), Self::Error> {
48        if msg.name().is_empty() {
49            return Err(AmqpError::InvalidOptions(
50                "routing key must not be empty; on the default exchange it names the target queue"
51                    .to_owned(),
52            ));
53        }
54        {
55            let mut txn = self.txn.lock().expect("transaction buffer mutex poisoned");
56            if let Some(buffer) = txn.as_mut() {
57                buffer.push((
58                    msg.name().to_owned(),
59                    Bytes::copy_from_slice(msg.payload()),
60                    msg.headers().clone(),
61                ));
62                return Ok(());
63            }
64        }
65        self.route(
66            msg.name(),
67            &Bytes::copy_from_slice(msg.payload()),
68            msg.headers(),
69        );
70        Ok(())
71    }
72}
73
74impl TransactionalPublisher for LapinTestPublisher {
75    /// Opens the buffering transaction; a no-op when one is already open.
76    ///
77    /// # Errors
78    ///
79    /// Never fails.
80    async fn begin_transaction(&self) -> Result<(), Self::Error> {
81        self.txn
82            .lock()
83            .expect("transaction buffer mutex poisoned")
84            .get_or_insert_with(Vec::new);
85        Ok(())
86    }
87
88    /// Replays the buffered publishes in order; a no-op when no transaction is open.
89    ///
90    /// # Errors
91    ///
92    /// Never fails.
93    async fn commit(&self) -> Result<(), Self::Error> {
94        let buffered = {
95            let mut txn = self.txn.lock().expect("transaction buffer mutex poisoned");
96            txn.take()
97        };
98        if let Some(buffered) = buffered {
99            for (queue, payload, headers) in buffered {
100                self.route(&queue, &payload, &headers);
101            }
102        }
103        Ok(())
104    }
105
106    /// Discards the buffered publishes; a no-op when no transaction is open.
107    ///
108    /// # Errors
109    ///
110    /// Never fails.
111    async fn abort(&self) -> Result<(), Self::Error> {
112        self.txn
113            .lock()
114            .expect("transaction buffer mutex poisoned")
115            .take();
116        Ok(())
117    }
118}