kafkit_client/consumer/

mod.rs

//! Group consumer API.
//!
//! ```no_run
//! # async fn example() -> kafkit_client::Result<()> {
//! use kafkit_client::KafkaClient;
//!
//! let consumer = KafkaClient::new("localhost:9092")
//!     .consumer("orders-reader")
//!     .with_topic("orders")
//!     .connect()
//!     .await?;
//!
//! let records = consumer.poll().await?;
//! consumer.commit(&records).await?;
//! consumer.shutdown().await?;
//! # Ok(())
//! # }
//! ```
//!
mod protocol;
mod runtime;
mod scheduler;
mod share;
mod state;

use std::collections::BTreeSet;
use std::time::Duration;

use tokio::sync::{mpsc, oneshot};
use tokio::task::JoinHandle;
use tracing::{debug, instrument};

use crate::config::ConsumerConfig;
use crate::types::{
    CommitOffset, ConsumerGroupMetadata, ConsumerRecords, SubscriptionPattern, TopicPartition,
    TopicPartitionInfo, TopicPartitionOffset, TopicPartitionOffsetAndTimestamp,
    TopicPartitionTimestamp,
};
use crate::{ConsumerError, Result};

use runtime::ConsumerRuntime;

/// A Kafka group consumer with explicit polling and commits.
pub struct KafkaConsumer {
    application_event_handler: ConsumerApplicationEventHandler,
    join: JoinHandle<()>,
    default_poll_timeout: Duration,
}

pub use share::{
    AcknowledgeType, AcknowledgementCommitCallback, KafkaShareConsumer, ShareAcknowledgementCommit,
    ShareAcquireMode, ShareConsumerOptions, ShareRecord, ShareRecords,
};

impl KafkaConsumer {
    #[instrument(
        name = "consumer.connect",
        level = "debug",
        skip(config),
        fields(
            bootstrap_server_count = config.bootstrap_servers.len(),
            client_id = %config.client_id,
            group_id = %config.group_id
        )
    )]
    /// Connects to Kafka and returns the client.
    pub async fn connect(config: ConsumerConfig) -> Result<Self> {
        let (tx, rx) = mpsc::channel(64);
        let runtime = ConsumerRuntime::new(config);
        let join = tokio::spawn(async move {
            runtime.run(rx).await;
        });

        let consumer = Self {
            application_event_handler: ConsumerApplicationEventHandler::new(tx),
            join,
            default_poll_timeout: Duration::from_millis(100),
        };
        if let Err(error) = consumer.warm_up().await {
            consumer.join.abort();
            return Err(error);
        }
        debug!("consumer connected");
        Ok(consumer)
    }

    /// Sets default poll timeout and returns the updated value.
    pub fn with_default_poll_timeout(mut self, timeout: Duration) -> Self {
        self.default_poll_timeout = timeout;
        self
    }

    #[instrument(
        name = "consumer.subscribe",
        level = "debug",
        skip(self, topics),
        fields(topic_count = topics.len())
    )]
    /// Subscribe.
    pub async fn subscribe(&self, topics: Vec<String>) -> Result<()> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::Subscribe {
                    topics,
                    reply: reply_tx,
                },
                ConsumerError::ThreadStoppedBefore {
                    operation: "subscribe",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "subscribe",
            })?
    }

    #[instrument(
        name = "consumer.subscribe_pattern",
        level = "debug",
        skip(self, pattern),
        fields(pattern = %pattern.pattern())
    )]
    /// Subscribe Pattern.
    pub async fn subscribe_pattern(&self, pattern: SubscriptionPattern) -> Result<()> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::SubscribePattern {
                    pattern,
                    reply: reply_tx,
                },
                ConsumerError::ThreadStoppedBefore {
                    operation: "subscribe_pattern",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "subscribe_pattern",
            })?
    }

    #[instrument(
        name = "consumer.subscribe_regex",
        level = "debug",
        skip(self, pattern),
        fields(pattern = %pattern)
    )]
    /// Subscribe Regex.
    pub async fn subscribe_regex(&self, pattern: String) -> Result<()> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::SubscribeRegex {
                    pattern,
                    reply: reply_tx,
                },
                ConsumerError::ThreadStoppedBefore {
                    operation: "subscribe_regex",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "subscribe_regex",
            })?
    }

    #[instrument(name = "consumer.unsubscribe", level = "debug", skip(self))]
    /// Unsubscribe.
    pub async fn unsubscribe(&self) -> Result<()> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::Unsubscribe { reply: reply_tx },
                ConsumerError::ThreadStoppedBefore {
                    operation: "unsubscribe",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "unsubscribe",
            })?
    }

    /// Polls Kafka for records.
    pub async fn poll(&self) -> Result<ConsumerRecords> {
        self.poll_for(self.default_poll_timeout).await
    }

    #[instrument(
        name = "consumer.poll",
        level = "debug",
        skip(self),
        fields(timeout_ms = timeout.as_millis())
    )]
    /// Poll For.
    pub async fn poll_for(&self, timeout: Duration) -> Result<ConsumerRecords> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::Poll {
                    timeout,
                    reply: reply_tx,
                },
                ConsumerError::ThreadStoppedBefore { operation: "poll" },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring { operation: "poll" })?
    }

    #[instrument(
        name = "consumer.assign",
        level = "debug",
        skip(self, partitions),
        fields(partition_count = partitions.len())
    )]
    /// Assign.
    pub async fn assign(&self, partitions: Vec<TopicPartition>) -> Result<()> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::Assign {
                    partitions,
                    reply: reply_tx,
                },
                ConsumerError::ThreadStoppedBefore {
                    operation: "assign",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "assign",
            })?
    }

    #[instrument(
        name = "consumer.seek",
        level = "debug",
        skip(self, partition),
        fields(topic = %partition.topic, partition_id = partition.partition, offset)
    )]
    /// Seek.
    pub async fn seek(&self, partition: TopicPartition, offset: i64) -> Result<()> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::Seek {
                    partition,
                    offset,
                    reply: reply_tx,
                },
                ConsumerError::ThreadStoppedBefore { operation: "seek" },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring { operation: "seek" })?
    }

    #[instrument(
        name = "consumer.seek_to_beginning",
        level = "debug",
        skip(self, partitions),
        fields(partition_count = partitions.len())
    )]
    /// Seek To Beginning.
    pub async fn seek_to_beginning(&self, partitions: Vec<TopicPartition>) -> Result<()> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::SeekToBeginning {
                    partitions,
                    reply: reply_tx,
                },
                ConsumerError::ThreadStoppedBefore {
                    operation: "seek_to_beginning",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "seek_to_beginning",
            })?
    }

    #[instrument(
        name = "consumer.seek_to_end",
        level = "debug",
        skip(self, partitions),
        fields(partition_count = partitions.len())
    )]
    /// Seek To End.
    pub async fn seek_to_end(&self, partitions: Vec<TopicPartition>) -> Result<()> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::SeekToEnd {
                    partitions,
                    reply: reply_tx,
                },
                ConsumerError::ThreadStoppedBefore {
                    operation: "seek_to_end",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "seek_to_end",
            })?
    }

    #[instrument(
        name = "consumer.seek_to_timestamp",
        level = "debug",
        skip(self, partitions),
        fields(partition_count = partitions.len())
    )]
    /// Seek To Timestamp.
    pub async fn seek_to_timestamp(&self, partitions: Vec<TopicPartitionTimestamp>) -> Result<()> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::SeekToTimestamp {
                    partitions,
                    reply: reply_tx,
                },
                ConsumerError::ThreadStoppedBefore {
                    operation: "seek_to_timestamp",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "seek_to_timestamp",
            })?
    }

    #[instrument(
        name = "consumer.position",
        level = "debug",
        skip(self, partition),
        fields(topic = %partition.topic, partition_id = partition.partition)
    )]
    /// Position.
    pub async fn position(&self, partition: TopicPartition) -> Result<i64> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::Position {
                    partition,
                    reply: reply_tx,
                },
                ConsumerError::ThreadStoppedBefore {
                    operation: "position",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "position",
            })?
    }

    #[instrument(
        name = "consumer.pause",
        level = "debug",
        skip(self, partitions),
        fields(partition_count = partitions.len())
    )]
    /// Pause.
    pub async fn pause(&self, partitions: Vec<TopicPartition>) -> Result<()> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::Pause {
                    partitions,
                    reply: reply_tx,
                },
                ConsumerError::ThreadStoppedBefore { operation: "pause" },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring { operation: "pause" })?
    }

    #[instrument(
        name = "consumer.resume",
        level = "debug",
        skip(self, partitions),
        fields(partition_count = partitions.len())
    )]
    /// Resume.
    pub async fn resume(&self, partitions: Vec<TopicPartition>) -> Result<()> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::Resume {
                    partitions,
                    reply: reply_tx,
                },
                ConsumerError::ThreadStoppedBefore {
                    operation: "resume",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "resume",
            })?
    }

    #[instrument(name = "consumer.group_metadata", level = "debug", skip(self))]
    /// Group Metadata.
    pub async fn group_metadata(&self) -> Result<ConsumerGroupMetadata> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::GroupMetadata { reply: reply_tx },
                ConsumerError::ThreadStoppedBefore {
                    operation: "group_metadata",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "group_metadata",
            })?
    }

    #[instrument(name = "consumer.assignment", level = "debug", skip(self))]
    /// Assignment.
    pub async fn assignment(&self) -> Result<BTreeSet<TopicPartition>> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::Assignment { reply: reply_tx },
                ConsumerError::ThreadStoppedBefore {
                    operation: "assignment",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "assignment",
            })?
    }

    #[instrument(
        name = "consumer.committed_offsets",
        level = "debug",
        skip(self, partitions),
        fields(partition_count = partitions.len())
    )]
    /// Committed.
    pub async fn committed(
        &self,
        partitions: Vec<TopicPartition>,
    ) -> Result<Vec<TopicPartitionOffset>> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::Committed {
                    partitions,
                    reply: reply_tx,
                },
                ConsumerError::ThreadStoppedBefore {
                    operation: "committed",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "committed",
            })?
    }

    #[instrument(
        name = "consumer.beginning_offsets",
        level = "debug",
        skip(self, partitions),
        fields(partition_count = partitions.len())
    )]
    /// Beginning Offsets.
    pub async fn beginning_offsets(
        &self,
        partitions: Vec<TopicPartition>,
    ) -> Result<Vec<TopicPartitionOffset>> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::BeginningOffsets {
                    partitions,
                    reply: reply_tx,
                },
                ConsumerError::ThreadStoppedBefore {
                    operation: "beginning_offsets",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "beginning_offsets",
            })?
    }

    #[instrument(
        name = "consumer.end_offsets",
        level = "debug",
        skip(self, partitions),
        fields(partition_count = partitions.len())
    )]
    /// End Offsets.
    pub async fn end_offsets(
        &self,
        partitions: Vec<TopicPartition>,
    ) -> Result<Vec<TopicPartitionOffset>> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::EndOffsets {
                    partitions,
                    reply: reply_tx,
                },
                ConsumerError::ThreadStoppedBefore {
                    operation: "end_offsets",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "end_offsets",
            })?
    }

    #[instrument(
        name = "consumer.offsets_for_times",
        level = "debug",
        skip(self, partitions),
        fields(partition_count = partitions.len())
    )]
    /// Offsets For Times.
    pub async fn offsets_for_times(
        &self,
        partitions: Vec<TopicPartitionTimestamp>,
    ) -> Result<Vec<TopicPartitionOffsetAndTimestamp>> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::OffsetsForTimes {
                    partitions,
                    reply: reply_tx,
                },
                ConsumerError::ThreadStoppedBefore {
                    operation: "offsets_for_times",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "offsets_for_times",
            })?
    }

    #[instrument(
        name = "consumer.partitions_for",
        level = "debug",
        skip(self),
        fields(topic = %topic)
    )]
    /// Partitions For.
    pub async fn partitions_for(&self, topic: String) -> Result<Vec<TopicPartitionInfo>> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::PartitionsFor {
                    topic,
                    reply: reply_tx,
                },
                ConsumerError::ThreadStoppedBefore {
                    operation: "partitions_for",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "partitions_for",
            })?
    }

    #[instrument(name = "consumer.list_topics", level = "debug", skip(self))]
    /// List Topics.
    pub async fn list_topics(&self) -> Result<Vec<String>> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::ListTopics { reply: reply_tx },
                ConsumerError::ThreadStoppedBefore {
                    operation: "list_topics",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "list_topics",
            })?
    }

    /// Commit.
    pub async fn commit(&self, records: &ConsumerRecords) -> Result<()> {
        self.commit_offsets(records.commit_offsets()).await
    }

    #[instrument(
        name = "consumer.commit_offsets",
        level = "debug",
        skip(self, offsets),
        fields(offset_count = offsets.len())
    )]
    /// Commit Offsets.
    pub async fn commit_offsets(&self, offsets: Vec<CommitOffset>) -> Result<()> {
        if offsets.is_empty() {
            return Ok(());
        }

        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::Commit {
                    offsets,
                    reply: reply_tx,
                },
                ConsumerError::ThreadStoppedBefore {
                    operation: "commit",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "commit",
            })?
    }

    #[instrument(name = "consumer.wakeup", level = "debug", skip(self))]
    /// Wakeup.
    pub async fn wakeup(&self) -> Result<()> {
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::Wakeup,
                ConsumerError::ThreadStoppedBefore {
                    operation: "wakeup",
                },
            )
            .await
    }

    #[instrument(name = "consumer.shutdown", level = "debug", skip(self))]
    /// Shuts the client down and waits for in-flight work to finish.
    pub async fn shutdown(self) -> Result<()> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::Shutdown { reply: reply_tx },
                ConsumerError::ThreadStoppedBefore {
                    operation: "shutdown",
                },
            )
            .await?;

        let result = reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "shutdown",
            })?;
        self.join.await.map_err(ConsumerError::Join)?;
        result
    }

    #[instrument(name = "consumer.warm_up", level = "trace", skip(self))]
    async fn warm_up(&self) -> Result<()> {
        let (reply_tx, reply_rx) = oneshot::channel();
        self.application_event_handler
            .send(
                ConsumerRuntimeEvent::WarmUp { reply: reply_tx },
                ConsumerError::ThreadStoppedBefore {
                    operation: "startup",
                },
            )
            .await?;
        reply_rx
            .await
            .map_err(|_| ConsumerError::ThreadStoppedDuring {
                operation: "startup",
            })?
    }
}

struct ConsumerApplicationEventHandler {
    tx: mpsc::Sender<ConsumerRuntimeEvent>,
}

impl ConsumerApplicationEventHandler {
    fn new(tx: mpsc::Sender<ConsumerRuntimeEvent>) -> Self {
        Self { tx }
    }

    async fn send(&self, event: ConsumerRuntimeEvent, stopped_error: ConsumerError) -> Result<()> {
        self.tx.send(event).await.map_err(|_| stopped_error.into())
    }
}

enum ConsumerRuntimeEvent {
    WarmUp {
        reply: oneshot::Sender<Result<()>>,
    },
    Subscribe {
        topics: Vec<String>,
        reply: oneshot::Sender<Result<()>>,
    },
    SubscribePattern {
        pattern: SubscriptionPattern,
        reply: oneshot::Sender<Result<()>>,
    },
    SubscribeRegex {
        pattern: String,
        reply: oneshot::Sender<Result<()>>,
    },
    Unsubscribe {
        reply: oneshot::Sender<Result<()>>,
    },
    Assign {
        partitions: Vec<TopicPartition>,
        reply: oneshot::Sender<Result<()>>,
    },
    Poll {
        timeout: Duration,
        reply: oneshot::Sender<Result<ConsumerRecords>>,
    },
    Seek {
        partition: TopicPartition,
        offset: i64,
        reply: oneshot::Sender<Result<()>>,
    },
    SeekToBeginning {
        partitions: Vec<TopicPartition>,
        reply: oneshot::Sender<Result<()>>,
    },
    SeekToEnd {
        partitions: Vec<TopicPartition>,
        reply: oneshot::Sender<Result<()>>,
    },
    SeekToTimestamp {
        partitions: Vec<TopicPartitionTimestamp>,
        reply: oneshot::Sender<Result<()>>,
    },
    Position {
        partition: TopicPartition,
        reply: oneshot::Sender<Result<i64>>,
    },
    Pause {
        partitions: Vec<TopicPartition>,
        reply: oneshot::Sender<Result<()>>,
    },
    Resume {
        partitions: Vec<TopicPartition>,
        reply: oneshot::Sender<Result<()>>,
    },
    GroupMetadata {
        reply: oneshot::Sender<Result<ConsumerGroupMetadata>>,
    },
    Assignment {
        reply: oneshot::Sender<Result<BTreeSet<TopicPartition>>>,
    },
    Committed {
        partitions: Vec<TopicPartition>,
        reply: oneshot::Sender<Result<Vec<TopicPartitionOffset>>>,
    },
    BeginningOffsets {
        partitions: Vec<TopicPartition>,
        reply: oneshot::Sender<Result<Vec<TopicPartitionOffset>>>,
    },
    EndOffsets {
        partitions: Vec<TopicPartition>,
        reply: oneshot::Sender<Result<Vec<TopicPartitionOffset>>>,
    },
    OffsetsForTimes {
        partitions: Vec<TopicPartitionTimestamp>,
        reply: oneshot::Sender<Result<Vec<TopicPartitionOffsetAndTimestamp>>>,
    },
    PartitionsFor {
        topic: String,
        reply: oneshot::Sender<Result<Vec<TopicPartitionInfo>>>,
    },
    ListTopics {
        reply: oneshot::Sender<Result<Vec<String>>>,
    },
    Commit {
        offsets: Vec<CommitOffset>,
        reply: oneshot::Sender<Result<()>>,
    },
    Wakeup,
    Shutdown {
        reply: oneshot::Sender<Result<()>>,
    },
}