ruststream_lapin/

requester.rs

//! Request/reply over `RabbitMQ` direct reply-to (`amq.rabbitmq.reply-to`).

// Without the `testing` feature a connection link has one variant, so a `match` on it has a
// single arm; the matches stay so that the in-process arm has its place when the feature is on.
#![cfg_attr(
    not(feature = "testing"),
    allow(clippy::infallible_destructuring_match)
)]

use std::collections::HashMap;
use std::future::{Future, ready};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, Weak};
use std::time::Duration;

use futures::StreamExt;
use lapin::Channel;
use lapin::options::{BasicConsumeOptions, BasicPublishOptions};
use lapin::types::{FieldTable, ShortString};
#[cfg(feature = "asyncapi")]
use ruststream::asyncapi::Bindings;
use ruststream::{Lend, OutgoingMessage, PairError, PublishPolicy, Publisher, RequestReply};
use tokio::sync::oneshot;

use crate::broker::{AmqpConnection, ConnectedLapinBroker, Link};
use crate::channel::{ChannelCell, Holds};
use crate::convert;
use crate::error::AmqpError;
#[cfg(feature = "testing")]
use crate::in_process;
use crate::message::LapinMessage;
use crate::publish_policy::{LapinPublishPolicy, PublishOptions, publish_policy_bindings};
use crate::publish_step::LapinPublishOptions;

/// The pseudo-queue `RabbitMQ` rewrites per-request for direct reply-to.
pub(crate) const REPLY_TO: &str = "amq.rabbitmq.reply-to";

type Pending = Mutex<HashMap<String, oneshot::Sender<LapinMessage>>>;

/// The request/reply policy: pure declaration, constructible anywhere, pairing into
/// [`LapinRequester`].
///
/// Requests are published transient (delivery mode 1) by default: a request nobody is waiting
/// for after the timeout gains nothing from surviving a broker restart. Opt into persistence
/// with [`persistent(true)`](Self::persistent).
///
/// # Examples
///
/// ```
/// use std::time::Duration;
///
/// use ruststream::OutgoingMessage;
/// use ruststream_lapin::prelude::*;
/// use serde::Deserialize;
///
/// #[derive(Deserialize)]
/// struct Order {
///     sku: String,
/// }
///
/// #[subscriber("orders")]
/// async fn place(order: &Order, Out(inventory): Out<impl RequestReply>) -> HandlerOutcome {
///     let ask = OutgoingMessage::new("inventory.check", order.sku.as_bytes());
///     match inventory.request(ask, Duration::from_secs(2)).await {
///         Ok(reply) if reply.payload() == b"in-stock" => HandlerOutcome::ack(),
///         Ok(_) => HandlerOutcome::drop(),
///         Err(_) => HandlerOutcome::retry(),
///     }
/// }
///
/// fn app() -> RustStream {
///     let broker = LapinBroker::new("amqp://localhost:5672");
///     RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(broker, |b| {
///         b.include(place)
///             .out(DefaultSlot, LapinRequest::default().exchange("rpc"))
///             .build();
///     })
/// }
/// ```
#[derive(Debug, Clone, PartialEq, Eq)]
#[must_use]
pub struct LapinRequest(PublishOptions);

impl Default for LapinRequest {
    fn default() -> Self {
        Self(PublishOptions {
            defaults: LapinPublishOptions {
                persistent: Some(false),
                ..LapinPublishOptions::default()
            },
            ..PublishOptions::default()
        })
    }
}

impl LapinRequest {
    /// What this policy hands the requester it pairs into.
    ///
    /// The requester takes the value by move at `bind`; this borrow is for the document, which
    /// reads it.
    #[cfg(feature = "asyncapi")]
    pub(crate) const fn publish_options(&self) -> &PublishOptions {
        &self.0
    }

    /// Publishes requests to `exchange` instead of the default exchange.
    pub fn exchange(mut self, exchange: impl Into<String>) -> Self {
        self.0.exchange = exchange.into();
        self
    }

    /// Whether requests are marked persistent (delivery mode 2). Defaults to `false`.
    ///
    /// One request departs from this with the publish builder's
    /// [`persistent`](crate::LapinPublishSteps::persistent) step.
    pub fn persistent(mut self, persistent: bool) -> Self {
        self.0.defaults.persistent = Some(persistent);
        self
    }

    /// The AMQP `priority` property requests carry. Unset by default.
    ///
    /// One request departs from this with the publish builder's
    /// [`priority`](crate::LapinPublishSteps::priority) step.
    pub fn priority(mut self, priority: u8) -> Self {
        self.0.defaults.priority = Some(priority);
        self
    }

    /// The AMQP per-message `expiration` (TTL) requests carry: the broker drops one once `ttl`
    /// has passed without it being consumed. Unset by default.
    ///
    /// Pair it with the per-request timeout: a request that expires on the broker fails with
    /// [`AmqpError::RequestTimeout`] at the caller, without a responder ever seeing it.
    ///
    /// One request departs from this with the publish builder's
    /// [`expiration`](crate::LapinPublishSteps::expiration) step.
    pub fn expiration(mut self, ttl: Duration) -> Self {
        self.0.defaults.expiration = Some(ttl);
        self
    }
}

impl PublishPolicy<ConnectedLapinBroker> for LapinRequest {
    type Live = LapinRequester;

    fn pair(
        self,
        connected: &ConnectedLapinBroker,
    ) -> impl Future<Output = Result<Self::Live, PairError>> {
        ready(Ok(self.bind(connected)))
    }

    publish_policy_bindings!();
}

impl LapinPublishPolicy for LapinRequest {
    fn bind(self, connected: &ConnectedLapinBroker) -> Self::Live {
        LapinRequester {
            link: connected.link().clone(),
            options: self.0,
            state: Arc::new(ChannelCell::new()),
            pending: Arc::new(Mutex::new(HashMap::new())),
            next_id: Arc::new(AtomicU64::new(0)),
        }
    }
}

/// The live request/reply client over `RabbitMQ` direct reply-to.
///
/// [`request`](RequestReply::request) publishes to the routing key named by
/// [`OutgoingMessage::name`] (on the default exchange unless the policy's
/// [`exchange`](LapinRequest::exchange) says otherwise) with `reply-to` set to the direct
/// reply-to pseudo-queue and a generated `correlation-id`; the responder replies by publishing
/// to the `reply-to` it received, echoing the `correlation-id`.
///
/// Direct reply-to is at-most-once: replies live in channel state on one broker node, so a
/// dropped requester channel loses in-flight replies. The per-request timeout is the recovery
/// mechanism.
///
/// Obtained by pairing [`LapinRequest`] with a connected broker
/// ([`ConnectedLapinBroker::requester`](crate::ConnectedLapinBroker::requester)). Clones share
/// the reply consumer and the pending-request table; like every live handle it aliases the
/// connection and reports [`AmqpError::Closed`] once the broker has shut down.
#[derive(Debug, Clone)]
pub struct LapinRequester {
    link: Link,
    options: PublishOptions,
    state: Arc<ChannelCell<ReqState>>,
    pending: Arc<Pending>,
    next_id: Arc<AtomicU64>,
}

/// The requester's channel and the reply consumer running on it.
#[derive(Debug, Clone)]
struct ReqState {
    channel: Channel,
}

impl Holds for ReqState {
    fn channel(&self) -> &Channel {
        &self.channel
    }
}

impl LapinRequester {
    /// Opens the requester channel and starts the reply consumer, and does it again once the
    /// broker closes that channel.
    ///
    /// The consumer MUST be up before the first publish carrying the direct reply-to address;
    /// `RabbitMQ` rejects such a publish with `PRECONDITION_FAILED` otherwise. Opening it on
    /// first use rather than at pairing time keeps pairing a synchronous constructor call and
    /// leaves an unused requester without a channel.
    ///
    /// A channel-level error closes the channel and takes the reply consumer with it. The
    /// requests in flight on it are lost, which is what the per-request timeout is for; what the
    /// requester must not do is stop working for every later request too.
    async fn state(&self, conn: &AmqpConnection, target: &str) -> Result<ReqState, AmqpError> {
        self.state
            .get(|| async {
                let channel = conn
                    .live_connection(target)?
                    .create_channel()
                    .await
                    .map_err(AmqpError::request)?;
                let consumer = channel
                    .basic_consume(
                        ShortString::from(REPLY_TO),
                        ShortString::default(),
                        BasicConsumeOptions {
                            no_ack: true,
                            ..BasicConsumeOptions::default()
                        },
                        FieldTable::default(),
                    )
                    .await
                    .map_err(AmqpError::request)?;

                // The task exits when the channel closes (consumer stream ends) or when every
                // requester clone is gone (Weak upgrade fails on the next reply). It runs on the
                // runtime the broker connected on: the request that starts it may come from a
                // handler on a dedicated thread, whose runtime may stop while later requests
                // still wait on this dispatcher.
                let pending = Arc::downgrade(&self.pending);
                conn.runtime().spawn(dispatch_replies(consumer, pending));

                Ok(ReqState { channel })
            })
            .await
    }
}

async fn dispatch_replies(mut consumer: lapin::Consumer, pending: Weak<Pending>) {
    while let Some(delivery) = consumer.next().await {
        let Ok(delivery) = delivery else {
            // The channel is failing; consuming further would spin. Outstanding requests fail
            // by timeout.
            return;
        };
        let Some(pending) = pending.upgrade() else {
            return;
        };
        let correlation_id = delivery
            .properties
            .correlation_id()
            .as_ref()
            .map(ShortString::as_str);
        let Some(correlation_id) = correlation_id else {
            tracing::debug!("dropping direct reply-to delivery without a correlation-id");
            continue;
        };
        let waiter = pending
            .lock()
            .expect("pending requests mutex poisoned")
            .remove(correlation_id);
        match waiter {
            // The receiver may have timed out concurrently; nothing to do then.
            Some(tx) => drop(tx.send(LapinMessage::from_delivery_no_ack(delivery))),
            None => {
                tracing::debug!(
                    correlation_id,
                    "dropping direct reply-to delivery with no waiter"
                );
            }
        }
    }
}

impl Publisher for LapinRequester {
    /// `basic_publish` takes the payload as `&[u8]` and copies it into the frame it sends, so
    /// this publisher reads the bytes and keeps nothing.
    type Payload = Lend;
    type Error = AmqpError;
    type Options = LapinPublishOptions;

    /// Publishes `msg` on the requester channel without expecting a reply.
    ///
    /// The AMQP basic properties are `options` resolved over the policy's defaults.
    ///
    /// # Errors
    ///
    /// Returns [`AmqpError::Closed`] once the broker has shut down and [`AmqpError::Publish`]
    /// when the channel rejects the frame.
    ///
    /// # Cancel safety
    ///
    /// Not cancel safe: dropping the future may leave the message published or not.
    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),
                );
            }
        };
        conn.ensure_live(msg.name())?;
        let state = self.state(conn, msg.name()).await?;
        let properties =
            convert::properties_for_publish(msg.headers(), &self.options.resolve(options))?;
        #[cfg(feature = "testing")]
        conn.routes().published(&self.options.exchange, msg.name());
        let _confirm = state
            .channel
            .basic_publish(
                convert::short(&self.options.exchange, "exchange name")?,
                convert::short(msg.name(), "routing key")?,
                BasicPublishOptions::default(),
                msg.payload(),
                properties,
            )
            .await
            .map_err(AmqpError::publish)?;
        Ok(())
    }
}

impl RequestReply for LapinRequester {
    type Reply = LapinMessage;

    /// Sends `msg` and awaits the correlated reply.
    ///
    /// # Errors
    ///
    /// Returns [`AmqpError::RequestTimeout`] when no reply arrives within `timeout`,
    /// [`AmqpError::Closed`] once the broker has shut down, and [`AmqpError::Request`] /
    /// [`AmqpError::Publish`] on channel failures.
    ///
    /// # Cancel safety
    ///
    /// Cancel safe for the caller's state: dropping the future abandons the pending slot and a
    /// late reply is discarded. The request itself may still have been published.
    async fn request(
        &self,
        msg: OutgoingMessage<'_>,
        timeout: Duration,
    ) -> Result<Self::Reply, Self::Error> {
        let conn = match &self.link {
            Link::Amqp(conn) => conn,
            #[cfg(feature = "testing")]
            Link::InProcess(bus) => {
                return in_process::request(bus, &self.options, msg, timeout).await;
            }
        };
        conn.ensure_live(msg.name())?;
        let state = self.state(conn, msg.name()).await?;

        let correlation_id = format!("rs-{}", self.next_id.fetch_add(1, Ordering::Relaxed));
        let (tx, rx) = oneshot::channel();
        {
            let mut pending = self
                .pending
                .lock()
                .expect("pending requests mutex poisoned");
            pending.insert(correlation_id.clone(), tx);
        }
        // Every failure path below must reclaim the slot, or it leaks until shutdown.
        let cleanup = || {
            let mut pending = self
                .pending
                .lock()
                .expect("pending requests mutex poisoned");
            pending.remove(&correlation_id);
        };

        let properties =
            match convert::properties_for_publish(msg.headers(), &self.options.resolve(None)) {
                Ok(properties) => properties
                    .with_reply_to(ShortString::from(REPLY_TO))
                    .with_correlation_id(ShortString::from(correlation_id.clone())),
                Err(err) => {
                    cleanup();
                    return Err(err);
                }
            };
        let exchange = match convert::short(&self.options.exchange, "exchange name") {
            Ok(exchange) => exchange,
            Err(err) => {
                cleanup();
                return Err(err);
            }
        };
        let routing_key = match convert::short(msg.name(), "routing key") {
            Ok(routing_key) => routing_key,
            Err(err) => {
                cleanup();
                return Err(err);
            }
        };

        #[cfg(feature = "testing")]
        conn.routes().published(&self.options.exchange, msg.name());
        let published = state
            .channel
            .basic_publish(
                exchange,
                routing_key,
                BasicPublishOptions::default(),
                msg.payload(),
                properties,
            )
            .await;
        if let Err(err) = published {
            cleanup();
            return Err(AmqpError::publish(err));
        }

        match tokio::time::timeout(timeout, rx).await {
            Ok(Ok(reply)) => Ok(reply),
            // The dispatch task dropped the sender: the reply channel died under us.
            Ok(Err(_)) => {
                cleanup();
                Err(AmqpError::Request(
                    "the reply consumer stopped before a reply arrived".into(),
                ))
            }
            Err(_) => {
                cleanup();
                Err(AmqpError::RequestTimeout(timeout))
            }
        }
    }
}