Skip to main content

ruststream_nats/
message.rs

1//! Delivered-message wrapper that implements [`IncomingMessage`].
2
3use std::fmt::{Debug, Formatter};
4use std::sync::OnceLock;
5use std::time::Duration;
6
7use async_nats::jetstream::AckKind;
8use ruststream::{AckError, Headers, IncomingMessage, Partitioned};
9
10use crate::convert::headers_from_nats;
11
12/// A NATS delivery. Two flavours: core NATS (no ack) and `JetStream` (real ack/nack/redelivery).
13///
14/// Both variants are boxed to keep the enum compact; the wrapped `async_nats` messages are large.
15pub enum NatsMessage {
16    /// A core NATS subject delivery. Acknowledgement is not supported.
17    Core(Box<CoreMessage>),
18    /// A `JetStream` pull-consumer delivery with full ack support.
19    JetStream(Box<JetStreamMessage>),
20}
21
22impl Debug for NatsMessage {
23    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
24        match self {
25            Self::Core(_) => f.debug_struct("NatsMessage::Core").finish_non_exhaustive(),
26            Self::JetStream(_) => f
27                .debug_struct("NatsMessage::JetStream")
28                .finish_non_exhaustive(),
29        }
30    }
31}
32
33/// Wrapper around an `async_nats::Message` from a core (non-JetStream) subscription.
34pub struct CoreMessage {
35    inner: async_nats::Message,
36    headers: Headers,
37}
38
39impl Debug for CoreMessage {
40    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
41        f.debug_struct("CoreMessage")
42            .field("subject", &self.inner.subject.as_str())
43            .field("payload_len", &self.inner.payload.len())
44            .finish_non_exhaustive()
45    }
46}
47
48impl CoreMessage {
49    pub(crate) fn new(inner: async_nats::Message) -> Self {
50        let mut headers = headers_from_nats(inner.headers.as_ref());
51        // NATS carries the request inbox in the wire-level `reply` field, not in a header.
52        // Surface it as the well-known `reply-to` header so framework handlers can respond
53        // (the in-memory testing broker already exposes it this way). The wire field is
54        // authoritative: it overrides a literal `reply-to` header if both are present.
55        // JetStream deliveries are excluded on purpose - there `reply` is the ack inbox.
56        if let Some(reply) = inner.reply.as_ref() {
57            headers.insert("reply-to", reply.as_str().to_owned());
58        }
59        Self { inner, headers }
60    }
61}
62
63/// Wrapper around an `async_nats::jetstream::Message` with ack semantics.
64pub struct JetStreamMessage {
65    inner: async_nats::jetstream::Message,
66    headers: Headers,
67}
68
69impl Debug for JetStreamMessage {
70    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
71        f.debug_struct("JetStreamMessage")
72            .field("subject", &self.inner.message.subject.as_str())
73            .field("payload_len", &self.inner.message.payload.len())
74            .finish_non_exhaustive()
75    }
76}
77
78impl JetStreamMessage {
79    pub(crate) fn new(inner: async_nats::jetstream::Message) -> Self {
80        let headers = headers_from_nats(inner.message.headers.as_ref());
81        Self { inner, headers }
82    }
83
84    /// The native `JetStream` delivery metadata (stream/consumer name and sequences, redelivery
85    /// count, pending count), parsed from the `$JS.ACK` reply subject.
86    ///
87    /// Returns `None` when the reply subject is absent or malformed - i.e. the underlying
88    /// `async_nats` parse failed - so a caller building a context can fall back to "no metadata"
89    /// rather than surfacing an error on the per-delivery hot path.
90    pub(crate) fn info(&self) -> Option<async_nats::jetstream::message::Info<'_>> {
91        self.inner.info().ok()
92    }
93}
94
95fn empty_headers() -> &'static Headers {
96    static EMPTY: OnceLock<Headers> = OnceLock::new();
97    EMPTY.get_or_init(Headers::new)
98}
99
100impl IncomingMessage for NatsMessage {
101    fn payload(&self) -> &[u8] {
102        match self {
103            Self::Core(m) => &m.inner.payload,
104            Self::JetStream(m) => &m.inner.message.payload,
105        }
106    }
107
108    fn headers(&self) -> &Headers {
109        match self {
110            Self::Core(m) => &m.headers,
111            Self::JetStream(m) => &m.headers,
112        }
113    }
114
115    async fn ack(self) -> Result<(), AckError> {
116        match self {
117            Self::Core(_) => Err(AckError::Unsupported),
118            Self::JetStream(m) => m
119                .inner
120                .ack()
121                .await
122                .map_err(|err| AckError::Broker(format_err(err))),
123        }
124    }
125
126    async fn nack(self, requeue: bool) -> Result<(), AckError> {
127        match self {
128            Self::Core(_) => Err(AckError::Unsupported),
129            Self::JetStream(m) => {
130                let kind = if requeue {
131                    AckKind::Nak(None)
132                } else {
133                    AckKind::Term
134                };
135                m.inner
136                    .ack_with(kind)
137                    .await
138                    .map_err(|err| AckError::Broker(format_err(err)))
139            }
140        }
141    }
142
143    /// Whether this delivery can honor a native delayed redelivery.
144    ///
145    /// `true` for every `JetStream` delivery: the protocol carries the delay in the negative
146    /// acknowledgement itself, so no opt-in infrastructure is needed. Core NATS has no
147    /// acknowledgement at all, so a core delivery reports `false` and the runtime applies its
148    /// broker-agnostic deferred re-publish instead.
149    fn supports_nack_after(&self) -> bool {
150        matches!(self, Self::JetStream(_))
151    }
152
153    /// Redelivers this message no sooner than `delay`, natively: `JetStream`'s negative
154    /// acknowledgement takes the delay as its argument (`-NAK {"delay": ns}`), so the server holds
155    /// the message for that long and then redelivers it on this consumer. Nothing is re-published
156    /// and no copy is made, so the delivery count, the stream sequence, and the payload all stay
157    /// the ones the message was first delivered with.
158    ///
159    /// # Errors
160    ///
161    /// Returns [`AckError::Unsupported`] on a core (non-JetStream) delivery, and
162    /// [`AckError::Broker`] when the acknowledgement cannot be sent.
163    async fn nack_after(self, delay: Duration) -> Result<(), AckError> {
164        match self {
165            Self::Core(_) => Err(AckError::Unsupported),
166            Self::JetStream(m) => m
167                .inner
168                .ack_with(AckKind::Nak(Some(delay)))
169                .await
170                .map_err(|err| AckError::Broker(format_err(err))),
171        }
172    }
173}
174
175/// The well-known header key for per-message routing / partitioning.
176///
177/// Set this header on outgoing messages to control key-based fan-out when the runtime is
178/// configured with `workers(N, by_key)`. The value is opaque bytes; the runtime hashes it to
179/// assign a dispatch lane.
180pub const PARTITION_KEY_HEADER: &str = "nats-partition-key";
181
182/// `Partitioned` lets the `workers(N, by_key)` runtime feature assign a dispatch lane based on
183/// a well-known message header. NATS has no native partition concept, so the key travels as the
184/// [`PARTITION_KEY_HEADER`] header value and the sender is responsible for setting it.
185impl Partitioned for NatsMessage {
186    fn partition_key(&self) -> Option<&[u8]> {
187        self.headers().get(PARTITION_KEY_HEADER)
188    }
189}
190
191fn format_err<E>(err: E) -> Box<dyn std::error::Error + Send + Sync>
192where
193    E: std::fmt::Display + Send + Sync + 'static,
194{
195    let msg = err.to_string();
196    Box::<dyn std::error::Error + Send + Sync>::from(msg)
197}
198
199#[allow(dead_code)]
200fn _empty_headers_keepalive() -> &'static Headers {
201    empty_headers()
202}
203
204#[cfg(test)]
205mod tests {
206    use super::*;
207
208    fn core_message(reply: Option<&str>) -> NatsMessage {
209        NatsMessage::Core(Box::new(CoreMessage::new(async_nats::Message {
210            subject: "subj".into(),
211            reply: reply.map(Into::into),
212            payload: bytes::Bytes::from_static(b"x"),
213            headers: None,
214            status: None,
215            description: None,
216            length: 1,
217        })))
218    }
219
220    #[test]
221    fn core_reply_inbox_surfaces_as_reply_to_header() {
222        let msg = core_message(Some("_INBOX.42"));
223        assert_eq!(msg.headers().reply_to(), Some("_INBOX.42"));
224    }
225
226    #[test]
227    fn core_message_without_reply_has_no_reply_to() {
228        assert_eq!(core_message(None).headers().reply_to(), None);
229    }
230
231    // Core NATS has no acknowledgement, so it must not claim the native delay: the runtime reads
232    // this to decide between the native `-NAK {"delay"}` and its own deferred re-publish. The
233    // JetStream arm answers `true` and is exercised against a real server (a
234    // `async_nats::jetstream::Message` has no in-process constructor).
235    #[test]
236    fn core_delivery_does_not_claim_native_delayed_redelivery() {
237        assert!(!core_message(None).supports_nack_after());
238    }
239
240    #[tokio::test]
241    async fn core_delivery_reports_nack_after_unsupported() {
242        let err = core_message(None)
243            .nack_after(Duration::from_secs(1))
244            .await
245            .expect_err("core NATS cannot honor a delayed redelivery");
246        assert!(matches!(err, AckError::Unsupported));
247    }
248}