use std::sync::{Arc, Mutex};
use bytes::Bytes;
use ruststream::{
Headers, OutgoingMessage, OwnedTransactions, PairError, PublishPolicy, Publisher, Transaction,
TransactionalPublisher,
};
use tracing::warn;
use super::broker::{ConnectedLapinTestBroker, TestBrokerState};
use crate::error::AmqpError;
type Buffered = (String, Bytes, Headers);
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
#[must_use]
pub struct LapinTestPublish;
impl LapinTestPublish {
#[must_use]
pub fn bind(self, connected: &ConnectedLapinTestBroker) -> LapinTestPublisher {
LapinTestPublisher {
state: connected.state(),
txn: Arc::new(Mutex::new(None)),
}
}
}
impl PublishPolicy<ConnectedLapinTestBroker> for LapinTestPublish {
type Live = LapinTestPublisher;
async fn pair(self, connected: &ConnectedLapinTestBroker) -> Result<Self::Live, PairError> {
Ok(self.bind(connected))
}
}
#[derive(Debug, Clone)]
pub struct LapinTestPublisher {
state: Arc<TestBrokerState>,
txn: Arc<Mutex<Option<Vec<Buffered>>>>,
}
impl LapinTestPublisher {
fn route(&self, queue: &str, payload: &Bytes, headers: &Headers) {
self.state
.router
.publish(queue, payload, headers, self.state.coordinator().as_ref());
}
}
impl Publisher for LapinTestPublisher {
type Error = AmqpError;
async fn publish(&self, msg: OutgoingMessage<'_>) -> Result<(), Self::Error> {
if msg.name().is_empty() {
return Err(AmqpError::InvalidOptions(
"routing key must not be empty; on the default exchange it names the target queue"
.to_owned(),
));
}
self.state.ensure_live(msg.name())?;
{
let mut txn = self.txn.lock().expect("transaction buffer mutex poisoned");
if let Some(buffer) = txn.as_mut() {
buffer.push((
msg.name().to_owned(),
Bytes::copy_from_slice(msg.payload()),
msg.headers().clone(),
));
return Ok(());
}
}
self.route(
msg.name(),
&Bytes::copy_from_slice(msg.payload()),
msg.headers(),
);
Ok(())
}
}
impl TransactionalPublisher for LapinTestPublisher {
async fn begin_transaction(&self) -> Result<(), Self::Error> {
let already_open = {
let mut txn = self.txn.lock().expect("transaction buffer mutex poisoned");
let open = txn.is_some();
if !open {
*txn = Some(Vec::new());
}
open
};
if already_open {
return Err(AmqpError::Transaction(
"a transaction is already open on this test publisher; commit or abort it before \
beginning another"
.to_owned(),
));
}
Ok(())
}
async fn commit(&self) -> Result<(), Self::Error> {
let buffered = {
let mut txn = self.txn.lock().expect("transaction buffer mutex poisoned");
txn.take()
};
let Some(buffered) = buffered else {
return Err(AmqpError::Transaction(
"commit with no open transaction on this test publisher".to_owned(),
));
};
for (queue, payload, headers) in buffered {
self.state.ensure_live(&queue)?;
self.route(&queue, &payload, &headers);
}
Ok(())
}
async fn abort(&self) -> Result<(), Self::Error> {
let discarded = self
.txn
.lock()
.expect("transaction buffer mutex poisoned")
.take();
if discarded.is_none() {
return Err(AmqpError::Transaction(
"abort with no open transaction on this test publisher".to_owned(),
));
}
Ok(())
}
}
impl OwnedTransactions for LapinTestPublisher {
type Transaction = LapinTestTransaction;
async fn transaction(&self) -> Result<Self::Transaction, Self::Error> {
Ok(LapinTestTransaction {
publisher: self.clone(),
buffered: Vec::new(),
settled: false,
})
}
}
#[must_use = "a transaction does nothing until settled with commit() or abort()"]
pub struct LapinTestTransaction {
publisher: LapinTestPublisher,
buffered: Vec<Buffered>,
settled: bool,
}
impl std::fmt::Debug for LapinTestTransaction {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("LapinTestTransaction")
.field("buffered", &self.buffered.len())
.field("settled", &self.settled)
.finish_non_exhaustive()
}
}
impl Drop for LapinTestTransaction {
fn drop(&mut self) {
if !self.settled {
warn!(
target: "ruststream_lapin",
buffered = self.buffered.len(),
"owned transaction dropped without commit or abort; its buffered messages are \
discarded"
);
}
}
}
impl Transaction for LapinTestTransaction {
type Error = AmqpError;
async fn publish(&mut self, msg: OutgoingMessage<'_>) -> Result<(), Self::Error> {
if msg.name().is_empty() {
return Err(AmqpError::InvalidOptions(
"routing key must not be empty; on the default exchange it names the target queue"
.to_owned(),
));
}
self.buffered.push((
msg.name().to_owned(),
Bytes::copy_from_slice(msg.payload()),
msg.headers().clone(),
));
Ok(())
}
async fn commit(mut self) -> Result<(), Self::Error> {
self.settled = true;
for (queue, payload, headers) in &self.buffered {
self.publisher.state.ensure_live(queue)?;
self.publisher.route(queue, payload, headers);
}
Ok(())
}
async fn abort(mut self) -> Result<(), Self::Error> {
self.settled = true;
Ok(())
}
}