#![cfg_attr(
not(feature = "testing"),
allow(clippy::infallible_destructuring_match)
)]
#[cfg(feature = "testing")]
use std::future::{Future, ready};
use std::num::NonZeroU16;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use lapin::options::{BasicConsumeOptions, BasicQosOptions};
use lapin::types::{FieldTable, ShortString};
#[cfg(feature = "testing")]
use lapin::uri::AMQPUri;
use lapin::{Channel, ChannelState, Connection, ConnectionProperties, ConnectionState, ErrorKind};
#[cfg(feature = "testing")]
use ruststream::testing::{Coordinator, InProcess, TestableBroker};
use ruststream::{
AddressedCopies, Broker, ConnectedBroker, DefaultPublish, DescribeServer, ServerSpec, Subscribe,
};
#[cfg(feature = "testing")]
use ruststream::{OutgoingMessage, RawMessage};
use tokio::runtime::Handle;
use crate::channel::ChannelCell;
use crate::convert;
use crate::delay::DelayContext;
use crate::error::AmqpError;
#[cfg(feature = "testing")]
use crate::in_process::{self, Bus};
use crate::publish_policy::{LapinPublish, LapinPublishPolicy};
#[cfg(feature = "testing")]
use crate::publish_step::LapinPublishOptions;
use crate::queue::{QueueDescriptor, RabbitQueue, declared_arguments};
use crate::requester::{LapinRequest, LapinRequester};
use crate::subscriber::LapinSubscriber;
use crate::topology;
pub(crate) const PROTOCOL: &str = "amqp";
pub(crate) const PROTOCOL_VERSION: &str = "0.9.1";
pub(crate) struct AmqpConnection {
connection: Connection,
publish_channel: ChannelCell<Channel>,
closed: AtomicBool,
runtime: Handle,
}
impl AmqpConnection {
fn new(connection: Connection, publish_channel: Channel, runtime: Handle) -> Arc<Self> {
Arc::new(Self {
connection,
publish_channel: ChannelCell::holding(publish_channel),
closed: AtomicBool::new(false),
runtime,
})
}
pub(crate) const fn runtime(&self) -> &Handle {
&self.runtime
}
pub(crate) fn live_connection(&self, target: &str) -> Result<&Connection, AmqpError> {
self.ensure_live(target)?;
Ok(&self.connection)
}
pub(crate) async fn live_publish_channel(&self, target: &str) -> Result<Channel, AmqpError> {
self.ensure_live(target)?;
self.publish_channel
.get(|| async {
self.connection
.create_channel()
.await
.map_err(AmqpError::publish)
})
.await
}
pub(crate) fn is_closed(&self) -> bool {
self.closed.load(Ordering::Acquire)
}
pub(crate) fn ensure_live(&self, target: &str) -> Result<(), AmqpError> {
if self.is_closed() {
return Err(AmqpError::closed(target));
}
Ok(())
}
}
impl std::fmt::Debug for AmqpConnection {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AmqpConnection")
.field("closed", &self.closed.load(Ordering::Relaxed))
.finish_non_exhaustive()
}
}
#[derive(Debug, Clone)]
pub(crate) enum Link {
Amqp(Arc<AmqpConnection>),
#[cfg(feature = "testing")]
InProcess(Arc<Bus>),
}
#[cfg(not(feature = "testing"))]
const _: () = assert!(size_of::<Link>() == size_of::<Arc<AmqpConnection>>());
#[derive(Debug, Clone)]
#[must_use]
pub struct LapinBroker {
uri: String,
connection_name: Option<String>,
prefetch: Option<NonZeroU16>,
declare: bool,
}
impl LapinBroker {
pub fn new(uri: impl Into<String>) -> Self {
Self {
uri: uri.into(),
connection_name: None,
prefetch: None,
declare: false,
}
}
pub fn connection_name(mut self, name: impl Into<String>) -> Self {
self.connection_name = Some(name.into());
self
}
pub fn prefetch(mut self, prefetch: NonZeroU16) -> Self {
self.prefetch = Some(prefetch);
self
}
pub fn declare_topology(mut self, declare: bool) -> Self {
self.declare = declare;
self
}
}
impl Broker for LapinBroker {
type Error = AmqpError;
type Connected = ConnectedLapinBroker;
async fn connect(self) -> Result<Self::Connected, Self::Error> {
let mut properties = ConnectionProperties::default();
if let Some(name) = &self.connection_name {
properties = properties.with_connection_name(name.as_str().into());
}
let connection = Connection::connect(&self.uri, properties)
.await
.map_err(AmqpError::connect)?;
let publish_channel = connection
.create_channel()
.await
.map_err(AmqpError::connect)?;
Ok(ConnectedLapinBroker {
link: Link::Amqp(AmqpConnection::new(
connection,
publish_channel,
Handle::current(),
)),
uri: self.uri,
prefetch: self.prefetch,
declare: self.declare,
})
}
}
#[cfg(feature = "testing")]
impl InProcess for LapinBroker {
fn connect_in_process(
self,
) -> impl Future<Output = Result<Self::Connected, Self::Error>> + Send {
let connected = self
.uri
.parse::<AMQPUri>()
.map_err(|err| AmqpError::Connect(err.into()))
.map(|_| ConnectedLapinBroker {
link: Link::InProcess(Bus::new(Handle::current())),
uri: self.uri,
prefetch: self.prefetch,
declare: self.declare,
});
ready(connected)
}
}
#[cfg(feature = "testing")]
ruststream::register_testable_broker!(LapinBroker);
impl DescribeServer for LapinBroker {
fn describe_server(&self) -> ServerSpec {
ServerSpec::from_url(&self.uri, PROTOCOL).protocol_version(PROTOCOL_VERSION)
}
}
#[derive(Debug)]
pub struct ConnectedLapinBroker {
link: Link,
uri: String,
prefetch: Option<NonZeroU16>,
declare: bool,
}
impl ConnectedLapinBroker {
pub(crate) fn link(&self) -> &Link {
&self.link
}
#[must_use]
pub fn server_spec(&self) -> ServerSpec {
ServerSpec::from_url(&self.uri, PROTOCOL).protocol_version(PROTOCOL_VERSION)
}
pub async fn subscribe(&self, def: impl QueueDescriptor) -> Result<LapinSubscriber, AmqpError> {
let conn = match &self.link {
Link::Amqp(conn) => conn,
#[cfg(feature = "testing")]
Link::InProcess(bus) => return in_process::subscribe(bus, &def, self.declare),
};
let spec = def.spec();
def.check_retry(self.declare)?;
let channel = conn
.live_connection(&spec.name)?
.create_channel()
.await
.map_err(AmqpError::subscribe)?;
if self.declare {
let arguments = declared_arguments(spec, def.declared_retry());
topology::declare(&channel, spec, arguments).await?;
}
if let Some(prefetch) = spec.prefetch_or(self.prefetch) {
channel
.basic_qos(prefetch.get(), BasicQosOptions::default())
.await
.map_err(AmqpError::subscribe)?;
}
let queue = spec.name.clone();
let delay = spec
.delay
.as_ref()
.map(|delay| DelayContext::new(channel.clone(), delay.target_for(&queue)));
let consumer = channel
.basic_consume(
convert::short(&queue, "queue name")?,
ShortString::default(),
BasicConsumeOptions::default(),
FieldTable::default(),
)
.await
.map_err(AmqpError::subscribe)?;
Ok(LapinSubscriber::new(
Arc::clone(conn),
channel,
consumer,
queue,
spec.batch_wait,
delay,
))
}
#[must_use]
pub fn publisher<P: LapinPublishPolicy>(&self, policy: P) -> P::Live {
policy.bind(self)
}
#[must_use]
pub fn requester(&self, policy: LapinRequest) -> LapinRequester {
policy.bind(self)
}
}
impl ConnectedBroker for ConnectedLapinBroker {
type Error = AmqpError;
type Closed = ClosedLapinBroker;
async fn shutdown(self) -> Result<Self::Closed, Self::Error> {
let conn = match self.link {
Link::Amqp(conn) => conn,
#[cfg(feature = "testing")]
Link::InProcess(bus) => {
bus.close();
return Ok(ClosedLapinBroker { handshake: true });
}
};
conn.closed.store(true, Ordering::Release);
let handshake = conn.connection.status().connected();
if handshake {
match conn.connection.close(200, ShortString::from("OK")).await {
Ok(()) => {}
Err(err) if already_closing(&err) => {}
Err(err) => return Err(AmqpError::connect(err)),
}
}
Ok(ClosedLapinBroker { handshake })
}
}
#[allow(clippy::use_self)]
impl Subscribe for ConnectedLapinBroker {
type Subscriber = LapinSubscriber;
type Copies = AddressedCopies;
async fn subscribe(&self, name: &str) -> Result<Self::Subscriber, Self::Error> {
ConnectedLapinBroker::subscribe(self, RabbitQueue::new(name)).await
}
}
impl DefaultPublish for ConnectedLapinBroker {
type Policy = LapinPublish;
}
#[cfg(feature = "testing")]
impl TestableBroker for ConnectedLapinBroker {
fn install_coordinator(&self, coordinator: Coordinator) {
if let Link::InProcess(bus) = &self.link {
bus.install(coordinator);
}
}
fn inject(&self, message: OutgoingMessage<'_>) {
let bus = self.bus("inject");
if let Err(err) = bus.publish(
"",
message.name(),
message.payload(),
message.headers(),
&LapinPublishOptions::default(),
) {
panic!(
"the injected message to {:?} is not one the server takes: {err}",
message.name()
);
}
}
fn published(&self, name: &str) -> Vec<RawMessage> {
self.bus("published").published(name)
}
fn routes(&self, destination: &str, subscriptions: &[&str]) -> Vec<usize> {
let candidates: Vec<usize> = subscriptions
.iter()
.enumerate()
.filter(|(_, name)| **name == destination)
.map(|(position, _)| position)
.collect();
let turn = match &self.link {
Link::InProcess(bus) => bus.next_consumer_of(destination),
Link::Amqp(_) => {
assert!(
candidates.len() <= 1,
"queue {destination:?} has {} subscriptions, and RabbitMQ picks which consumer \
takes a message where the harness cannot see it: a live test opens one \
subscription per queue",
candidates.len()
);
Some(0)
}
};
turn.and_then(|turn| candidates.get(turn).copied())
.into_iter()
.collect()
}
}
#[cfg(feature = "testing")]
impl ConnectedLapinBroker {
fn bus(&self, what: &str) -> &Arc<Bus> {
match &self.link {
Link::InProcess(bus) => bus,
Link::Amqp(_) => panic!(
"TestableBroker::{what} reached a broker connected with `connect`; the harness \
drives the connection `connect_in_process` produces"
),
}
}
}
fn already_closing(err: &lapin::Error) -> bool {
matches!(
err.kind(),
ErrorKind::InvalidChannelState(ChannelState::Closing | ChannelState::Closed, _)
| ErrorKind::InvalidConnectionState(ConnectionState::Closing | ConnectionState::Closed)
)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ClosedLapinBroker {
handshake: bool,
}
impl ClosedLapinBroker {
#[must_use]
pub const fn handshake(&self) -> bool {
self.handshake
}
}
#[cfg(test)]
mod tests {
use super::{DescribeServer, LapinBroker};
fn described_host(uri: &str) -> String {
LapinBroker::new(uri)
.describe_server()
.host
.expect("an AMQP URI always names a host")
}
#[test]
fn the_description_is_the_host_alone_whatever_the_uri_carries() {
assert_eq!(described_host("amqp://localhost:5672"), "localhost:5672");
assert_eq!(
described_host("amqp://user:pass@rabbit:5672/prod"),
"rabbit:5672"
);
assert_eq!(described_host("amqps://rabbit/vhost"), "rabbit");
assert_eq!(described_host("rabbit:5672"), "rabbit:5672");
assert_eq!(described_host("amqp://rabbit:5672/my@vhost"), "rabbit:5672");
assert_eq!(
described_host("amqp://user:p@ss@rabbit:5672/my@vhost"),
"rabbit:5672"
);
assert_eq!(
described_host("amqp://rabbit:5672/prod?heartbeat=30"),
"rabbit:5672"
);
}
#[test]
fn new_performs_no_io_and_describes_the_configured_address() {
let spec = LapinBroker::new("amqp://127.0.0.1:5672").describe_server();
assert_eq!(spec.protocol, "amqp");
assert_eq!(spec.host.as_deref(), Some("127.0.0.1:5672"));
}
}