ruststream_rdkafka/

message.rs

//! The delivery type yielded by [`KafkaSubscriber`](crate::KafkaSubscriber).

use std::convert::Infallible;
use std::fmt;
use std::future::{Future, ready};
use std::io::Write as _;
use std::sync::{Arc, OnceLock};

#[cfg(feature = "schema-registry")]
use bytes::Bytes;
use ruststream::{AckError, HeaderMap, IncomingMessage, Partitioned, Positioned, Str};

use crate::convert;
use crate::record::{HeldRecord, Shared};
use crate::seek::{KafkaPosition, KafkaSeeker};
use crate::tracker::PartitionSlot;

/// Header carrying a message's partition key, mapped onto Kafka's native record key.
///
/// On publish, this header becomes the record key (so Kafka itself routes deliveries that share
/// a key to the same partition) and is not duplicated as a wire header. On consume, the header
/// always mirrors the native record key - a same-named wire header from a foreign producer is
/// not preserved, because the record key is Kafka's source of truth for partitioning. Keyed
/// worker lanes (`workers(n, by_key)`) read it through
/// [`IncomingMessage::partition_key`]; [`Partitioned`] mirrors it as the capability surface.
pub const PARTITION_KEY_HEADER: &str = "kafka-partition-key";

/// How this delivery settles when acked.
///
/// The tracked forms carry the partition slot the delivery was pulled in, so a settle that
/// arrives after the subscription was repositioned (a seek, a rebalance) is dropped instead of
/// moving a position that no longer describes what this consumer reads.
pub(crate) enum Settlement {
    /// `Commit::Auto`: librdkafka owns the committed position; `ack`/`nack` are advisory.
    Advisory,
    /// `Commit::Tracked`: an ack advances the shared watermark and stores the new position,
    /// through the tracker and the consumer the delivery's own record keeps alive.
    Tracked { slot: PartitionSlot },
    /// `Commit::Transactional`: an ack advances the shared watermark only - the EOS pipeline
    /// commits positions through the producer transaction, so nothing is stored here.
    Transactional { slot: PartitionSlot },
}

/// How a delivery answers [`IncomingMessage::partition_key`], which most deliveries are never
/// asked.
///
/// Both forms answer out of what the delivery already holds, so a subscription with no keyed
/// lanes - the default - pays nothing for a key nobody reads.
pub(crate) enum Lane {
    /// The record key, which the delivery carries as its [`PARTITION_KEY_HEADER`] header.
    RecordKey,
    /// The source partition, written into the delivery's own bytes on the first ask.
    Partition(OnceLock<PartitionText>),
}

impl Lane {
    /// This delivery's lane key.
    fn of<'a>(&'a self, msg: &'a KafkaMessage) -> Option<&'a [u8]> {
        match self {
            Self::RecordKey => msg.record.key(),
            Self::Partition(text) => Some(
                text.get_or_init(|| PartitionText::of(msg.partition))
                    .as_bytes(),
            ),
        }
    }
}

impl fmt::Debug for Lane {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::RecordKey => f.write_str("RecordKey"),
            Self::Partition(_) => f.write_str("Partition"),
        }
    }
}

/// A partition number as the decimal text a lane key is, held inline: the widest `i32` is eleven
/// bytes, which is shorter than the pointer a heap copy of it would cost. The length is a byte
/// because every delivery carries this cell whether or not anything reads it.
#[derive(Debug)]
pub(crate) struct PartitionText {
    bytes: [u8; 11],
    len: u8,
}

impl PartitionText {
    fn of(partition: i32) -> Self {
        let mut bytes = [0u8; 11];
        let mut cursor = &mut bytes[..];
        write!(cursor, "{partition}").expect("eleven bytes hold the widest i32");
        let len = u8::try_from(11 - cursor.len()).expect("eleven bytes hold the widest i32");
        Self { bytes, len }
    }

    fn as_bytes(&self) -> &[u8] {
        &self.bytes[..self.len as usize]
    }
}

/// One Kafka delivery: an owned snapshot of the record plus its settlement handle.
///
/// Settlement mapping depends on the [`Commit`](crate::Commit) mode of the subscription:
///
/// Under `Commit::Auto` (the default) librdkafka owns the committed position - it is stored
/// the moment a message is handed to the application - so `ack` and `nack(false)` are advisory
/// no-ops, and `nack(true)` reports [`AckError::Unsupported`]: nothing brings the record back.
///
/// Under `Commit::Tracked`:
///
/// - [`ack`](IncomingMessage::ack) settles the offset and advances the stored position across
///   everything settled below it.
/// - [`nack(false)`](IncomingMessage::nack) drops the message: the offset settles so the
///   position can move past it.
/// - [`nack(true)`](IncomingMessage::nack) hands the delivery back to its subscription, which
///   delivers it again ahead of what it fetches next. Its offset stays unsettled meanwhile, so
///   the committed position stays below it: were the process to stop first, Kafka redelivers
///   the record from there to whoever reads the partition next. A rebalance or a seek that
///   moves the partition before the delivery comes back drops it for the same reason. Until it
///   settles, the unsettled offset also blocks the position, keeping every later ack
///   uncommitted.
///
/// Under `Commit::Transactional` the settlements report to the exactly-once pipeline instead,
/// and `nack(true)` keeps the open window from committing: it aborts, and its replay delivers
/// the record again with the rest of the window.
///
/// Wire headers map name for name; a null-valued Kafka header arrives with an empty value
/// (presence preserved).
pub struct KafkaMessage {
    /// The record librdkafka fetched, kept where it put it: the body, the key and the wire
    /// headers are read out of the fetch buffer, and the consumer that owns that buffer stays
    /// open for as long as this delivery does.
    record: HeldRecord,
    /// The payload the registry middleware transcoded, when there is one. Without it the body
    /// is the record's own.
    #[cfg(feature = "schema-registry")]
    transcoded: Option<Bytes>,
    /// The `RustStream` view of the record's headers, built on the first read: a delivery
    /// nothing asks headers of never pays for them.
    headers: OnceLock<HeaderMap>,
    /// The record's topic, when the subscription reads several: a Kafka consumer reads a handful
    /// of topics and delivers millions of records, so the name is minted once per topic and
    /// every delivery of it takes a reference count. A subscription over one topic leaves it
    /// empty, and the delivery reads the name its shared block holds, for no count of its own.
    topic: Option<Str>,
    partition: i32,
    offset: i64,
    settlement: Settlement,
    /// How this delivery answers `partition_key()`.
    lane: Lane,
}

impl fmt::Debug for KafkaMessage {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut out = f.debug_struct("KafkaMessage");
        out.field("record", &self.record);
        #[cfg(feature = "schema-registry")]
        out.field("transcoded", &self.transcoded);
        out.field("headers", &self.headers)
            .field("topic", self.topic_name())
            .field("partition", &self.partition)
            .field("offset", &self.offset)
            .field("settlement", &self.settlement)
            .field("lane", &self.lane)
            .finish()
    }
}

impl fmt::Debug for Settlement {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Advisory => f.write_str("Advisory"),
            Self::Tracked { .. } => f.debug_struct("Tracked").finish_non_exhaustive(),
            Self::Transactional { .. } => f.debug_struct("Transactional").finish_non_exhaustive(),
        }
    }
}

impl KafkaMessage {
    // An internal constructor mirroring the record's natural fields; grouping them into
    // intermediate structs would only add indirection for the one caller.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        record: HeldRecord,
        headers: OnceLock<HeaderMap>,
        topic: Option<Str>,
        partition: i32,
        offset: i64,
        settlement: Settlement,
        lane: Lane,
    ) -> Self {
        Self {
            record,
            #[cfg(feature = "schema-registry")]
            transcoded: None,
            headers,
            topic,
            partition,
            offset,
            settlement,
            lane,
        }
    }

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

    /// The subscription's reposition handle, for the context built off a batch.
    pub(crate) fn seeker_handle(&self) -> Arc<KafkaSeeker> {
        Arc::clone(&self.record.shared().seeker)
    }

    /// What the subscription shares, for the context built off this delivery: one count holds
    /// the seeker and the topic name alike.
    pub(crate) fn shared_handle(&self) -> Arc<Shared> {
        Arc::clone(self.record.shared())
    }

    /// The topic name this delivery carries itself, when its subscription reads several.
    pub(crate) fn own_topic(&self) -> Option<Str> {
        self.topic.clone()
    }

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

    /// The partition this 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.
    ///
    /// Read from the record on the ask rather than carried by the delivery: librdkafka answers
    /// it out of the fetch buffer this delivery holds open, and most deliveries are never asked.
    #[must_use]
    pub fn timestamp_millis(&self) -> Option<i64> {
        self.record.timestamp_millis()
    }

    /// The record key, surfaced from Kafka's native key (see [`PARTITION_KEY_HEADER`]).
    #[must_use]
    pub fn key(&self) -> Option<&[u8]> {
        self.record.key()
    }

    /// Replaces the payload with its registry-transcoded form (the subscriber's async
    /// middleware), before the delivery is handed on.
    #[cfg(feature = "schema-registry")]
    pub(crate) fn replace_payload(&mut self, payload: Bytes) {
        self.transcoded = Some(payload);
    }

    /// Whether this delivery still belongs to its partition's read position. One handed back
    /// with `nack(true)` is delivered again only while it does: after a seek or a revoke the
    /// replay or the partition's next owner delivers its record instead.
    pub(crate) fn is_current(&self) -> bool {
        match &self.settlement {
            Settlement::Advisory => true,
            Settlement::Tracked { slot } | Settlement::Transactional { slot } => {
                self.record.shared().tracker.is_current(*slot)
            }
        }
    }

    fn settle(self) -> Result<(), AckError> {
        match &self.settlement {
            Settlement::Advisory => Ok(()),
            Settlement::Tracked { slot } => self
                .record
                .shared()
                .tracker
                .settle_with(*slot, self.offset, |position, target| {
                    self.record.store(self.offset, position, target)
                })
                .map_err(|err| AckError::Broker(Box::new(err))),
            Settlement::Transactional { slot } => {
                let infallible: Result<(), Infallible> = self.record.shared().tracker.settle_with(
                    *slot,
                    self.offset,
                    |_position, _target| Ok(()),
                );
                infallible.expect("no-op store cannot fail");
                Ok(())
            }
        }
    }
}

impl IncomingMessage for KafkaMessage {
    fn payload(&self) -> &[u8] {
        #[cfg(feature = "schema-registry")]
        if let Some(transcoded) = &self.transcoded {
            return transcoded;
        }
        self.record.payload()
    }

    fn headers(&self) -> &HeaderMap {
        self.headers
            .get_or_init(|| convert::headers_from_message(&self.record))
    }

    /// Marks the offset processed (see the type-level settlement mapping).
    ///
    /// # Errors
    ///
    /// Returns [`AckError::Broker`] when librdkafka refuses the stored position, which is a
    /// client-state failure rather than a network one - the partition is no longer this
    /// consumer's, say. A descriptor passthrough that would take the offset store back is not
    /// one of those cases and never reaches here: librdkafka would accept the store and keep
    /// owning the position, so [`Commit::Tracked`](crate::Commit::Tracked) and
    /// [`Commit::Transactional`](crate::Commit::Transactional) refuse such a subscription at
    /// startup instead.
    ///
    /// # Cancel safety
    ///
    /// Cancel safe: the watermark update is synchronous, so the future either completed or did
    /// nothing.
    fn ack(self) -> impl Future<Output = Result<(), AckError>> {
        ready(self.settle())
    }

    /// Settles negatively. `requeue = false` drops the delivery: the offset settles so the
    /// committed position can move past it. `requeue = true` hands the delivery back to its
    /// subscription, which delivers it again; its offset stays unsettled until then. Under
    /// `Commit::Transactional` it aborts the open window, whose replay delivers it again. Under
    /// `Commit::Auto` `requeue = false` is an advisory no-op and `requeue = true` is unsupported
    /// (see the type-level settlement mapping).
    ///
    /// Kafka counts no deliveries of its own, so a registration that caps its attempts with
    /// `max_attempts(..)` gets the framework's count instead: the runtime republishes the
    /// delivery through the registration's retry publisher so the count travels with it.
    ///
    /// # Errors
    ///
    /// Returns [`AckError::Unsupported`] for `requeue = true` under `Commit::Auto`, where the
    /// position was stored when the record was handed over, and [`AckError::Broker`] under the
    /// same conditions as [`ack`](Self::ack).
    ///
    /// # Cancel safety
    ///
    /// Cancel safe: the watermark update is synchronous, so the future either completed or did
    /// nothing.
    fn nack(self, requeue: bool) -> impl Future<Output = Result<(), AckError>> {
        ready(match (requeue, &self.settlement) {
            // Auto-commit stored the position as the record was handed over, so no requeue is
            // left to perform: a success here would promise a redelivery that never comes.
            (true, Settlement::Advisory) => Err(AckError::Unsupported),
            // The subscription delivers it again. Its offset stays unsettled meanwhile, which is
            // what brings the record back from Kafka should the process stop first.
            (true, Settlement::Tracked { .. }) => {
                let tracker = Arc::clone(&self.record.shared().tracker);
                tracker.attention.requeue(self);
                Ok(())
            }
            // Under an exactly-once pipeline the unsettled offset keeps the open window from
            // committing: it aborts, and its seek back delivers the whole window again, this
            // record included. Delivering it again inside the window would publish its output
            // twice in one transaction.
            (true, Settlement::Transactional { .. }) => Ok(()),
            (false, _) => self.settle(),
        })
    }

    /// The keyed-lane key, so keyed worker lanes see it without a `Partitioned` bound: the
    /// source partition (the default), or the record key under
    /// [`LaneKey::RecordKey`](crate::LaneKey::RecordKey).
    fn partition_key(&self) -> Option<&[u8]> {
        self.lane.of(self)
    }
}

impl Positioned for KafkaMessage {
    type Position = KafkaPosition;

    /// This delivery's own coordinates: seeking to them redelivers exactly this record (and the
    /// ordered suffix behind it on the partition).
    fn position(&self) -> Self::Position {
        KafkaPosition::topic_offset(&**self.topic_name(), self.partition, self.offset)
    }
}

impl Partitioned for KafkaMessage {
    /// The keyed-lane key (see [`IncomingMessage::partition_key`] on this type): the source
    /// partition (the default), or the record key under
    /// [`LaneKey::RecordKey`](crate::LaneKey::RecordKey).
    fn partition_key(&self) -> Option<&[u8]> {
        self.lane.of(self)
    }
}

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

    #[test]
    fn a_partition_lane_key_is_the_partition_in_decimal() {
        // The widest `i32` is eleven bytes, which is what the inline buffer and its length are
        // sized for; a lane key that came out short would quietly merge two partitions' lanes.
        for partition in [0, 7, 42, i32::MAX, i32::MIN] {
            assert_eq!(
                PartitionText::of(partition).as_bytes(),
                partition.to_string().as_bytes(),
                "the lane key of partition {partition}",
            );
        }
    }
}