#![cfg_attr(
not(feature = "testing"),
allow(clippy::infallible_destructuring_match)
)]
use std::future::{Future, ready};
use std::sync::{Arc, Mutex};
use bytes::Bytes;
use lapin::options::{BasicPublishOptions, ConfirmSelectOptions};
use lapin::{BasicProperties, Channel};
use lapin::{Confirmation, PublisherConfirm};
use ruststream::{
HeaderMap, Lend, OutgoingFor, OutgoingMessage, Publisher, Take, TransactionalPublisher,
};
use crate::broker::{AmqpConnection, ConnectedLapinBroker, Link};
use crate::channel::ChannelCell;
use crate::convert;
use crate::error::AmqpError;
#[cfg(feature = "testing")]
use crate::in_process::Bus;
use crate::publish_policy::PublishOptions;
use crate::publish_step::LapinPublishOptions;
#[derive(Debug, Clone)]
pub(crate) struct Buffered {
pub(crate) routing_key: String,
pub(crate) payload: Bytes,
pub(crate) headers: HeaderMap,
pub(crate) options: LapinPublishOptions,
}
impl Buffered {
pub(crate) fn lent(msg: OutgoingMessage<'_>, options: Option<&LapinPublishOptions>) -> Self {
let (routing_key, payload, headers) = msg.into_parts();
Self {
routing_key: routing_key.to_owned(),
payload: Bytes::copy_from_slice(payload),
headers,
options: options.copied().unwrap_or_default(),
}
}
pub(crate) fn taken(msg: OutgoingFor<'_, Take>, options: Option<&LapinPublishOptions>) -> Self {
let (routing_key, payload, headers) = msg.into_parts();
Self {
routing_key: routing_key.to_owned(),
payload: payload.freeze(),
headers,
options: options.copied().unwrap_or_default(),
}
}
}
pub(crate) async fn do_publish(
channel: &Channel,
exchange: &str,
routing_key: &str,
payload: &[u8],
properties: BasicProperties,
) -> Result<PublisherConfirm, AmqpError> {
channel
.basic_publish(
convert::short(exchange, "exchange name")?,
convert::short(routing_key, "routing key")?,
BasicPublishOptions::default(),
payload,
properties,
)
.await
.map_err(AmqpError::publish)
}
#[derive(Debug, Clone)]
pub struct LapinPublisher {
link: Link,
options: PublishOptions,
}
impl LapinPublisher {
pub(crate) fn new(connected: &ConnectedLapinBroker, options: PublishOptions) -> Self {
Self {
link: connected.link().clone(),
options,
}
}
}
impl Publisher for LapinPublisher {
type Payload = Lend;
type Error = AmqpError;
type Options = LapinPublishOptions;
async fn publish(
&self,
msg: OutgoingMessage<'_>,
options: Option<&Self::Options>,
) -> Result<(), Self::Error> {
let conn = match &self.link {
Link::Amqp(conn) => conn,
#[cfg(feature = "testing")]
Link::InProcess(bus) => {
return bus.publish(
&self.options.exchange,
msg.name(),
msg.payload(),
msg.headers(),
&self.options.resolve(options),
);
}
};
let channel = conn.live_publish_channel(msg.name()).await?;
let properties =
convert::properties_for_publish(msg.headers(), &self.options.resolve(options))?;
let _confirm = do_publish(
&channel,
&self.options.exchange,
msg.name(),
msg.payload(),
properties,
)
.await?;
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct ConfirmsPublisher {
link: Link,
options: PublishOptions,
channel: Arc<ChannelCell<Channel>>,
txn: Arc<Mutex<Option<Vec<Buffered>>>>,
}
impl ConfirmsPublisher {
pub(crate) fn new(connected: &ConnectedLapinBroker, options: PublishOptions) -> Self {
Self {
link: connected.link().clone(),
options,
channel: Arc::new(ChannelCell::new()),
txn: Arc::new(Mutex::new(None)),
}
}
async fn channel(&self, conn: &AmqpConnection, target: &str) -> Result<Channel, AmqpError> {
self.channel
.get(|| async {
let channel = conn
.live_connection(target)?
.create_channel()
.await
.map_err(AmqpError::publish)?;
channel
.confirm_select(ConfirmSelectOptions::default())
.await
.map_err(AmqpError::publish)?;
Ok(channel)
})
.await
}
pub(crate) async fn flush_owned(&self, buffered: &[Buffered]) -> Result<(), AmqpError> {
let Some(first) = buffered.first() else {
return Ok(());
};
let conn = match &self.link {
Link::Amqp(conn) => conn,
#[cfg(feature = "testing")]
Link::InProcess(bus) => return self.options.flush_in_process(bus, buffered),
};
conn.ensure_live(&first.routing_key)?;
let channel = self.channel(conn, &first.routing_key).await?;
let mut confirms = Vec::with_capacity(buffered.len());
for entry in buffered {
let properties = convert::properties_for_publish(
&entry.headers,
&self.options.resolve(Some(&entry.options)),
)?;
let confirm = do_publish(
&channel,
&self.options.exchange,
&entry.routing_key,
&entry.payload,
properties,
)
.await?;
confirms.push((&entry.routing_key, confirm));
}
for (routing_key, confirm) in confirms {
let confirmation = confirm.await.map_err(AmqpError::publish)?;
confirmation_ok(&confirmation, routing_key)?;
}
Ok(())
}
async fn publish_confirmed(
&self,
routing_key: &str,
payload: &[u8],
headers: &HeaderMap,
options: Option<&LapinPublishOptions>,
) -> Result<(), AmqpError> {
let conn = match &self.link {
Link::Amqp(conn) => conn,
#[cfg(feature = "testing")]
Link::InProcess(bus) => {
return bus.publish(
&self.options.exchange,
routing_key,
payload,
headers,
&self.options.resolve(options),
);
}
};
conn.ensure_live(routing_key)?;
let channel = self.channel(conn, routing_key).await?;
let properties = convert::properties_for_publish(headers, &self.options.resolve(options))?;
let confirm = do_publish(
&channel,
&self.options.exchange,
routing_key,
payload,
properties,
)
.await?
.await
.map_err(AmqpError::publish)?;
confirmation_ok(&confirm, routing_key)
}
}
fn confirmation_ok(confirmation: &Confirmation, routing_key: &str) -> Result<(), AmqpError> {
if confirmation.is_nack() {
return Err(AmqpError::Publish(
format!("the broker negatively confirmed the publish to {routing_key:?}").into(),
));
}
Ok(())
}
impl Publisher for ConfirmsPublisher {
type Payload = Lend;
type Error = AmqpError;
type Options = LapinPublishOptions;
async fn publish(
&self,
msg: OutgoingMessage<'_>,
options: Option<&Self::Options>,
) -> Result<(), Self::Error> {
{
let mut txn = self.txn.lock().expect("transaction buffer mutex poisoned");
if let Some(buffer) = txn.as_mut() {
buffer.push(Buffered::lent(msg, options));
return Ok(());
}
}
self.publish_confirmed(msg.name(), msg.payload(), msg.headers(), options)
.await
}
}
impl TransactionalPublisher for ConfirmsPublisher {
fn begin_transaction(&self) -> impl Future<Output = 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 ready(Err(AmqpError::Transaction(
"a transaction is already open on this confirms publisher; commit or abort it \
before beginning another"
.to_owned(),
)));
}
ready(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 confirms publisher".to_owned(),
));
};
if buffered.is_empty() {
return Ok(());
}
let conn = match &self.link {
Link::Amqp(conn) => conn,
#[cfg(feature = "testing")]
Link::InProcess(bus) => return self.options.flush_in_process(bus, &buffered),
};
let target = buffered[0].routing_key.as_str();
conn.ensure_live(target)?;
let channel = self.channel(conn, target).await?;
let mut confirms = Vec::with_capacity(buffered.len());
for entry in &buffered {
let properties = convert::properties_for_publish(
&entry.headers,
&self.options.resolve(Some(&entry.options)),
)?;
let confirm = do_publish(
&channel,
&self.options.exchange,
&entry.routing_key,
&entry.payload,
properties,
)
.await?;
confirms.push((&entry.routing_key, confirm));
}
for (routing_key, confirm) in confirms {
let confirmation = confirm.await.map_err(AmqpError::publish)?;
confirmation_ok(&confirmation, routing_key)?;
}
Ok(())
}
fn abort(&self) -> impl Future<Output = Result<(), Self::Error>> {
let discarded = self
.txn
.lock()
.expect("transaction buffer mutex poisoned")
.take();
if discarded.is_none() {
return ready(Err(AmqpError::Transaction(
"abort with no open transaction on this confirms publisher".to_owned(),
)));
}
ready(Ok(()))
}
}
#[derive(Debug, Clone)]
pub struct ServerTxPublisher {
link: Link,
options: PublishOptions,
channel: Arc<ChannelCell<Channel>>,
open: Arc<Mutex<bool>>,
#[cfg(feature = "testing")]
staged: Arc<Mutex<Vec<Buffered>>>,
}
impl ServerTxPublisher {
pub(crate) fn new(connected: &ConnectedLapinBroker, options: PublishOptions) -> Self {
Self {
link: connected.link().clone(),
options,
channel: Arc::new(ChannelCell::new()),
open: Arc::new(Mutex::new(false)),
#[cfg(feature = "testing")]
staged: Arc::new(Mutex::new(Vec::new())),
}
}
async fn tx_channel(&self, conn: &AmqpConnection, target: &str) -> Result<Channel, AmqpError> {
if self.is_open() {
return self.channel.held().ok_or_else(|| {
AmqpError::Transaction(
"the transactional channel is gone; the open transaction cannot be published \
into or committed"
.to_owned(),
)
});
}
self.channel
.get(|| async {
let channel = conn
.live_connection(target)?
.create_channel()
.await
.map_err(AmqpError::publish)?;
channel.tx_select().await.map_err(AmqpError::publish)?;
Ok(channel)
})
.await
}
fn is_open(&self) -> bool {
*self.open.lock().expect("transaction state mutex poisoned")
}
fn set_open(&self, open: bool) {
*self.open.lock().expect("transaction state mutex poisoned") = open;
}
}
impl Publisher for ServerTxPublisher {
type Payload = Lend;
type Error = AmqpError;
type Options = LapinPublishOptions;
async fn publish(
&self,
msg: OutgoingMessage<'_>,
options: Option<&Self::Options>,
) -> Result<(), Self::Error> {
let conn = match &self.link {
Link::Amqp(conn) => conn,
#[cfg(feature = "testing")]
Link::InProcess(bus) => {
let resolved = self.options.resolve(options);
if !self.is_open() {
return bus.publish(
&self.options.exchange,
msg.name(),
msg.payload(),
msg.headers(),
&resolved,
);
}
bus.check(&self.options.exchange, msg.name(), msg.headers(), &resolved)?;
self.staged
.lock()
.expect("staged transaction mutex poisoned")
.push(Buffered::lent(msg, options));
return Ok(());
}
};
let properties =
convert::properties_for_publish(msg.headers(), &self.options.resolve(options))?;
let channel = if self.is_open() {
conn.ensure_live(msg.name())?;
self.tx_channel(conn, msg.name()).await?.clone()
} else {
conn.live_publish_channel(msg.name()).await?
};
let _confirm = do_publish(
&channel,
&self.options.exchange,
msg.name(),
msg.payload(),
properties,
)
.await?;
Ok(())
}
}
const TX_TARGET: &str = "the transactional channel";
impl TransactionalPublisher for ServerTxPublisher {
async fn begin_transaction(&self) -> Result<(), Self::Error> {
if self.is_open() {
return Err(AmqpError::Transaction(
"a transaction is already open on this server-transactional publisher; commit or \
abort it before beginning another"
.to_owned(),
));
}
match &self.link {
Link::Amqp(conn) => {
conn.ensure_live(TX_TARGET)?;
self.tx_channel(conn, TX_TARGET).await?;
}
#[cfg(feature = "testing")]
Link::InProcess(bus) => bus.ensure_live(TX_TARGET)?,
}
self.set_open(true);
Ok(())
}
async fn commit(&self) -> Result<(), Self::Error> {
if !self.is_open() {
return Err(AmqpError::Transaction(
"commit with no open transaction on this server-transactional publisher".to_owned(),
));
}
let conn = match &self.link {
Link::Amqp(conn) => conn,
#[cfg(feature = "testing")]
Link::InProcess(bus) => {
bus.ensure_live(TX_TARGET)?;
let staged = std::mem::take(
&mut *self
.staged
.lock()
.expect("staged transaction mutex poisoned"),
);
self.set_open(false);
for entry in &staged {
Bus::check_size(&entry.routing_key, entry.payload.len())?;
}
return self.options.flush_in_process(bus, &staged);
}
};
conn.ensure_live(TX_TARGET)?;
let channel = self.tx_channel(conn, TX_TARGET).await?;
channel.tx_commit().await.map_err(AmqpError::publish)?;
self.set_open(false);
Ok(())
}
async fn abort(&self) -> Result<(), Self::Error> {
if !self.is_open() {
return Err(AmqpError::Transaction(
"abort with no open transaction on this server-transactional publisher".to_owned(),
));
}
let conn = match &self.link {
Link::Amqp(conn) => conn,
#[cfg(feature = "testing")]
Link::InProcess(bus) => {
bus.ensure_live(TX_TARGET)?;
self.staged
.lock()
.expect("staged transaction mutex poisoned")
.clear();
self.set_open(false);
return Ok(());
}
};
conn.ensure_live(TX_TARGET)?;
let channel = self.tx_channel(conn, TX_TARGET).await?;
channel.tx_rollback().await.map_err(AmqpError::publish)?;
self.set_open(false);
Ok(())
}
}
#[cfg(test)]
mod tests {
use ruststream::{BytesMut, OutgoingFor, OutgoingMessage, Take};
use super::Buffered;
#[test]
fn a_buffered_message_keeps_the_buffer_the_framework_wrote() {
let payload = BytesMut::from(&br#"{"id":1}"#[..]);
let written = payload.as_ptr();
let msg: OutgoingFor<'_, Take> = OutgoingMessage::produced("orders", payload);
let buffered = Buffered::taken(msg, None);
assert_eq!(
buffered.payload.as_ptr(),
written,
"a transaction keeps the message until the commit, so it takes the buffer rather \
than copying out of it"
);
}
}