ruststream_nats/

context.rs

//! Native `JetStream` metadata of a delivery, read by a handler through `Ctx<K>`.
//!
//! A handler binds one field as a parameter with the core's `Ctx<K>` extractor and a key from
//! [`keys`], or types its context parameter as `Context<'_, JetStreamContext>`.
//!
//! - [JetStreamContext][context]: [`JetStreamContext`], the stream, the consumer, the stream and
//!   consumer sequences, the delivery count and the pending count of a delivery, under the keys
//!   [`keys::Stream`], [`keys::Consumer`], [`keys::StreamSequence`], [`keys::ConsumerSequence`],
//!   [`keys::Delivered`] and [`keys::Pending`]. Every key reads `None` on a Core NATS delivery.
//! - [Expectations][expectations]: [`keys::StreamSequence`] is the position a write to the
//!   subject a handler read can expect.
//!
//! [context]: https://powersemmi.github.io/ruststream-nats/latest/components/context/
//! [expectations]: https://powersemmi.github.io/ruststream-nats/latest/components/publishing/jetstream-publish/expectations/

use ruststream::BuildContext;

use crate::message::NatsMessage;

/// Native `JetStream` per-delivery metadata, built once per delivery from the broker message.
///
/// Read its fields by the compile-time [`keys`] (for example
/// [`keys::STREAM_SEQUENCE`]). On a core (non-JetStream) delivery there is no such metadata, so the
/// context is empty and every key reads `None`; the same handler therefore works on both kinds of
/// subscription.
///
/// The context owns its fields (it copies them out of the delivery), so it does not borrow the
/// message: numbers are a stack copy and the stream/consumer names are cloned once per delivery.
///
/// # Examples
///
/// ```
/// # mod demo {
/// use ruststream::schemars::JsonSchema;
/// use ruststream_nats::context::{JetStreamContext, keys};
/// use ruststream_nats::prelude::*;
/// use serde::Deserialize;
///
/// #[derive(Deserialize, JsonSchema)]
/// # #[schemars(crate = "ruststream::schemars")]
/// struct Order {
///     id: u64,
/// }
///
/// struct Ledger;
///
/// impl Handle<Order, (), (), JetStreamContext> for Ledger {
///     async fn handle(
///         &self,
///         order: &Order,
///         _outs: &(),
///         ctx: &mut Context<'_, JetStreamContext>,
///     ) -> Result<(), HandlerOutcome> {
///         if let Some(sequence) = ctx.context(keys::STREAM_SEQUENCE) {
///             tracing::info!(order.id, sequence, "entered in the ledger");
///         }
///         Ok(())
///     }
/// }
///
/// #[ruststream::app]
/// fn app() -> impl App {
///     RustStream::new(AppInfo::new("ledger", "0.1.0")).with_broker(
///         NatsBroker::new("nats://localhost:4222"),
///         |b| {
///             b.include(subscriber(
///                 JetStreamSubject::new("orders.*", "ORDERS").durable("ledger"),
///                 Ledger,
///             )
///             .build());
///         },
///     )
/// }
/// # }
/// # fn main() {}
/// ```
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct JetStreamContext {
    /// `None` for a core delivery, or when the `JetStream` reply subject could not be parsed.
    info: Option<JetStreamInfo>,
}

/// The owned snapshot of one `JetStream` delivery's native metadata.
#[derive(Debug, Clone, PartialEq, Eq)]
struct JetStreamInfo {
    stream: String,
    consumer: String,
    stream_sequence: u64,
    consumer_sequence: u64,
    delivered: i64,
    pending: u64,
}

impl BuildContext<NatsMessage> for JetStreamContext {
    fn build(msg: &NatsMessage) -> Self {
        let info = match msg {
            NatsMessage::JetStream(m) => m.metadata().map(|i| JetStreamInfo {
                stream: i.stream.to_owned(),
                consumer: i.consumer.to_owned(),
                stream_sequence: i.stream_sequence,
                consumer_sequence: i.consumer_sequence,
                delivered: i.delivered,
                pending: i.pending,
            }),
            NatsMessage::Core(_) => None,
        };
        Self { info }
    }
}

// Compile-time guarantee that the context is buildable from the subscriber's actual delivery type,
// so a handler declaring `Context<'_, JetStreamContext>` satisfies the runtime's
// `Cx: BuildContext<S::Message>` mount bound.
const _: fn() = || {
    fn assert_build_context<C: BuildContext<M>, M>() {}
    assert_build_context::<
        JetStreamContext,
        <crate::NatsSubscriber as ruststream::Subscriber>::Message,
    >();
};

/// Compile-time [`Field`](ruststream::Field) keys, one per native `JetStream` metadatum.
///
/// Each key is a zero-sized selector exported both as a type (for naming in bounds) and as a
/// `const` value (for use at the call site, `ctx.context(keys::STREAM_SEQUENCE)`). Every key reads
/// `None` on a core delivery or when the `JetStream` reply subject could not be parsed.
///
/// The fields come from the `JetStream` acknowledgement subject, not from the payload or the
/// headers, so these keys are the one way to reach them.
///
/// # Examples
///
/// ```
/// # mod demo {
/// use ruststream_nats::context::keys::{Delivered, StreamSequence};
/// use ruststream_nats::prelude::*;
/// use serde::Deserialize;
///
/// #[derive(Deserialize)]
/// struct Order {
///     id: u64,
/// }
///
/// #[subscriber(JetStreamSubject::new("orders.*", "ORDERS").durable("worker"))]
/// async fn handle(
///     order: &Order,
///     Ctx(sequence): Ctx<StreamSequence>,
///     Ctx(delivered): Ctx<Delivered>,
/// ) -> HandlerOutcome {
///     // `None` on a core delivery; the stream sequence on a JetStream one.
///     if let Some(sequence) = sequence {
///         tracing::info!(order.id, sequence, "stored");
///     }
///     // The server-side delivery count tells a redelivery from a first delivery.
///     if delivered.is_some_and(|n| n > 1) {
///         tracing::warn!(order.id, "redelivered");
///     }
///     HandlerOutcome::ack()
/// }
///
/// #[ruststream::app]
/// fn app() -> impl App {
///     RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(
///         NatsBroker::new("nats://localhost:4222"),
///         |b| {
///             b.include(handle);
///         },
///     )
/// }
/// # }
/// # fn main() {}
/// ```
pub mod keys {
    use ruststream::{ContextField, Field};

    use super::JetStreamContext;

    /// Selector for the stream sequence number; see [`STREAM_SEQUENCE`].
    #[derive(Debug, Clone, Copy, Default)]
    pub struct StreamSequence;

    /// The monotonically increasing sequence the message holds within its stream.
    pub const STREAM_SEQUENCE: StreamSequence = StreamSequence;

    impl Field<JetStreamContext> for StreamSequence {
        type Value<'a> = Option<u64>;
        fn get(self, cx: &JetStreamContext) -> Option<u64> {
            cx.info.as_ref().map(|i| i.stream_sequence)
        }
    }

    impl ContextField for StreamSequence {
        type Context = JetStreamContext;
        type Value = Option<u64>;
        fn read(self, cx: &JetStreamContext) -> Option<u64> {
            cx.info.as_ref().map(|i| i.stream_sequence)
        }
    }

    /// Selector for the consumer sequence number; see [`CONSUMER_SEQUENCE`].
    #[derive(Debug, Clone, Copy, Default)]
    pub struct ConsumerSequence;

    /// The sequence the message holds within this consumer's delivery stream.
    pub const CONSUMER_SEQUENCE: ConsumerSequence = ConsumerSequence;

    impl Field<JetStreamContext> for ConsumerSequence {
        type Value<'a> = Option<u64>;
        fn get(self, cx: &JetStreamContext) -> Option<u64> {
            cx.info.as_ref().map(|i| i.consumer_sequence)
        }
    }

    impl ContextField for ConsumerSequence {
        type Context = JetStreamContext;
        type Value = Option<u64>;
        fn read(self, cx: &JetStreamContext) -> Option<u64> {
            cx.info.as_ref().map(|i| i.consumer_sequence)
        }
    }

    /// Selector for the server-side delivery count; see [`DELIVERED`].
    #[derive(Debug, Clone, Copy, Default)]
    pub struct Delivered;

    /// The number of times the server has delivered this message.
    ///
    /// `1` on first delivery, higher on a redelivery (after a `nack` or an `ack_wait` timeout).
    /// This is the native redelivery count, not reconstructable from the payload or headers.
    pub const DELIVERED: Delivered = Delivered;

    impl Field<JetStreamContext> for Delivered {
        type Value<'a> = Option<i64>;
        fn get(self, cx: &JetStreamContext) -> Option<i64> {
            cx.info.as_ref().map(|i| i.delivered)
        }
    }

    impl ContextField for Delivered {
        type Context = JetStreamContext;
        type Value = Option<i64>;
        fn read(self, cx: &JetStreamContext) -> Option<i64> {
            cx.info.as_ref().map(|i| i.delivered)
        }
    }

    /// Selector for the pending count; see [`PENDING`].
    #[derive(Debug, Clone, Copy, Default)]
    pub struct Pending;

    /// The number of messages the server still has pending for this consumer behind this one.
    pub const PENDING: Pending = Pending;

    impl Field<JetStreamContext> for Pending {
        type Value<'a> = Option<u64>;
        fn get(self, cx: &JetStreamContext) -> Option<u64> {
            cx.info.as_ref().map(|i| i.pending)
        }
    }

    impl ContextField for Pending {
        type Context = JetStreamContext;
        type Value = Option<u64>;
        fn read(self, cx: &JetStreamContext) -> Option<u64> {
            cx.info.as_ref().map(|i| i.pending)
        }
    }

    /// Selector for the stream name; see [`STREAM`].
    #[derive(Debug, Clone, Copy, Default)]
    pub struct Stream;

    /// The name of the stream this message was delivered from.
    pub const STREAM: Stream = Stream;

    impl Field<JetStreamContext> for Stream {
        type Value<'a> = Option<&'a str>;
        fn get(self, cx: &JetStreamContext) -> Option<&str> {
            cx.info.as_ref().map(|i| i.stream.as_str())
        }
    }

    impl ContextField for Stream {
        type Context = JetStreamContext;
        type Value = Option<String>;
        fn read(self, cx: &JetStreamContext) -> Option<String> {
            cx.info.as_ref().map(|i| i.stream.clone())
        }
    }

    /// Selector for the consumer name; see [`CONSUMER`].
    #[derive(Debug, Clone, Copy, Default)]
    pub struct Consumer;

    /// The name of the durable (or ephemeral) consumer this message was delivered through.
    pub const CONSUMER: Consumer = Consumer;

    impl Field<JetStreamContext> for Consumer {
        type Value<'a> = Option<&'a str>;
        fn get(self, cx: &JetStreamContext) -> Option<&str> {
            cx.info.as_ref().map(|i| i.consumer.as_str())
        }
    }

    impl ContextField for Consumer {
        type Context = JetStreamContext;
        type Value = Option<String>;
        fn read(self, cx: &JetStreamContext) -> Option<String> {
            cx.info.as_ref().map(|i| i.consumer.clone())
        }
    }
}

#[cfg(test)]
mod tests {
    use ruststream::{BuildContext, Field};

    use super::keys::{CONSUMER, CONSUMER_SEQUENCE, DELIVERED, PENDING, STREAM, STREAM_SEQUENCE};
    use super::{JetStreamContext, JetStreamInfo};
    use crate::message::{CoreMessage, NatsMessage};

    fn populated() -> JetStreamContext {
        JetStreamContext {
            info: Some(JetStreamInfo {
                stream: "ORDERS".to_owned(),
                consumer: "orders-worker".to_owned(),
                stream_sequence: 42,
                consumer_sequence: 7,
                delivered: 3,
                pending: 5,
            }),
        }
    }

    #[test]
    fn keys_read_populated_jetstream_fields() {
        let cx = populated();
        assert_eq!(STREAM_SEQUENCE.get(&cx), Some(42));
        assert_eq!(CONSUMER_SEQUENCE.get(&cx), Some(7));
        assert_eq!(DELIVERED.get(&cx), Some(3));
        assert_eq!(PENDING.get(&cx), Some(5));
        assert_eq!(STREAM.get(&cx), Some("ORDERS"));
        assert_eq!(CONSUMER.get(&cx), Some("orders-worker"));
    }

    #[test]
    fn context_field_keys_yield_owned_values() {
        use ruststream::ContextField;

        use super::keys::{Consumer, Delivered, Stream, StreamSequence};

        let cx = populated();
        assert_eq!(StreamSequence.read(&cx), Some(42));
        assert_eq!(Delivered.read(&cx), Some(3));
        assert_eq!(Stream.read(&cx), Some("ORDERS".to_owned()));
        assert_eq!(Consumer.read(&cx), Some("orders-worker".to_owned()));
    }

    #[test]
    fn keys_read_none_without_jetstream_info() {
        let cx = JetStreamContext::default();
        assert_eq!(STREAM_SEQUENCE.get(&cx), None);
        assert_eq!(CONSUMER_SEQUENCE.get(&cx), None);
        assert_eq!(DELIVERED.get(&cx), None);
        assert_eq!(PENDING.get(&cx), None);
        assert_eq!(STREAM.get(&cx), None);
        assert_eq!(CONSUMER.get(&cx), None);
    }

    #[test]
    fn build_over_core_message_has_no_info() {
        // A core delivery is constructible in-process; a JetStream `async_nats::jetstream::Message`
        // is not, so the JetStream build path is compile-validated and exercised against a real
        // server only. Core must yield an empty context (no native JetStream metadata).
        let msg = NatsMessage::Core(Box::new(CoreMessage::new(async_nats::Message {
            subject: "orders.created".into(),
            reply: None,
            payload: bytes::Bytes::from_static(b"{}"),
            headers: None,
            status: None,
            description: None,
            length: 2,
        })));
        let cx = JetStreamContext::build(&msg);
        assert_eq!(cx, JetStreamContext::default());
        assert_eq!(STREAM_SEQUENCE.get(&cx), None);
    }
}