use std::fmt;
use std::time::Duration;
use bytes::Bytes;
use lapin::Channel;
use lapin::options::BasicPublishOptions;
use lapin::types::ShortString;
use ruststream::runtime::RETRY_COUNT_HEADER;
use ruststream::{HeaderMap, Str};
use crate::convert;
use crate::error::AmqpError;
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum Delay {
DlxTtl {
waiting_queue: Option<String>,
},
#[cfg(feature = "plugin-dme")]
DelayedMessageExchange {
exchange: Option<String>,
},
}
impl Delay {
#[must_use]
pub const fn dlx_ttl() -> Self {
Self::DlxTtl {
waiting_queue: None,
}
}
#[must_use]
pub fn dlx_ttl_named(name: impl Into<String>) -> Self {
Self::DlxTtl {
waiting_queue: Some(name.into()),
}
}
#[cfg(feature = "plugin-dme")]
#[must_use]
pub const fn plugin_dme() -> Self {
Self::DelayedMessageExchange { exchange: None }
}
#[cfg(feature = "plugin-dme")]
#[must_use]
pub fn plugin_dme_named(name: impl Into<String>) -> Self {
Self::DelayedMessageExchange {
exchange: Some(name.into()),
}
}
pub(crate) fn target_for(&self, origin: &str) -> DelayTarget {
match self {
Self::DlxTtl {
waiting_queue: Some(name),
} => DelayTarget::WaitingQueue {
waiting_queue: name.clone(),
},
Self::DlxTtl {
waiting_queue: None,
} => DelayTarget::WaitingQueue {
waiting_queue: format!("{origin}.retry"),
},
#[cfg(feature = "plugin-dme")]
Self::DelayedMessageExchange { exchange } => DelayTarget::DelayedExchange {
exchange: exchange
.clone()
.unwrap_or_else(|| format!("{origin}.delay")),
routing_key: origin.to_owned(),
},
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum DelayTarget {
WaitingQueue { waiting_queue: String },
#[cfg(feature = "plugin-dme")]
DelayedExchange {
exchange: String,
routing_key: String,
},
}
#[derive(Clone)]
pub(crate) struct DelayContext {
channel: Channel,
target: DelayTarget,
}
impl DelayContext {
pub(crate) fn new(channel: Channel, target: DelayTarget) -> Self {
Self { channel, target }
}
pub(crate) async fn republish(
&self,
payload: &[u8],
headers: &HeaderMap,
delay: Duration,
) -> Result<(), AmqpError> {
let options = convert::redelivery_options(headers);
let headers = &counted_again(headers);
match &self.target {
DelayTarget::WaitingQueue { waiting_queue } => {
let properties = convert::properties_for_publish(headers, &options)?
.with_expiration(convert::expiration_millis(delay));
self.channel
.basic_publish(
ShortString::default(),
convert::short(waiting_queue, "waiting queue name")?,
BasicPublishOptions::default(),
payload,
properties,
)
.await
.map_err(AmqpError::publish)?;
}
#[cfg(feature = "plugin-dme")]
DelayTarget::DelayedExchange {
exchange,
routing_key,
} => {
use lapin::types::{AMQPValue, FieldTable};
let mut properties = convert::properties_for_publish(headers, &options)?;
let mut table = properties
.headers()
.clone()
.unwrap_or_else(FieldTable::default);
let millis = i64::try_from(delay.as_millis()).unwrap_or(i64::MAX);
table.insert(ShortString::from("x-delay"), AMQPValue::LongLongInt(millis));
properties = properties.with_headers(table);
self.channel
.basic_publish(
convert::short(exchange, "delayed exchange name")?,
convert::short(routing_key, "routing key")?,
BasicPublishOptions::default(),
payload,
properties,
)
.await
.map_err(AmqpError::publish)?;
}
}
Ok(())
}
}
pub(crate) fn counted_again(headers: &HeaderMap) -> HeaderMap {
let spent = headers
.get_str(RETRY_COUNT_HEADER)
.and_then(|count| count.parse::<u64>().ok())
.unwrap_or(0);
let mut headers = headers.clone();
headers.insert(
Str::from_static(RETRY_COUNT_HEADER),
Bytes::from(spent.saturating_add(1).to_string()),
);
headers
}
impl fmt::Debug for DelayContext {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("DelayContext")
.field("target", &self.target)
.finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests {
use super::{Delay, DelayTarget};
#[test]
fn dlx_ttl_target_defaults_to_origin_dot_retry() {
assert_eq!(
Delay::dlx_ttl().target_for("orders"),
DelayTarget::WaitingQueue {
waiting_queue: "orders.retry".to_owned()
}
);
assert_eq!(
Delay::dlx_ttl_named("orders.wait").target_for("orders"),
DelayTarget::WaitingQueue {
waiting_queue: "orders.wait".to_owned()
}
);
}
#[cfg(feature = "plugin-dme")]
#[test]
fn dme_target_defaults_to_origin_dot_delay() {
assert_eq!(
Delay::plugin_dme().target_for("orders"),
DelayTarget::DelayedExchange {
exchange: "orders.delay".to_owned(),
routing_key: "orders".to_owned(),
}
);
}
}