ruststream_rdkafka/

seek.rs

//! Repositioning a live subscription: the Kafka position vocabulary and the seeker that
//! applies it.
//!
//! Kafka is a replayable log, so a running consumer can be moved: back to reprocess, forward to
//! skip a poison region, or to the first record at a wall-clock time. The reposition applies to
//! the partitions **this consumer instance** currently holds, not to the group - other members
//! keep reading where they were, and nothing is committed on their behalf.

use std::sync::Arc;
use std::time::Duration;

use rdkafka::consumer::{Consumer, StreamConsumer};
use rdkafka::error::RDKafkaErrorCode;
use rdkafka::{Offset, TopicPartitionList};
use ruststream::{Seeker, Str};
use tokio::task;

use crate::broker::ConnState;
use crate::error::KafkaError;
#[cfg(feature = "testing")]
use crate::in_process::Member;
use crate::tracker::{CommitTracker, TrackedConsumer, TrackingContext};

/// How long a reposition waits for librdkafka (the seek itself, and the timestamp lookup).
const SEEK_TIMEOUT: Duration = Duration::from_secs(10);

/// Where a subscription should resume reading.
///
/// The stream-wide variants ([`Earliest`](Self::Earliest), [`Latest`](Self::Latest),
/// [`Timestamp`](Self::Timestamp)) apply to every partition currently assigned to this
/// consumer; [`Offset`](Self::Offset) names one partition. Build them with the constructors
/// ([`earliest`](Self::earliest), [`offset`](Self::offset), ...) - the variants are what the
/// seeker matches on.
///
/// # Examples
///
/// A read model rebuilt from the start of the retained log on every startup, whatever the group
/// committed before:
///
/// ```
/// use ruststream_rdkafka::prelude::*;
/// # #[derive(serde::Deserialize)]
/// # struct Order { id: u64 }
///
/// #[subscriber(
///     KafkaTopic::new("orders").group("orders-projection"),
///     start_at(KafkaPosition::earliest())
/// )]
/// async fn project(order: &Order) -> HandlerOutcome {
///     tracing::info!(order.id, "projecting");
///     HandlerOutcome::ack()
/// }
///
/// fn app() -> RustStream {
///     RustStream::new(AppInfo::new("projection", "0.1.0"))
///         .with_broker(KafkaBroker::new(["localhost:9092"]), |b| {
///             b.include(project);
///         })
/// }
/// # let _ = app;
/// ```
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum KafkaPosition {
    /// The earliest offset still retained, on every assigned partition.
    Earliest,
    /// The end of the log, on every assigned partition, as the broker reports it while the seek
    /// runs: records published after that point arrive normally, everything before is skipped.
    Latest,
    /// The first record at or after this timestamp (epoch milliseconds), resolved per assigned
    /// partition. A partition with no such record resumes at its end.
    Timestamp(i64),
    /// An absolute offset on one partition.
    Offset {
        /// The topic, when the position names one. [`Positioned`](ruststream::Positioned)
        /// captures the delivery's own topic here; [`offset`](Self::offset) leaves it unset,
        /// which repositions that partition index on every assigned topic (one topic being the
        /// usual case).
        topic: Option<String>,
        /// The partition to reposition.
        partition: i32,
        /// The offset to resume from: the record at this offset is delivered next.
        offset: i64,
    },
}

impl KafkaPosition {
    /// Every assigned partition, at the earliest retained offset.
    #[must_use]
    pub const fn earliest() -> Self {
        Self::Earliest
    }

    /// Every assigned partition, at the end of the log.
    #[must_use]
    pub const fn latest() -> Self {
        Self::Latest
    }

    /// One partition, at an absolute offset.
    #[must_use]
    pub const fn offset(partition: i32, offset: i64) -> Self {
        Self::Offset {
            topic: None,
            partition,
            offset,
        }
    }

    /// One partition of one topic, at an absolute offset. The form a delivery's
    /// [`position`](ruststream::Positioned::position) returns, and what a multi-topic
    /// subscription needs to name a partition unambiguously.
    #[must_use]
    pub fn topic_offset(topic: impl Into<String>, partition: i32, offset: i64) -> Self {
        Self::Offset {
            topic: Some(topic.into()),
            partition,
            offset,
        }
    }

    /// Every assigned partition, at the first record whose timestamp is at or after
    /// `when_millis` (epoch milliseconds).
    #[must_use]
    pub const fn timestamp(when_millis: i64) -> Self {
        Self::Timestamp(when_millis)
    }
}

/// What a [`KafkaSeeker`] repositions.
///
/// A live consumer moves its assigned partitions through librdkafka; under the `testing`
/// feature a member of the in-process cluster moves its partitions there, by the same rules.
/// Without the feature there is one variant and no branch.
#[derive(Clone)]
enum Repositioner {
    /// A librdkafka consumer: the reposition is a real seek over the partitions this member
    /// holds, plus the offset bookkeeping the new read position invalidates.
    Live {
        consumer: Arc<TrackedConsumer>,
        tracker: Arc<CommitTracker>,
    },
    /// A member of the in-process cluster.
    #[cfg(feature = "testing")]
    InProcess(Arc<Member>),
}

#[cfg(not(feature = "testing"))]
const _: () = assert!(size_of::<Repositioner>() == 2 * size_of::<usize>());

/// Repositions a live [`KafkaSubscriber`](crate::KafkaSubscriber), minted by
/// [`Seekable::seeker`](ruststream::Seekable::seeker).
///
/// Cheap to clone and usable while the subscription's stream runs, which is the point: the
/// runtime owns the subscriber, so a handler reaches its subscription through the `SeekHandle`
/// context key ([`Ctx(seeker): Ctx<SeekHandle>`](crate::context::keys::SeekHandle)) or through a
/// token minted at the mount site.
///
/// # Scope
///
/// A seek moves **this consumer instance**, over the partitions it currently holds. It is not a
/// group operation: it commits nothing, and other members of the group are unaffected.
///
/// # Rebalances discard a seek
///
/// The reposition lives in the assignment it was applied to. When a rebalance revokes those
/// partitions - a member joining or leaving, a session timeout, a topic-metadata change - the
/// seek goes with them: whoever gets the partitions next (this instance included) resumes from
/// the group's committed offsets. Repositioning is therefore an operational action on a running
/// consumer, not durable state; a position that must survive restarts belongs in the
/// subscription descriptor ([`StartOffset`](crate::StartOffset)) or in a `start_at(..)` clause,
/// which reapplies it on every startup.
#[derive(Clone)]
pub struct KafkaSeeker {
    inner: Repositioner,
    /// The connection the subscription was opened on. A seeker outlives the shutdown in a
    /// handler's hands, and a seek through it must then error rather than move a consumer the
    /// service no longer runs.
    connection: Arc<ConnState>,
    /// The subscription's name, for the error a seek after shutdown reports.
    subscription: Str,
}

impl std::fmt::Debug for KafkaSeeker {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let transport = match &self.inner {
            Repositioner::Live { .. } => "live",
            #[cfg(feature = "testing")]
            Repositioner::InProcess(_) => "in-process",
        };
        f.debug_struct("KafkaSeeker")
            .field("transport", &transport)
            .finish_non_exhaustive()
    }
}

impl KafkaSeeker {
    pub(crate) const fn new(
        consumer: Arc<TrackedConsumer>,
        tracker: Arc<CommitTracker>,
        connection: Arc<ConnState>,
        subscription: Str,
    ) -> Self {
        Self {
            inner: Repositioner::Live { consumer, tracker },
            connection,
            subscription,
        }
    }

    #[cfg(feature = "testing")]
    pub(crate) const fn in_process(
        member: Arc<Member>,
        connection: Arc<ConnState>,
        subscription: Str,
    ) -> Self {
        Self {
            inner: Repositioner::InProcess(member),
            connection,
            subscription,
        }
    }
}

impl Seeker for KafkaSeeker {
    type Position = KafkaPosition;
    type Error = KafkaError;

    /// Moves this consumer's assigned partitions to `to`; the next delivery comes from there.
    ///
    /// The offset bookkeeping moves with the read position: the tracked watermark of every
    /// repositioned partition is dropped, so a later commit cannot advance past a record the
    /// seek replayed but nobody handled, and an exactly-once window that was open when the seek
    /// landed aborts instead of committing offsets from the position it replaced.
    ///
    /// A consumer that holds no partitions yet is not an error: a group can open a subscription
    /// with none, when it has more members than partitions or the topic does not exist yet. The
    /// position is kept and applied to the first assignment the group hands over, ahead of its
    /// first record.
    ///
    /// # Errors
    ///
    /// Returns [`KafkaError::Closed`] once the broker the subscription was opened on has shut
    /// down, [`KafkaError::InvalidOptions`] when the position names no partition among the ones
    /// this consumer holds and [`KafkaError::Consume`] when librdkafka rejects the reposition or
    /// the timestamp lookup fails.
    ///
    /// # Cancel safety
    ///
    /// Not cancel safe: dropping the future may leave the consumer repositioned, its bookkeeping
    /// already reset.
    async fn seek(&self, to: Self::Position) -> Result<(), Self::Error> {
        self.connection.ensure_open(&self.subscription)?;
        match &self.inner {
            Repositioner::Live { consumer, tracker } => {
                let consumer = Arc::clone(consumer);
                let tracker = Arc::clone(tracker);
                // Every librdkafka call on this path blocks (the assignment wait, the timestamp
                // lookup, the seek itself), so the whole reposition runs on the blocking pool.
                task::spawn_blocking(move || reposition(&consumer, &tracker, &to))
                    .await
                    .map_err(|err| KafkaError::Consume(Box::new(err)))?
            }
            #[cfg(feature = "testing")]
            Repositioner::InProcess(member) => member.seek(&to),
        }
    }
}

/// Resolves `to` against the current assignment, resets the bookkeeping of every partition it
/// names, and moves the consumer.
///
/// A consumer with nothing assigned keeps the position instead: a group can open a subscription
/// with no partitions, and a `start_at(..)` on it still has to take effect. [`TrackingContext`]
/// applies it to the assignment when it lands.
fn reposition(
    consumer: &StreamConsumer<TrackingContext>,
    tracker: &CommitTracker,
    to: &KafkaPosition,
) -> Result<(), KafkaError> {
    let assignment = consumer.assignment().map_err(KafkaError::consume)?;
    if assignment.count() == 0 {
        consumer.context().hold_start(to.clone());
        return Ok(());
    }
    let targets = apply_position(consumer, tracker, &assignment, to)?;
    // A subscription whose stream has not started yet resumes from here when it does.
    consumer.context().startup.moved(&targets);
    Ok(())
}

/// Moves `assignment`'s partitions to `to`, resetting the offset bookkeeping of each one first:
/// the reposition of a running subscription.
fn apply_position<C>(
    consumer: &C,
    tracker: &CommitTracker,
    assignment: &TopicPartitionList,
    to: &KafkaPosition,
) -> Result<TopicPartitionList, KafkaError>
where
    C: Consumer<TrackingContext>,
{
    let targets = targets_in(consumer, assignment, to)?;

    // The bookkeeping is reset before the consumer moves: between the two, a delivery pulled
    // from the old position could otherwise settle into the new one and commit past records the
    // replay has not handled yet.
    for element in targets.elements() {
        tracker.reposition(&Str::from(element.topic()), element.partition());
    }
    clear_stored_offsets(consumer, &targets)?;

    let outcome = consumer
        .seek_partitions(targets.clone(), SEEK_TIMEOUT)
        .map_err(KafkaError::consume)?;
    for element in outcome.elements() {
        element.error().map_err(KafkaError::consume)?;
    }
    Ok(targets)
}

/// Writes `to` into an assignment the group has just handed over, before librdkafka takes it, so
/// fetching begins there; the offset bookkeeping of each partition it names is reset first.
///
/// This is how a start position reaches a group subscription. Seeking the new partitions instead
/// would race librdkafka: it refuses a seek (`Local: Erroneous state`) until it has fetched the
/// group's committed offset for the partition, and that fetch is still in flight when the
/// rebalance callback runs. An offset carried by the assignment itself needs no committed offset.
pub(crate) fn position_assignment<C>(
    consumer: &C,
    tracker: &CommitTracker,
    assignment: &mut TopicPartitionList,
    to: &KafkaPosition,
) -> Result<(), KafkaError>
where
    C: Consumer<TrackingContext>,
{
    let targets = targets_in(consumer, assignment, to)?;
    for element in targets.elements() {
        tracker.reposition(&Str::from(element.topic()), element.partition());
        assignment
            .set_partition_offset(element.topic(), element.partition(), element.offset())
            .map_err(KafkaError::consume)?;
    }
    Ok(())
}

/// Resolves `to` against `assignment`, refusing a position that names none of its partitions.
fn targets_in<C>(
    consumer: &C,
    assignment: &TopicPartitionList,
    to: &KafkaPosition,
) -> Result<TopicPartitionList, KafkaError>
where
    C: Consumer<TrackingContext>,
{
    let targets = resolve(consumer, assignment, to)?;
    if targets.count() == 0 {
        return Err(KafkaError::InvalidOptions(format!(
            "{to:?} names no partition assigned to this consumer; a seek moves the partitions \
             this instance holds, and its assignment is {}",
            describe(assignment),
        )));
    }
    Ok(targets)
}

/// Clears librdkafka's own stored offsets for the repositioned partitions.
///
/// The crate's watermark is not the only offset bookkeeping in play: librdkafka keeps a stored
/// position per partition and auto-commit flushes it, periodically and once more when the
/// consumer closes. That store still holds the position from before the seek, so leaving it
/// would let a commit carry the group past everything the seek replayed - the acks that built
/// it describe a read position this subscription no longer has. A logical `Invalid` offset is
/// how librdkafka itself clears the store when a partition is revoked.
pub(crate) fn clear_stored_offsets<C>(
    consumer: &C,
    targets: &TopicPartitionList,
) -> Result<(), KafkaError>
where
    C: Consumer<TrackingContext>,
{
    let mut cleared = TopicPartitionList::new();
    for element in targets.elements() {
        cleared
            .add_partition_offset(element.topic(), element.partition(), Offset::Invalid)
            .map_err(KafkaError::consume)?;
    }
    match consumer.store_offsets(&cleared) {
        Ok(()) => Ok(()),
        // `Commit::Auto` leaves librdkafka's own offset store enabled, which is exactly the
        // mode where it refuses application stores - and where it owns the position anyway, so
        // there is nothing of ours to clear.
        Err(err) if err.rdkafka_error_code() == Some(RDKafkaErrorCode::InvalidArgument) => Ok(()),
        Err(err) => Err(KafkaError::consume(err)),
    }
}

/// Builds the seek list: which assigned partitions move, and to which offsets.
fn resolve<C>(
    consumer: &C,
    assignment: &TopicPartitionList,
    to: &KafkaPosition,
) -> Result<TopicPartitionList, KafkaError>
where
    C: Consumer<TrackingContext>,
{
    let mut targets = TopicPartitionList::new();
    match to {
        // The log ends are resolved here rather than handed to librdkafka as logical offsets:
        // a logical end resolves whenever the fetcher next asks the broker, which would put
        // records published in the meantime on the wrong side of the reposition. Asking now
        // pins "the end of the log" to the moment the seek ran.
        KafkaPosition::Earliest | KafkaPosition::Latest => {
            for element in assignment.elements() {
                let (low, high) = consumer
                    .fetch_watermarks(element.topic(), element.partition(), SEEK_TIMEOUT)
                    .map_err(KafkaError::consume)?;
                let offset = if matches!(to, KafkaPosition::Earliest) {
                    low
                } else {
                    high
                };
                targets
                    .add_partition_offset(
                        element.topic(),
                        element.partition(),
                        Offset::Offset(offset),
                    )
                    .map_err(KafkaError::consume)?;
            }
        }
        KafkaPosition::Timestamp(when_millis) => {
            let mut query = TopicPartitionList::new();
            for element in assignment.elements() {
                query
                    .add_partition_offset(
                        element.topic(),
                        element.partition(),
                        Offset::Offset(*when_millis),
                    )
                    .map_err(KafkaError::consume)?;
            }
            let resolved = consumer
                .offsets_for_times(query, SEEK_TIMEOUT)
                .map_err(KafkaError::consume)?;
            for element in resolved.elements() {
                // A partition with no record at or after the timestamp answers with no usable
                // offset; the log end is the honest resume point (everything older is before
                // the requested time), pinned now for the reason the arm above pins it.
                let offset = match element.offset() {
                    Offset::Offset(offset) if offset >= 0 => offset,
                    _ => {
                        consumer
                            .fetch_watermarks(element.topic(), element.partition(), SEEK_TIMEOUT)
                            .map_err(KafkaError::consume)?
                            .1
                    }
                };
                let offset = Offset::Offset(offset);
                targets
                    .add_partition_offset(element.topic(), element.partition(), offset)
                    .map_err(KafkaError::consume)?;
            }
        }
        KafkaPosition::Offset {
            topic,
            partition,
            offset,
        } => {
            for element in assignment.elements() {
                let named = topic.as_ref().is_none_or(|name| name == element.topic());
                if named && element.partition() == *partition {
                    targets
                        .add_partition_offset(
                            element.topic(),
                            element.partition(),
                            Offset::Offset(*offset),
                        )
                        .map_err(KafkaError::consume)?;
                }
            }
        }
    }
    Ok(targets)
}

/// A human-readable assignment for the "nothing to seek" error and for the rebalance log line.
pub(crate) fn describe(assignment: &TopicPartitionList) -> String {
    if assignment.count() == 0 {
        return "empty".to_owned();
    }
    assignment
        .elements()
        .iter()
        .map(|element| format!("{}[{}]", element.topic(), element.partition()))
        .collect::<Vec<_>>()
        .join(", ")
}

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

    #[test]
    fn constructors_build_the_variants_the_seeker_matches_on() {
        assert_eq!(KafkaPosition::earliest(), KafkaPosition::Earliest);
        assert_eq!(KafkaPosition::latest(), KafkaPosition::Latest);
        assert_eq!(KafkaPosition::timestamp(42), KafkaPosition::Timestamp(42));
        assert_eq!(
            KafkaPosition::offset(3, 1_024),
            KafkaPosition::Offset {
                topic: None,
                partition: 3,
                offset: 1_024,
            },
        );
        assert_eq!(
            KafkaPosition::topic_offset("orders", 3, 1_024),
            KafkaPosition::Offset {
                topic: Some("orders".to_owned()),
                partition: 3,
                offset: 1_024,
            },
        );
    }
}