ruststream_rdkafka/

context.rs

//! Per-delivery and per-batch context fields exposed to handlers.
//!
//! A handler types its context parameter as `Context<'_, KafkaContext>`, or binds one field as a
//! parameter with the core's `Ctx<K>` extractor and a key from [`keys`].
//!
//! - [Kafka context][context]: [`KafkaContext`], the record's topic, partition, offset, timestamp
//!   and key, under the keys [`keys::Topic`], [`keys::Partition`], [`keys::Offset`],
//!   [`keys::TimestampMillis`] and [`keys::Key`].
//! - [KafkaBatchContext][batch]: [`KafkaBatchContext`], a batch handler's context, which carries
//!   only the seek handle; each element of a `&[Message<H, T>]` batch carries its own
//!   coordinates.
//! - [Seek][seek]: [`keys::SeekHandle`] and [`keys::Position`] move a running subscription.
//! - [EosPipeline][eos]: [`keys::Source`] gives a hand-written exactly-once publish its
//!   coordinates.
//!
//! [context]: https://powersemmi.github.io/ruststream-rdkafka/latest/components/context/
//! [batch]: https://powersemmi.github.io/ruststream-rdkafka/latest/components/context/batch-context/
//! [seek]: https://powersemmi.github.io/ruststream-rdkafka/latest/components/positions/seek/
//! [eos]: https://powersemmi.github.io/ruststream-rdkafka/latest/components/transactions/exactly-once/eos-pipeline/

use std::fmt;
use std::sync::Arc;

use bytes::Bytes;
use ruststream::{BuildBatchContext, BuildContext, Field, Str};

use crate::message::KafkaMessage;
use crate::record::Shared;
use crate::seek::{KafkaPosition, KafkaSeeker};

/// Native Kafka delivery metadata plus this subscription's reposition handle, built once per
/// delivery.
#[derive(Clone)]
pub struct KafkaContext {
    /// What the subscription shares, the seeker and the topic name among it, behind one
    /// reference count.
    shared: Arc<Shared>,
    /// The record's own topic, when the subscription reads several.
    topic: Option<Str>,
    partition: i32,
    offset: i64,
    timestamp_millis: Option<i64>,
    key: Option<Bytes>,
}

impl KafkaContext {
    /// The topic the record was consumed from.
    #[must_use]
    pub fn topic(&self) -> &str {
        self.topic_name()
    }

    /// The topic as the shared string: the record's own, or the subscription's.
    fn topic_name(&self) -> &Str {
        self.shared.topic(self.topic.as_ref())
    }

    /// The partition the record was consumed from.
    #[must_use]
    pub fn partition(&self) -> i32 {
        self.partition
    }

    /// The record's offset within its partition.
    #[must_use]
    pub fn offset(&self) -> i64 {
        self.offset
    }

    /// The record's timestamp in milliseconds since the epoch, when the broker provided one.
    #[must_use]
    pub fn timestamp_millis(&self) -> Option<i64> {
        self.timestamp_millis
    }

    /// The record key, or `None` for keyless records.
    #[must_use]
    pub fn key(&self) -> Option<&[u8]> {
        self.key.as_deref()
    }

    /// The topic name in the form the subscription minted it: a field key reading it hands it
    /// over for a reference count instead of copying it.
    pub(crate) fn shared_topic(&self) -> Str {
        self.topic_name().clone()
    }

    /// The record key in the form the delivery carries it, handed over the same way.
    pub(crate) fn shared_key(&self) -> Option<Bytes> {
        self.key.clone()
    }

    /// This delivery's own coordinates: seeking to them redelivers exactly this record, and the
    /// ordered suffix behind it on the partition.
    #[must_use]
    pub fn position(&self) -> KafkaPosition {
        KafkaPosition::topic_offset(self.topic(), self.partition, self.offset)
    }

    /// The handle repositioning the subscription this delivery came from.
    #[must_use]
    pub fn seeker(&self) -> &KafkaSeeker {
        &self.shared.seeker
    }
}

impl fmt::Debug for KafkaContext {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("KafkaContext")
            .field("topic", self.topic_name())
            .field("partition", &self.partition)
            .field("offset", &self.offset)
            .field("timestamp_millis", &self.timestamp_millis)
            .field("key", &self.key)
            .field("seeker", &self.shared.seeker)
            .finish()
    }
}

impl BuildContext<KafkaMessage> for KafkaContext {
    fn build(msg: &KafkaMessage) -> Self {
        Self {
            shared: msg.shared_handle(),
            topic: msg.own_topic(),
            partition: msg.partition(),
            offset: msg.offset(),
            timestamp_millis: msg.timestamp_millis(),
            key: msg.key().map(Bytes::copy_from_slice),
        }
    }
}

/// The subscription-scoped context of a batch handler: the reposition handle every delivery of
/// the batch shares.
///
/// Built once per dispatched batch, off its first delivery. A batch has no single position - it
/// spans many records - so the coordinates a body reacts to travel with the elements, and
/// keeping this a type of its own is what rejects a batch body naming [`KafkaContext`] at compile
/// time.
///
/// # Examples
///
/// ```
/// # #[cfg(feature = "json")]
/// # mod demo {
/// use ruststream_rdkafka::context::KafkaBatchContext;
/// use ruststream_rdkafka::prelude::*;
/// # #[derive(serde::Deserialize)]
/// # struct Order { id: u64, resume_at: Option<i64>, partition: i32 }
///
/// /// The batch is settled first; the reposition then opens the next batch at the target the
/// /// producer marked on one of the elements.
/// #[subscriber(KafkaTopic::new("orders").group("orders-svc"))]
/// async fn reprocess(orders: &[Order], ctx: &mut Context<'_, KafkaBatchContext>) -> HandlerOutcome {
///     let target = orders
///         .iter()
///         .find_map(|order| order.resume_at.map(|at| (order.partition, at)));
///     if let Some((partition, at)) = target
///         && ctx
///             .context(SeekHandle)
///             .seek(KafkaPosition::offset(partition, at))
///             .await
///             .is_err()
///     {
///         return HandlerOutcome::retry();
///     }
///     HandlerOutcome::ack()
/// }
///
/// fn app() -> RustStream {
///     RustStream::new(AppInfo::new("orders", "0.1.0"))
///         .with_broker(KafkaBroker::new(["localhost:9092"]), |b| {
///             b.include(reprocess.batch(nonzero!(100usize)));
///         })
/// }
/// # }
/// ```
#[derive(Debug, Clone)]
pub struct KafkaBatchContext {
    seeker: Arc<KafkaSeeker>,
}

impl KafkaBatchContext {
    /// The handle repositioning the subscription this batch came from.
    #[must_use]
    pub fn seeker(&self) -> &KafkaSeeker {
        &self.seeker
    }
}

impl BuildBatchContext<KafkaMessage> for KafkaBatchContext {
    fn build(first: &KafkaMessage) -> Self {
        Self {
            seeker: first.seeker_handle(),
        }
    }
}

/// Zero-sized [`Field`] keys reading one [`KafkaContext`] field each.
pub mod keys {
    use bytes::Bytes;
    use ruststream::{ContextField, Str};

    use super::{Field, KafkaBatchContext, KafkaContext, KafkaPosition, KafkaSeeker};

    use crate::eos::SourceOffset;

    /// Reads the source topic name.
    ///
    /// Both forms hand over the name the subscription minted: the borrow is a [`Str`], so owning
    /// it is a reference count rather than a copy.
    #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
    pub struct Topic;

    impl Field<KafkaContext> for Topic {
        type Value<'a> = &'a Str;

        fn get(self, src: &KafkaContext) -> &Str {
            src.topic_name()
        }
    }

    impl ContextField for Topic {
        type Context = KafkaContext;
        type Value = Str;
        fn read(self, src: &KafkaContext) -> Str {
            src.shared_topic()
        }
    }

    /// Reads the source partition.
    #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
    pub struct Partition;

    impl Field<KafkaContext> for Partition {
        type Value<'a> = i32;

        fn get(self, src: &KafkaContext) -> i32 {
            src.partition()
        }
    }

    impl ContextField for Partition {
        type Context = KafkaContext;
        type Value = i32;
        fn read(self, src: &KafkaContext) -> i32 {
            src.partition()
        }
    }

    /// Reads the record offset.
    #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
    pub struct Offset;

    impl Field<KafkaContext> for Offset {
        type Value<'a> = i64;

        fn get(self, src: &KafkaContext) -> i64 {
            src.offset()
        }
    }

    impl ContextField for Offset {
        type Context = KafkaContext;
        type Value = i64;
        fn read(self, src: &KafkaContext) -> i64 {
            src.offset()
        }
    }

    /// Reads the record timestamp in milliseconds since the epoch, when present.
    #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
    pub struct TimestampMillis;

    impl Field<KafkaContext> for TimestampMillis {
        type Value<'a> = Option<i64>;

        fn get(self, src: &KafkaContext) -> Option<i64> {
            src.timestamp_millis()
        }
    }

    impl ContextField for TimestampMillis {
        type Context = KafkaContext;
        type Value = Option<i64>;
        fn read(self, src: &KafkaContext) -> Option<i64> {
            src.timestamp_millis()
        }
    }

    /// Reads the record key, when present.
    ///
    /// The owned form is the delivery's own buffer: the first read promotes it to a shared block,
    /// one allocation per delivery, and every read after that is a reference count.
    #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
    pub struct Key;

    impl Field<KafkaContext> for Key {
        type Value<'a> = Option<&'a [u8]>;

        fn get(self, src: &KafkaContext) -> Option<&[u8]> {
            src.key()
        }
    }

    impl ContextField for Key {
        type Context = KafkaContext;
        type Value = Option<Bytes>;
        fn read(self, src: &KafkaContext) -> Option<Bytes> {
            src.shared_key()
        }
    }

    /// Reads the delivery's source coordinates as one value, the form
    /// [`EosPipeline::publish`](crate::EosPipeline::publish) takes.
    #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
    pub struct Source;

    impl Field<KafkaContext> for Source {
        type Value<'a> = SourceOffset;

        fn get(self, src: &KafkaContext) -> SourceOffset {
            SourceOffset::new(src.shared_topic(), src.partition(), src.offset())
        }
    }

    impl ContextField for Source {
        type Context = KafkaContext;
        type Value = SourceOffset;
        fn read(self, src: &KafkaContext) -> SourceOffset {
            SourceOffset::new(src.shared_topic(), src.partition(), src.offset())
        }
    }

    /// Reads this delivery's own log coordinates, in the form
    /// [`Seeker::seek`](ruststream::Seeker::seek) takes: seeking to them redelivers exactly this
    /// record.
    ///
    /// # Examples
    ///
    /// ```
    /// use ruststream_rdkafka::prelude::*;
    /// # #[derive(serde::Deserialize)]
    /// # struct Order { id: u64 }
    ///
    /// /// Rewinds to the record it is holding when a downstream dependency is unavailable, so
    /// /// the partition resumes here rather than from the group's committed position.
    /// #[subscriber("orders")]
    /// async fn place(
    ///     order: &Order,
    ///     Ctx(here): Ctx<Position>,
    ///     Ctx(seeker): Ctx<SeekHandle>,
    /// ) -> HandlerOutcome {
    ///     if order.id == 0 && seeker.seek(here).await.is_err() {
    ///         return HandlerOutcome::retry();
    ///     }
    ///     HandlerOutcome::ack()
    /// }
    /// ```
    #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
    pub struct Position;

    impl Field<KafkaContext> for Position {
        type Value<'a> = KafkaPosition;

        fn get(self, src: &KafkaContext) -> KafkaPosition {
            src.position()
        }
    }

    impl ContextField for Position {
        type Context = KafkaContext;
        type Value = KafkaPosition;
        fn read(self, src: &KafkaContext) -> KafkaPosition {
            src.position()
        }
    }

    /// Reads the subscription's reposition handle.
    ///
    /// The same key serves the per-delivery [`KafkaContext`] and the batch-scoped
    /// [`KafkaBatchContext`], because a seek is a subscription operation either way.
    ///
    /// # Examples
    ///
    /// ```
    /// use ruststream_rdkafka::prelude::*;
    /// # #[derive(serde::Deserialize)]
    /// # struct Job { id: u64, resume_at: Option<i64> }
    ///
    /// /// Jumps the partition past a run the producer marked unprocessable.
    /// #[subscriber("jobs")]
    /// async fn work(
    ///     job: &Job,
    ///     Ctx(partition): Ctx<Partition>,
    ///     Ctx(seeker): Ctx<SeekHandle>,
    /// ) -> HandlerOutcome {
    ///     if let Some(resume_at) = job.resume_at
    ///         && seeker
    ///             .seek(KafkaPosition::offset(partition, resume_at))
    ///             .await
    ///             .is_err()
    ///     {
    ///         return HandlerOutcome::retry();
    ///     }
    ///     HandlerOutcome::ack()
    /// }
    /// ```
    #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
    pub struct SeekHandle;

    impl Field<KafkaContext> for SeekHandle {
        type Value<'a> = &'a KafkaSeeker;

        fn get(self, src: &KafkaContext) -> &KafkaSeeker {
            src.seeker()
        }
    }

    impl ContextField for SeekHandle {
        type Context = KafkaContext;
        type Value = KafkaSeeker;
        fn read(self, src: &KafkaContext) -> KafkaSeeker {
            src.seeker().clone()
        }
    }

    impl Field<KafkaBatchContext> for SeekHandle {
        type Value<'a> = &'a KafkaSeeker;

        fn get(self, src: &KafkaBatchContext) -> &KafkaSeeker {
            src.seeker()
        }
    }
}