ruststream_nats/
subscriber.rs

1//! Unified NATS subscriber wrapping either a Core or a `JetStream` pull consumer.
2
3use async_nats::jetstream::consumer::{PullConsumer, pull::Stream as PullStream};
4use futures::stream::{poll_fn, unfold};
5use futures::{Stream, StreamExt, future::Either};
6use ruststream::{BatchSubscriber, Subscriber};
7use std::fmt::{Debug, Formatter};
8use std::{pin::Pin, task::Poll, time::Duration};
9use tracing::warn;
10
11use crate::{
12    error::NatsError,
13    message::{CoreMessage, JetStreamMessage, NatsMessage},
14};
15
16/// Cap on a Core NATS batch: [`BatchSubscriber::batches`] drains only what the client has
17/// already buffered locally, and this bounds one drain.
18const CORE_BATCH_LIMIT: usize = 256;
19
20enum SubscriberKind {
21    Core { inner: async_nats::Subscriber },
22    // Box the JetStream variant: PullConsumer is large (~1400 bytes) and the enum would otherwise
23    // penalise the Core path with the same footprint.
24    JetStream(Box<JetStreamKind>),
25}
26
27struct JetStreamKind {
28    inner: Pin<Box<PullStream>>,
29    consumer: PullConsumer,
30    stream_name: String,
31    pull_batch: usize,
32    pull_expires: Duration,
33}
34
35/// A NATS subscription.
36///
37/// Backed transparently by either a Core subscription (no ack) or a `JetStream` pull consumer
38/// (full ack/nack/term). Construct via [`ConnectedNatsBroker::subscribe_with`] with
39/// [`SubscribeOptions`], or let the runtime resolve a [`SubscribeOptions`] source at startup.
40///
41/// [`ConnectedNatsBroker::subscribe_with`]: crate::ConnectedNatsBroker::subscribe_with
42/// [`SubscribeOptions`]: crate::SubscribeOptions
43pub struct NatsSubscriber {
44    subject: String,
45    kind: SubscriberKind,
46}
47
48impl Debug for NatsSubscriber {
49    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
50        let mut s = f.debug_struct("NatsSubscriber");
51        s.field("subject", &self.subject);
52        match &self.kind {
53            SubscriberKind::Core { .. } => {
54                s.field("kind", &"core");
55            }
56            SubscriberKind::JetStream(js) => {
57                s.field("kind", &"jetstream")
58                    .field("stream", &js.stream_name);
59            }
60        }
61        s.finish_non_exhaustive()
62    }
63}
64
65impl NatsSubscriber {
66    pub(crate) const fn from_core(subject: String, inner: async_nats::Subscriber) -> Self {
67        Self {
68            subject,
69            kind: SubscriberKind::Core { inner },
70        }
71    }
72
73    pub(crate) fn from_jetstream(
74        subject: String,
75        stream_name: String,
76        inner: PullStream,
77        consumer: PullConsumer,
78        pull_batch: usize,
79        pull_expires: Duration,
80    ) -> Self {
81        Self {
82            subject,
83            kind: SubscriberKind::JetStream(Box::new(JetStreamKind {
84                inner: Box::pin(inner),
85                consumer,
86                stream_name,
87                pull_batch,
88                pull_expires,
89            })),
90        }
91    }
92}
93
94fn core_message(msg: async_nats::Message) -> NatsMessage {
95    NatsMessage::Core(Box::new(CoreMessage::new(msg)))
96}
97
98fn jetstream_message(msg: async_nats::jetstream::Message) -> NatsMessage {
99    NatsMessage::JetStream(Box::new(JetStreamMessage::new(msg)))
100}
101
102impl Subscriber for NatsSubscriber {
103    type Message = NatsMessage;
104    type Error = NatsError;
105
106    fn stream(&mut self) -> impl Stream<Item = Result<Self::Message, Self::Error>> + Send + '_ {
107        // Poll the inner subscription in place rather than moving it into the returned stream,
108        // so `stream` can be called again after the returned stream is dropped (the runtime and
109        // the conformance helpers re-enter it per call).
110        match &mut self.kind {
111            SubscriberKind::Core { inner } => Either::Left(
112                poll_fn(move |cx| Pin::new(&mut *inner).poll_next(cx))
113                    .map(|msg| Ok(core_message(msg))),
114            ),
115            SubscriberKind::JetStream(js) => Either::Right(
116                poll_fn(move |cx| js.inner.as_mut().poll_next(cx)).map(|item| match item {
117                    Ok(msg) => Ok(jetstream_message(msg)),
118                    Err(err) => {
119                        warn!(target: "ruststream::nats", error = %err, "jetstream fetch error");
120                        Err(NatsError::JetStream(Box::new(err)))
121                    }
122                }),
123            ),
124        }
125    }
126}
127
128impl BatchSubscriber for NatsSubscriber {
129    type Batch = Vec<NatsMessage>;
130
131    /// Returns a stream of message batches.
132    ///
133    /// `JetStream` batches natively: one stream item is one `fetch` of up to
134    /// [`pull_batch`](crate::SubscribeOptions::pull_batch) messages, waiting at most
135    /// [`pull_expires`](crate::SubscribeOptions::pull_expires) before delivering a partial batch
136    /// (an empty fetch is retried, so the stream never yields empty batches). Core NATS has no
137    /// wire-level batching; there a batch is whatever the client has already buffered locally
138    /// (at least one message, at most 256), with no added latency.
139    ///
140    /// Drive a subscriber through either [`Subscriber::stream`] or `batches`, not both at once:
141    /// on `JetStream` each issues its own pull requests, so deliveries would be split between
142    /// them.
143    ///
144    /// # Cancel safety
145    ///
146    /// Dropping the returned stream between items is allowed. On `JetStream`, dropping it
147    /// mid-fetch can leave already-fetched, undelivered messages to be redelivered after the
148    /// consumer's `ack_wait`.
149    fn batches(&mut self) -> impl Stream<Item = Result<Self::Batch, Self::Error>> + Send + '_ {
150        match &mut self.kind {
151            SubscriberKind::Core { inner } => Either::Left(poll_fn(move |cx| {
152                let first = match Pin::new(&mut *inner).poll_next(cx) {
153                    Poll::Pending => return Poll::Pending,
154                    Poll::Ready(None) => return Poll::Ready(None),
155                    Poll::Ready(Some(msg)) => msg,
156                };
157                let mut batch = vec![core_message(first)];
158                while batch.len() < CORE_BATCH_LIMIT {
159                    match Pin::new(&mut *inner).poll_next(cx) {
160                        Poll::Ready(Some(msg)) => batch.push(core_message(msg)),
161                        Poll::Ready(None) | Poll::Pending => break,
162                    }
163                }
164                Poll::Ready(Some(Ok(batch)))
165            })),
166            SubscriberKind::JetStream(js) => {
167                let max = js.pull_batch;
168                let expires = js.pull_expires;
169                Either::Right(unfold(&mut js.consumer, move |consumer| async move {
170                    loop {
171                        let fetch = consumer
172                            .fetch()
173                            .max_messages(max)
174                            .expires(expires)
175                            .messages()
176                            .await;
177                        let mut messages = match fetch {
178                            Ok(messages) => messages,
179                            // BatchError is a concrete sized type; wrap it.
180                            Err(err) => {
181                                return Some((Err(NatsError::JetStream(Box::new(err))), consumer));
182                            }
183                        };
184                        let mut batch = Vec::new();
185                        while let Some(item) = messages.next().await {
186                            match item {
187                                Ok(msg) => batch.push(jetstream_message(msg)),
188                                // crate::Error is already Box<dyn StdError + ...>; use directly.
189                                Err(err) => {
190                                    if batch.is_empty() {
191                                        return Some((Err(NatsError::JetStream(err)), consumer));
192                                    }
193                                    warn!(
194                                        target: "ruststream::nats",
195                                        error = %err,
196                                        "jetstream fetch error mid-batch; delivering the partial batch",
197                                    );
198                                    break;
199                                }
200                            }
201                        }
202                        if !batch.is_empty() {
203                            return Some((Ok(batch), consumer));
204                        }
205                        // An empty fetch only means `expires` elapsed with nothing pending.
206                    }
207                }))
208            }
209        }
210    }
211}
ruststream_nats/subscriber.rs

ruststream_nats/
subscriber.rs
