use std::{
sync::{
Arc,
atomic::{AtomicU64, Ordering},
},
time::{Duration, SystemTime, UNIX_EPOCH},
};
use async_trait::async_trait;
use futures::StreamExt;
use lapin::{
Channel, Connection, Consumer,
message::Delivery as LapinDelivery,
options::{BasicAckOptions, BasicConsumeOptions, BasicQosOptions, BasicRejectOptions},
types::{FieldTable, MAX_SHORT_STRING_LENGTH},
};
use queuey_core::{Backend, Delivery, DeliveryStream, Envelope, QueueConfig, Result};
use tracing::{debug, error, info, warn};
use crate::{
codec,
connection::{ConnectionHandle, REPLY_SUCCESS, declare_topology},
delivery::RabbitMqDelivery,
error::{RabbitMqError, amqp, short_string},
options::RabbitMqOptions,
publisher::{Hold, Publisher},
reconnect::{Attempt, Rebuilding},
topology,
};
#[derive(Debug)]
pub struct RabbitMqBackend {
connection: Arc<ConnectionHandle>,
publisher: Publisher,
options: Arc<RabbitMqOptions>,
}
impl RabbitMqBackend {
pub async fn connect(uri: &str) -> Result<Self> {
Self::with_options(uri, RabbitMqOptions::default()).await
}
pub async fn with_options(uri: &str, options: RabbitMqOptions) -> Result<Self> {
let options = Arc::new(options);
let connection = ConnectionHandle::connect(uri, Arc::clone(&options)).await?;
let live = connection.ensure_connected().await?;
let channel = Publisher::open_confirm_channel(&live).await?;
let declare_channel = live.create_channel().await.map_err(amqp)?;
info!(
channel = channel.id(),
declare_channel = declare_channel.id(),
reconnects = connection.reconnects(),
"rabbitmq backend connected"
);
Ok(Self {
publisher: Publisher::new(
Arc::clone(&connection),
channel,
declare_channel,
Arc::clone(&options),
),
connection,
options,
})
}
#[must_use]
pub fn is_connected(&self) -> bool {
self.connection.is_connected()
}
#[must_use]
pub fn options(&self) -> &RabbitMqOptions {
&self.options
}
#[must_use]
pub fn dead_queue_name(&self, queue: &str) -> String {
topology::dead_queue_name(queue, &self.options.dead_suffix)
}
#[must_use]
pub fn deferred_queue_name(&self, queue: &str, ttl_ms: u32) -> String {
topology::deferred_queue_name(queue, &self.options.deferred_suffix, ttl_ms)
}
async fn declare_one(&self, channel: &Channel, config: &QueueConfig) -> Result<()> {
declare_topology(channel, config, &self.options).await?;
self.connection.remember(config);
debug!(queue = %config.name, "topology declared");
Ok(())
}
}
#[async_trait]
impl Backend for RabbitMqBackend {
async fn declare(&self, queues: &[QueueConfig]) -> Result<()> {
if queues.is_empty() {
return Ok(());
}
for config in queues {
check_deferrable_name(&config.name, &self.options.deferred_suffix)?;
}
let channel = self
.connection
.ensure_connected()
.await?
.create_channel()
.await
.map_err(amqp)?;
let result: Result<()> = async {
for config in queues {
self.declare_one(&channel, config).await?;
}
Ok(())
}
.await;
if channel.status().connected()
&& let Err(error) = channel.close(REPLY_SUCCESS, "OK".into()).await
{
debug!(%error, "closing the declaration channel failed");
}
result
}
async fn publish(&self, envelope: &Envelope, delay: Option<Duration>) -> Result<()> {
match delay {
None => self.publisher.publish_envelope(envelope).await,
Some(delay) => {
self.publisher
.publish_held(envelope, delay, Hold::Retry)
.await
}
}
}
async fn defer(&self, envelope: &Envelope, delay: Duration) -> Result<()> {
self.publisher
.publish_held(envelope, delay, Hold::Deferral)
.await
}
async fn consume(&self, queue: &QueueConfig) -> Result<DeliveryStream> {
let connection = self.connection.ensure_connected().await?;
let generation = self.connection.generation();
let (channel, consumer) = subscribe(&connection, queue).await?;
let stream: DeliveryStream = Box::pin(delivery_stream(ConsumeState {
consumer,
channel,
connection: Arc::clone(&self.connection),
publisher: self.publisher.clone(),
config: queue.clone(),
generation,
}));
Ok(stream)
}
async fn close(&self) -> Result<()> {
self.connection.mark_closing();
if let Err(error) = self.publisher.close().await {
debug!(%error, "closing the publishing channel failed");
}
self.connection.close().await?;
info!("rabbitmq backend closed");
Ok(())
}
}
fn check_deferrable_name(queue: &str, deferred_suffix: &str) -> Result<()> {
let hold = topology::deferred_queue_name(queue, deferred_suffix, topology::MAX_DEFERRAL_MS);
if hold.len() <= MAX_SHORT_STRING_LENGTH {
return Ok(());
}
Err(RabbitMqError::DeferredNameTooLong {
queue: queue.to_owned(),
length: hold.len(),
hold,
limit: MAX_SHORT_STRING_LENGTH,
}
.into_core())
}
pub(crate) fn is_benign_close_error(error: &lapin::Error) -> bool {
use lapin::{ChannelState, ConnectionState, ErrorKind};
matches!(
error.kind(),
ErrorKind::InvalidChannelState(ChannelState::Closing | ChannelState::Closed, _)
| ErrorKind::InvalidConnectionState(ConnectionState::Closing | ConnectionState::Closed)
)
}
async fn subscribe(connection: &Connection, config: &QueueConfig) -> Result<(Channel, Consumer)> {
let channel = connection.create_channel().await.map_err(amqp)?;
channel
.basic_qos(config.prefetch, BasicQosOptions { global: false })
.await
.map_err(amqp)?;
let tag = consumer_tag(&config.name);
let consumer = channel
.basic_consume(
short_string(&config.name)?,
short_string(&tag)?,
BasicConsumeOptions {
no_local: false,
no_ack: false,
exclusive: false,
nowait: false,
},
FieldTable::default(),
)
.await
.map_err(amqp)?;
info!(queue = %config.name, prefetch = config.prefetch, %tag, "consuming");
Ok((channel, consumer))
}
struct ConsumeState {
consumer: Consumer,
channel: Channel,
connection: Arc<ConnectionHandle>,
publisher: Publisher,
config: QueueConfig,
generation: u64,
}
enum Resubscribed {
Yes,
Stop,
}
fn delivery_stream(
state: ConsumeState,
) -> impl futures::Stream<Item = Result<Box<dyn Delivery>>> + Send {
futures::stream::unfold(state, |mut state| async move {
loop {
let next = match state.consumer.next().await {
Some(Ok(delivery)) => Some(delivery),
Some(Err(error)) => {
warn!(
queue = %state.config.name,
%error,
"consumer failed"
);
None
}
None => None,
};
let Some(delivery) = next else {
match recover(&mut state).await {
Ok(Resubscribed::Yes) => continue,
Ok(Resubscribed::Stop) => return None,
Err(error) => return Some((Err(error), state)),
}
};
match Envelope::from_bytes(&delivery.data) {
Ok(envelope) => {
let boxed: Box<dyn Delivery> = Box::new(RabbitMqDelivery::new(
envelope,
delivery.acker.clone(),
state.publisher.clone(),
));
return Some((Ok(boxed), state));
}
Err(error) => {
warn!(
queue = %state.config.name,
delivery_tag = delivery.delivery_tag,
bytes = delivery.data.len(),
%error,
"dropping message whose body is not a valid envelope"
);
discard_malformed(&state.publisher, &state.config.name, &delivery).await;
}
}
}
})
}
async fn recover(state: &mut ConsumeState) -> Result<Resubscribed> {
if state.connection.is_closing() {
debug!(queue = %state.config.name, "consumer stopped; the backend is closing");
return Ok(Resubscribed::Stop);
}
let Some(policy) = state.connection.policy() else {
info!(
queue = %state.config.name,
"consumer stopped and reconnection is disabled; ending the stream"
);
return Ok(Resubscribed::Stop);
};
close_consumer_channel(&state.channel).await;
let mut failures: u32 = 0;
loop {
let connection = state.connection.ensure_connected().await?;
let generation = state.connection.generation();
match subscribe(&connection, &state.config).await {
Ok((channel, consumer)) => {
info!(
queue = %state.config.name,
generation,
previous_generation = state.generation,
attempts = failures + 1,
"consumer resubscribed"
);
state.channel = channel;
state.consumer = consumer;
state.generation = generation;
return Ok(Resubscribed::Yes);
}
Err(error) => {
failures += 1;
let Some(delay) =
policy.next_delay(Attempt::after(Rebuilding::Consumer, failures, &error))
else {
error!(
queue = %state.config.name,
%error,
attempts = failures,
"giving up on resubscribing the consumer; the policy declined another \
attempt"
);
return Err(error);
};
warn!(
queue = %state.config.name,
%error,
failures,
?delay,
"resubscribing the consumer failed; will try again"
);
state.connection.sleep_unless_closing(delay).await?;
}
}
}
}
async fn close_consumer_channel(channel: &Channel) {
if !channel.status().connected() {
return;
}
if let Err(error) = channel.close(REPLY_SUCCESS, "OK".into()).await {
debug!(%error, channel = channel.id(), "closing a replaced consumer channel failed");
}
}
async fn discard_malformed(publisher: &Publisher, queue: &str, delivery: &LapinDelivery) {
if publisher.options().declare_dead_letter_queues {
match publisher
.publish_malformed(queue, &delivery.data, codec::REASON_MALFORMED)
.await
{
Ok(()) => match delivery.acker.ack(BasicAckOptions::default()).await {
Ok(true) => return,
Ok(false) => {
warn!(
queue,
delivery_tag = delivery.delivery_tag,
"a malformed message was already settled; nothing left to ack"
);
return;
}
Err(error) => {
warn!(%error, queue, delivery_tag = delivery.delivery_tag,
"acking a malformed message failed; falling back to reject");
}
},
Err(error) => {
warn!(%error, queue, "forwarding a malformed message to the dead-letter queue failed");
}
}
}
match delivery
.acker
.reject(BasicRejectOptions { requeue: false })
.await
{
Ok(true) => {}
Ok(false) => {
warn!(
queue,
delivery_tag = delivery.delivery_tag,
"a malformed message was already settled; nothing left to reject"
);
}
Err(error) => {
error!(
%error, queue, delivery_tag = delivery.delivery_tag,
"rejecting a malformed message failed; it stays unacknowledged and holds a \
prefetch slot until the consumer channel closes"
);
}
}
}
fn consumer_tag(queue: &str) -> String {
static NEXT: AtomicU64 = AtomicU64::new(0);
let sequence = NEXT.fetch_add(1, Ordering::Relaxed);
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|since| since.as_nanos())
.unwrap_or_default();
let queue = codec::truncate_at_boundary(queue, 160);
format!("queuey.{queue}.{nanos:x}.{sequence}")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn closing_state_errors_are_benign_on_close() {
use lapin::{ChannelState, ConnectionState, ErrorKind};
for kind in [
ErrorKind::InvalidChannelState(ChannelState::Closing, "channel.close"),
ErrorKind::InvalidChannelState(ChannelState::Closed, "channel.close"),
ErrorKind::InvalidConnectionState(ConnectionState::Closing),
ErrorKind::InvalidConnectionState(ConnectionState::Closed),
] {
assert!(is_benign_close_error(&lapin::Error::from(kind)));
}
}
#[test]
fn other_state_errors_are_not_benign_on_close() {
use lapin::{ChannelState, ConnectionState, ErrorKind};
for kind in [
ErrorKind::InvalidChannelState(ChannelState::Initial, "channel.close"),
ErrorKind::InvalidChannelState(ChannelState::Error, "channel.close"),
ErrorKind::InvalidConnectionState(ConnectionState::Error),
ErrorKind::InvalidChannel(7),
] {
assert!(!is_benign_close_error(&lapin::Error::from(kind)));
}
}
#[test]
fn consumer_tags_are_unique_and_name_the_queue() {
let first = consumer_tag("myapp.emails");
let second = consumer_tag("myapp.emails");
assert_ne!(first, second);
assert!(first.starts_with("queuey.myapp.emails."));
assert!(second.starts_with("queuey.myapp.emails."));
}
#[test]
fn consumer_tags_fit_in_a_short_string() {
let tag = consumer_tag(&"q".repeat(1000));
assert!(
tag.len() <= MAX_SHORT_STRING_LENGTH,
"len was {}",
tag.len()
);
assert!(short_string(&tag).is_ok());
}
#[test]
fn a_queue_name_with_room_for_its_hold_queues_is_accepted() {
assert!(check_deferrable_name("myapp.emails", topology::DEFAULT_DEFERRED_SUFFIX).is_ok());
let longest = "q".repeat(MAX_SHORT_STRING_LENGTH - ".deferred.".len() - 10);
assert_eq!(longest.len(), 235);
assert!(check_deferrable_name(&longest, topology::DEFAULT_DEFERRED_SUFFIX).is_ok());
}
#[test]
fn a_queue_name_that_leaves_no_room_for_hold_queues_is_refused_at_declare() {
let name = "q".repeat(250);
assert!(
short_string(&name).is_ok(),
"the queue name itself is legal"
);
let error = check_deferrable_name(&name, topology::DEFAULT_DEFERRED_SUFFIX)
.expect_err("a name with no room for hold queues must be refused");
let text = error.to_string();
assert!(
text.contains("leaves no room for its hold queues"),
"{text}"
);
assert!(text.contains("270 bytes"), "{text}");
assert!(text.contains("255-byte"), "{text}");
}
#[test]
fn the_hold_queue_name_check_uses_the_configured_suffix() {
let name = "q".repeat(240);
assert!(check_deferrable_name(&name, topology::DEFAULT_DEFERRED_SUFFIX).is_err());
assert!(check_deferrable_name(&name, "-h").is_ok());
}
#[test]
fn declare_options_never_auto_delete() {
let durable = topology::declare_options(true);
assert!(durable.durable);
assert!(!durable.auto_delete);
assert!(!durable.exclusive);
assert!(!durable.passive);
assert!(!durable.nowait);
let transient = topology::declare_options(false);
assert!(!transient.durable);
assert!(!transient.auto_delete);
}
}