ruststream_nats/

message.rs

//! Delivered-message wrapper that implements [`IncomingMessage`].

// Without the `testing` feature a `JetStream` delivery has one variant; see `broker.rs`.
#![cfg_attr(
    not(feature = "testing"),
    allow(clippy::infallible_destructuring_match)
)]

use std::fmt::{Debug, Formatter};
use std::sync::OnceLock;
use std::time::Duration;

use async_nats::jetstream::AckKind;
use ruststream::{AckError, HeaderMap, IncomingMessage, Partitioned, Str};

use crate::convert::headers_from_nats;
#[cfg(feature = "testing")]
use crate::in_process::{JetStreamDelivery, Release};

/// A NATS delivery. Two flavours: core NATS (no ack) and `JetStream` (real ack/nack/redelivery).
///
/// Both variants are boxed to keep the enum compact; the wrapped `async_nats` messages are large.
pub enum NatsMessage {
    /// A core NATS subject delivery. Acknowledgement is not supported.
    Core(Box<CoreMessage>),
    /// A `JetStream` pull-consumer delivery with full ack support.
    JetStream(Box<JetStreamMessage>),
}

impl Debug for NatsMessage {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Core(_) => f.debug_struct("NatsMessage::Core").finish_non_exhaustive(),
            Self::JetStream(_) => f
                .debug_struct("NatsMessage::JetStream")
                .finish_non_exhaustive(),
        }
    }
}

/// Wrapper around an `async_nats::Message` from a core (non-JetStream) subscription.
pub struct CoreMessage {
    inner: async_nats::Message,
    /// `None` where the delivery carried nothing to put in a map, so a subscription of headerless
    /// deliveries builds none at all; [`headers`](IncomingMessage::headers) answers the crate's
    /// one empty map instead.
    headers: Option<HeaderMap>,
    /// Keeps a delivery of the in-process transport counted with the test harness until it
    /// drops. The field is there only with the `testing` feature.
    #[cfg(feature = "testing")]
    _release: Release,
}

impl Debug for CoreMessage {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CoreMessage")
            .field("subject", &self.inner.subject.as_str())
            .field("payload_len", &self.inner.payload.len())
            .finish_non_exhaustive()
    }
}

impl CoreMessage {
    pub(crate) fn new(inner: async_nats::Message) -> Self {
        if inner.headers.is_none() && inner.reply.is_none() {
            return Self {
                inner,
                headers: None,
                #[cfg(feature = "testing")]
                _release: Release::uncounted(),
            };
        }
        let mut headers = headers_from_nats(inner.headers.as_ref());
        // NATS carries the request inbox in the wire-level `reply` field, not in a header.
        // Surface it as the well-known `reply-to` header so framework handlers can respond
        // (the in-memory testing broker already exposes it this way). The wire field is
        // authoritative: it overrides a literal `reply-to` header if both are present.
        // JetStream deliveries are excluded on purpose - there `reply` is the ack inbox.
        if let Some(reply) = inner.reply.as_ref() {
            headers.insert(Str::from_static(REPLY_TO_HEADER), reply.as_str().to_owned());
        }
        Self {
            inner,
            headers: Some(headers),
            #[cfg(feature = "testing")]
            _release: Release::uncounted(),
        }
    }

    /// A delivery of the in-process transport: the client's own message, counted with the test
    /// harness until it drops.
    #[cfg(feature = "testing")]
    pub(crate) fn in_process(inner: async_nats::Message, release: Release) -> Self {
        Self {
            _release: release,
            ..Self::new(inner)
        }
    }
}

/// Wrapper around an `async_nats::jetstream::Message` with ack semantics.
pub struct JetStreamMessage {
    inner: JetStreamInner,
    /// `None` where the delivery carried no headers; see [`CoreMessage::headers`].
    headers: Option<HeaderMap>,
}

/// A consumer's delivery: the client's, or, under the `testing` feature, the in-process
/// transport's, which settles through its own consumer.
// The size difference exists only with the `testing` feature; boxing the production variant to
// even it out would cost every production delivery an allocation.
#[cfg_attr(feature = "testing", allow(clippy::large_enum_variant))]
enum JetStreamInner {
    Nats(async_nats::jetstream::Message),
    #[cfg(feature = "testing")]
    InProcess(Box<JetStreamDelivery>),
}

// A production build holds the client's delivery and nothing else.
#[cfg(not(feature = "testing"))]
const _: () = assert!(size_of::<JetStreamInner>() == size_of::<async_nats::jetstream::Message>());

/// The native `JetStream` metadata of one delivery, as the server states it.
pub(crate) struct DeliveryMetadata<'a> {
    pub(crate) stream: &'a str,
    pub(crate) consumer: &'a str,
    pub(crate) stream_sequence: u64,
    pub(crate) consumer_sequence: u64,
    pub(crate) delivered: i64,
    pub(crate) pending: u64,
}

impl Debug for JetStreamMessage {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("JetStreamMessage")
            .field("subject", &self.message().subject.as_str())
            .field("payload_len", &self.message().payload.len())
            .finish_non_exhaustive()
    }
}

impl JetStreamMessage {
    pub(crate) fn new(inner: async_nats::jetstream::Message) -> Self {
        let headers = inner
            .message
            .headers
            .as_ref()
            .map(|wire| headers_from_nats(Some(wire)));
        Self {
            inner: JetStreamInner::Nats(inner),
            headers,
        }
    }

    /// A delivery of the in-process transport.
    #[cfg(feature = "testing")]
    pub(crate) fn in_process(delivery: JetStreamDelivery) -> Self {
        let headers = delivery
            .message
            .headers
            .as_ref()
            .map(|wire| headers_from_nats(Some(wire)));
        Self {
            inner: JetStreamInner::InProcess(Box::new(delivery)),
            headers,
        }
    }

    /// The message as the client read it.
    fn message(&self) -> &async_nats::Message {
        match &self.inner {
            JetStreamInner::Nats(inner) => &inner.message,
            #[cfg(feature = "testing")]
            JetStreamInner::InProcess(delivery) => &delivery.message,
        }
    }

    /// The native `JetStream` delivery metadata (stream/consumer name and sequences, redelivery
    /// count, pending count), parsed from the `$JS.ACK` reply subject.
    ///
    /// Returns `None` when the reply subject is absent or malformed - i.e. the underlying
    /// `async_nats` parse failed - so a caller building a context can fall back to "no metadata"
    /// rather than surfacing an error on the per-delivery hot path.
    pub(crate) fn metadata(&self) -> Option<DeliveryMetadata<'_>> {
        match &self.inner {
            JetStreamInner::Nats(inner) => inner.info().ok().map(|info| DeliveryMetadata {
                stream: info.stream,
                consumer: info.consumer,
                stream_sequence: info.stream_sequence,
                consumer_sequence: info.consumer_sequence,
                delivered: info.delivered,
                pending: info.pending,
            }),
            #[cfg(feature = "testing")]
            JetStreamInner::InProcess(delivery) => Some(DeliveryMetadata {
                stream: &delivery.info.stream,
                consumer: &delivery.info.consumer,
                stream_sequence: delivery.info.stream_sequence,
                consumer_sequence: delivery.info.consumer_sequence,
                delivered: i64::try_from(delivery.info.delivered).unwrap_or(i64::MAX),
                pending: delivery.info.pending,
            }),
        }
    }
}

/// The header this crate surfaces a request's wire-level reply subject under.
///
/// NATS carries the inbox in a protocol field rather than a header, so a handler that answers a
/// request reads it here. The generated `AsyncAPI` document names the same header as the runtime
/// expression a client reads a reply address from.
///
/// The name is static, so surfacing it on a delivery costs no allocation.
pub(crate) const REPLY_TO_HEADER: &str = "reply-to";

/// The same header, as the runtime expression an `AsyncAPI` document reports a reply address at.
/// The specification spells the pointer out, so the header name is repeated here rather than
/// composed.
#[cfg(feature = "asyncapi")]
pub(crate) const REPLY_ADDRESS_LOCATION: &str = "$message.header#/reply-to";

/// The map a delivery with no headers of its own answers with.
///
/// The signature `headers()` has to satisfy is `-> &HeaderMap`, and building an empty map per
/// delivery to satisfy it is the whole cost this removes.
fn empty_headers() -> &'static HeaderMap {
    static EMPTY: OnceLock<HeaderMap> = OnceLock::new();
    EMPTY.get_or_init(HeaderMap::new)
}

impl IncomingMessage for NatsMessage {
    fn payload(&self) -> &[u8] {
        match self {
            Self::Core(m) => &m.inner.payload,
            Self::JetStream(m) => &m.message().payload,
        }
    }

    fn headers(&self) -> &HeaderMap {
        let carried = match self {
            Self::Core(m) => m.headers.as_ref(),
            Self::JetStream(m) => m.headers.as_ref(),
        };
        carried.unwrap_or_else(|| empty_headers())
    }

    /// The [`PARTITION_KEY_HEADER`] value the sender set, which the runtime's keyed worker lanes
    /// (`workers(n, by_key)`) order by.
    fn partition_key(&self) -> Option<&[u8]> {
        self.headers().get(PARTITION_KEY_HEADER)
    }

    async fn ack(self) -> Result<(), AckError> {
        match self {
            Self::Core(_) => Err(AckError::Unsupported),
            Self::JetStream(m) => match m.inner {
                JetStreamInner::Nats(inner) => inner
                    .ack()
                    .await
                    .map_err(|err| AckError::Broker(format_err(err))),
                #[cfg(feature = "testing")]
                JetStreamInner::InProcess(delivery) => delivery.ack(),
            },
        }
    }

    async fn nack(self, requeue: bool) -> Result<(), AckError> {
        match self {
            Self::Core(_) => Err(AckError::Unsupported),
            Self::JetStream(m) => {
                let inner = match m.inner {
                    JetStreamInner::Nats(inner) => inner,
                    #[cfg(feature = "testing")]
                    JetStreamInner::InProcess(delivery) => return delivery.nack(requeue),
                };
                let kind = if requeue {
                    AckKind::Nak(None)
                } else {
                    AckKind::Term
                };
                inner
                    .ack_with(kind)
                    .await
                    .map_err(|err| AckError::Broker(format_err(err)))
            }
        }
    }

    /// How many times the server has delivered this message, counting this one.
    ///
    /// A `JetStream` consumer keeps that count and reports it on every delivery, so a cap
    /// declared at the mount site counts the server's own redeliveries - an `ack_wait` that ran
    /// out, a negative acknowledgement - and not only the copies this process published. Core
    /// NATS neither stores a message nor redelivers one, so a core delivery has no count and the
    /// framework's retry-count header is the whole tally.
    fn redelivery_count(&self) -> Option<u64> {
        match self {
            Self::Core(_) => None,
            Self::JetStream(m) => m.metadata().map(|info| info.delivered.unsigned_abs()),
        }
    }

    /// Whether this delivery can honor a native delayed redelivery.
    ///
    /// `true` for every `JetStream` delivery: the protocol carries the delay in the negative
    /// acknowledgement itself, so no opt-in infrastructure is needed. Core NATS has no
    /// acknowledgement at all, so a core delivery reports `false` and the runtime applies its
    /// broker-agnostic deferred re-publish instead.
    fn supports_nack_after(&self) -> bool {
        matches!(self, Self::JetStream(_))
    }

    /// Redelivers this message no sooner than `delay`, natively: `JetStream`'s negative
    /// acknowledgement takes the delay as its argument (`-NAK {"delay": ns}`), so the server holds
    /// the message for that long and then redelivers it on this consumer. Nothing is re-published
    /// and no copy is made, so the delivery count, the stream sequence, and the payload all stay
    /// the ones the message was first delivered with.
    ///
    /// # Errors
    ///
    /// Returns [`AckError::Unsupported`] on a core (non-JetStream) delivery, and
    /// [`AckError::Broker`] when the acknowledgement cannot be sent.
    async fn nack_after(self, delay: Duration) -> Result<(), AckError> {
        match self {
            Self::Core(_) => Err(AckError::Unsupported),
            Self::JetStream(m) => match m.inner {
                JetStreamInner::Nats(inner) => inner
                    .ack_with(AckKind::Nak(Some(delay)))
                    .await
                    .map_err(|err| AckError::Broker(format_err(err))),
                #[cfg(feature = "testing")]
                JetStreamInner::InProcess(delivery) => delivery.nack_after(delay),
            },
        }
    }
}

/// The well-known header key for per-message routing / partitioning.
///
/// Set this header on outgoing messages to control key-based fan-out when the runtime is
/// configured with `workers(N, by_key)`. The runtime hashes the value to assign a dispatch lane.
///
/// The value is text, because a NATS header is: a key that is not UTF-8, or that carries a
/// newline, fails the publish rather than arriving without the header that decides its lane.
pub const PARTITION_KEY_HEADER: &str = "nats-partition-key";

/// `Partitioned` lets the `workers(N, by_key)` runtime feature assign a dispatch lane based on
/// a well-known message header. NATS has no native partition concept, so the key travels as the
/// [`PARTITION_KEY_HEADER`] header value and the sender is responsible for setting it.
impl Partitioned for NatsMessage {
    fn partition_key(&self) -> Option<&[u8]> {
        IncomingMessage::partition_key(self)
    }
}

fn format_err<E>(err: E) -> Box<dyn std::error::Error + Send + Sync>
where
    E: std::fmt::Display + Send + Sync + 'static,
{
    let msg = err.to_string();
    Box::<dyn std::error::Error + Send + Sync>::from(msg)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn core_message(reply: Option<&str>) -> NatsMessage {
        NatsMessage::Core(Box::new(CoreMessage::new(async_nats::Message {
            subject: "subj".into(),
            reply: reply.map(Into::into),
            payload: bytes::Bytes::from_static(b"x"),
            headers: None,
            status: None,
            description: None,
            length: 1,
        })))
    }

    /// A delivery that carries nothing to put in a map answers with the one empty map the crate
    /// keeps, so a subscription of headerless deliveries builds none at all. Two deliveries
    /// answering the same address is what says so.
    #[test]
    fn a_headerless_delivery_answers_with_the_shared_empty_map() {
        let first = core_message(None);
        let second = core_message(None);

        assert!(
            std::ptr::eq(first.headers(), second.headers()),
            "a delivery with neither wire headers nor a reply inbox has nothing to build a map \
             from, so both answer the crate's own empty one",
        );
    }

    #[test]
    fn core_reply_inbox_surfaces_as_reply_to_header() {
        let msg = core_message(Some("_INBOX.42"));
        assert_eq!(msg.headers().reply_to(), Some("_INBOX.42"));
    }

    #[test]
    fn core_message_without_reply_has_no_reply_to() {
        assert_eq!(core_message(None).headers().reply_to(), None);
    }

    // Core NATS has no acknowledgement, so it must not claim the native delay: the runtime reads
    // this to decide between the native `-NAK {"delay"}` and its own deferred re-publish. The
    // JetStream arm answers `true`, which `tests/in_process_nats.rs` and the live suite exercise.
    #[test]
    fn core_delivery_does_not_claim_native_delayed_redelivery() {
        assert!(!core_message(None).supports_nack_after());
    }

    #[tokio::test]
    async fn core_delivery_reports_nack_after_unsupported() {
        let err = core_message(None)
            .nack_after(Duration::from_secs(1))
            .await
            .expect_err("core NATS cannot honor a delayed redelivery");
        assert!(matches!(err, AckError::Unsupported));
    }
}