datum-mq 0.9.2

use std::{
    collections::{BTreeSet, VecDeque},
    sync::mpsc::{Receiver, RecvTimeoutError, SyncSender, TrySendError, sync_channel},
    sync::{
        Arc, Condvar, Mutex,
        atomic::{AtomicBool, Ordering},
    },
    thread,
    time::{Duration, Instant},
};

use bytes::Bytes;
use datum::{Source, SourceWithContext, StreamError};
use rdkafka::{
    ClientContext,
    consumer::{BaseConsumer, CommitMode, Consumer, ConsumerContext, Rebalance},
    error::KafkaError,
    message::Message,
    statistics::Statistics,
    topic_partition_list::{Offset, TopicPartitionList},
    util::Timeout,
};

use crate::offset::{OffsetCommitter, OffsetPosition, OffsetTracker};
use crate::profile::{self, ProfileBucket};
use crate::{
    CommitPolicy, ConsumerRecord, KafkaBatchOffset, KafkaConsumerSettings, KafkaMetrics,
    KafkaOffset, KafkaTimestamp, MqError, MqResult, Subscription, TopicPartition,
};

type PollMessage = Result<Vec<ConsumerRecord>, String>;
type PayloadPollMessage = Result<PayloadBatchRaw, String>;

/// Kafka source entry points.
pub struct KafkaSource;

#[derive(Debug, Clone)]
struct KafkaEnvelope {
    record: ConsumerRecord,
    offset: KafkaOffset,
}

/// Payload-only Kafka record emitted by [`KafkaSource::committable_payload_batches`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct KafkaPayloadRecord {
    pub partition: i32,
    pub offset: i64,
    pub timestamp: KafkaTimestamp,
    payload_start: usize,
    payload_end: usize,
}

impl KafkaPayloadRecord {
    #[must_use]
    pub fn payload<'a>(&self, batch: &'a KafkaPayloadBatch) -> &'a [u8] {
        &batch.payloads[self.payload_start..self.payload_end]
    }
}

/// A payload batch plus one committable Kafka watermark per touched partition.
#[derive(Debug, Clone)]
pub struct KafkaPayloadBatch {
    records: Vec<KafkaPayloadRecord>,
    payloads: Bytes,
    offset: KafkaBatchOffset,
}

impl KafkaPayloadBatch {
    #[must_use]
    pub fn records(&self) -> &[KafkaPayloadRecord] {
        &self.records
    }

    #[must_use]
    pub fn len(&self) -> usize {
        self.records.len()
    }

    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.records.is_empty()
    }

    #[must_use]
    pub fn payload(&self, record: &KafkaPayloadRecord) -> &[u8] {
        record.payload(self)
    }

    #[must_use]
    pub fn offset(&self) -> &KafkaBatchOffset {
        &self.offset
    }

    /// Commits this batch after the downstream batch checkpoint succeeds.
    pub fn commit(&self) -> MqResult<()> {
        self.offset.commit()
    }
}

impl KafkaSource {
    /// Creates a plain source of owned Kafka records.
    ///
    /// Use [`KafkaSource::committable`] when records should advance consumer-group
    /// offsets after downstream checkpoints.
    #[must_use]
    pub fn plain(
        settings: KafkaConsumerSettings,
        subscription: Subscription,
    ) -> Source<ConsumerRecord, KafkaControl> {
        envelope_source(settings, subscription).map(|envelope| envelope.record)
    }

    /// Creates a source whose records carry committable Kafka offsets as context.
    ///
    /// Call [`KafkaOffset::commit`](crate::KafkaOffset::commit) only after the
    /// corresponding downstream checkpoint or side effect succeeds.
    #[must_use]
    pub fn committable(
        settings: KafkaConsumerSettings,
        subscription: Subscription,
    ) -> SourceWithContext<ConsumerRecord, KafkaOffset, KafkaControl> {
        envelope_source(settings, subscription)
            .as_source_with_context(|envelope| envelope.offset.clone())
            .map(|envelope| envelope.record)
    }

    /// Creates a payload-only batched source with committable partition watermarks.
    ///
    /// This is the throughput-oriented counterpart to [`KafkaSource::committable`].
    /// It emits one Datum element per Kafka poll batch, copies payload bytes into a
    /// batch buffer, and exposes a single [`KafkaBatchOffset`] for the highest
    /// offset in each touched topic/partition. Downstream code should call
    /// [`KafkaPayloadBatch::commit`] only after the whole batch checkpoint succeeds.
    #[must_use]
    pub fn committable_payload_batches(
        settings: KafkaConsumerSettings,
        subscription: Subscription,
    ) -> Source<KafkaPayloadBatch, KafkaControl> {
        payload_batch_source(settings, subscription)
    }
}

fn envelope_source(
    settings: KafkaConsumerSettings,
    subscription: Subscription,
) -> Source<KafkaEnvelope, KafkaControl> {
    let metrics = KafkaMetrics::default();
    let control = KafkaControl::new(metrics);
    Source::unfold_resource(
        {
            let settings = settings.clone();
            let subscription = subscription.clone();
            let control = control.clone();
            move || {
                KafkaSourceResource::start(settings.clone(), subscription.clone(), control.clone())
                    .map_err(StreamError::from)
            }
        },
        KafkaSourceResource::read_next,
        KafkaSourceResource::close,
    )
    .map_materialized_value(move |_| control.clone())
}

fn payload_batch_source(
    settings: KafkaConsumerSettings,
    subscription: Subscription,
) -> Source<KafkaPayloadBatch, KafkaControl> {
    let metrics = KafkaMetrics::default();
    let control = KafkaControl::new(metrics);
    Source::unfold_resource(
        {
            let settings = settings.clone();
            let subscription = subscription.clone();
            let control = control.clone();
            move || {
                KafkaPayloadBatchSourceResource::start(
                    settings.clone(),
                    subscription.clone(),
                    control.clone(),
                )
                .map_err(StreamError::from)
            }
        },
        KafkaPayloadBatchSourceResource::read_next,
        KafkaPayloadBatchSourceResource::close,
    )
    .map_materialized_value(move |_| control.clone())
}

/// Materialized control for a Kafka source.
#[derive(Debug, Clone)]
pub struct KafkaControl {
    state: Arc<SourceControlState>,
}

#[derive(Debug)]
struct SourceControlState {
    draining: AtomicBool,
    shutdown_now: AtomicBool,
    initialized: AtomicBool,
    metrics: KafkaMetrics,
    wait_lock: Mutex<()>,
    wait_cv: Condvar,
}

impl KafkaControl {
    fn new(metrics: KafkaMetrics) -> Self {
        Self {
            state: Arc::new(SourceControlState {
                draining: AtomicBool::new(false),
                shutdown_now: AtomicBool::new(false),
                initialized: AtomicBool::new(false),
                metrics,
                wait_lock: Mutex::new(()),
                wait_cv: Condvar::new(),
            }),
        }
    }

    #[must_use]
    pub fn metrics(&self) -> KafkaMetrics {
        self.state.metrics.clone()
    }

    #[must_use]
    pub fn is_draining(&self) -> bool {
        self.state.draining.load(Ordering::SeqCst)
    }

    /// Stops new emission and waits for already emitted offsets to be committed.
    ///
    /// The source iterator completes on its next pull once outstanding offsets are
    /// drained. This is the Datum equivalent of the Alpakka `DrainingControl`
    /// hook; `datum-agent` jobs can call this before stopping the materialized
    /// graph.
    pub fn drain_and_shutdown(&self, timeout: Duration) -> MqResult<()> {
        self.state.draining.store(true, Ordering::SeqCst);
        self.state.wait_cv.notify_all();
        let deadline = Instant::now() + timeout;
        let mut guard = self
            .state
            .wait_lock
            .lock()
            .map_err(|_| MqError::Failed("Kafka source drain lock poisoned".to_owned()))?;
        loop {
            if self.metrics().snapshot().outstanding == 0 {
                return Ok(());
            }
            let now = Instant::now();
            if now >= deadline {
                return Err(MqError::DrainTimeout);
            }
            let remaining = deadline.saturating_duration_since(now);
            let (next_guard, result) = self
                .state
                .wait_cv
                .wait_timeout(guard, remaining)
                .map_err(|_| MqError::Failed("Kafka source drain wait poisoned".to_owned()))?;
            guard = next_guard;
            if result.timed_out() && self.metrics().snapshot().outstanding != 0 {
                return Err(MqError::DrainTimeout);
            }
        }
    }

    /// Stops the source without claiming uncommitted offsets.
    pub fn shutdown_now(&self) {
        self.state.shutdown_now.store(true, Ordering::SeqCst);
        self.state.wait_cv.notify_all();
    }

    fn mark_initialized(&self) {
        self.state.initialized.store(true, Ordering::SeqCst);
    }

    fn notify_progress(&self) {
        self.state.wait_cv.notify_all();
    }
}

struct KafkaSourceResource {
    consumer: Arc<BaseConsumer<DatumConsumerContext>>,
    committer: Arc<ConsumerCommitter>,
    control: KafkaControl,
    settings: KafkaConsumerSettings,
    receiver: Receiver<PollMessage>,
    buffer: VecDeque<ConsumerRecord>,
    poller: Option<thread::JoinHandle<()>>,
}

impl KafkaSourceResource {
    fn start(
        settings: KafkaConsumerSettings,
        subscription: Subscription,
        control: KafkaControl,
    ) -> MqResult<Self> {
        subscription.validate()?;
        control.mark_initialized();
        let shared = Arc::new(ConsumerSharedState {
            metrics: control.metrics(),
            lost: Mutex::new(BTreeSet::new()),
            tracker: Mutex::new(OffsetTracker::default()),
            last_flush: Mutex::new(Instant::now()),
            commit_policy: settings.commit_policy,
            commit_batch_size: settings.commit_batch_size,
            commit_interval: settings.commit_interval,
        });
        let context = DatumConsumerContext {
            shared: Arc::clone(&shared),
        };
        let consumer: BaseConsumer<DatumConsumerContext> =
            settings.to_client_config().create_with_context(context)?;
        subscribe(&consumer, &subscription)?;
        let consumer = Arc::new(consumer);
        let committer = Arc::new(ConsumerCommitter {
            consumer: Arc::clone(&consumer),
            shared,
            control: control.clone(),
            commit_mode: if settings.commit_sync {
                CommitMode::Sync
            } else {
                CommitMode::Async
            },
        });
        let paused = Arc::new(AtomicBool::new(false));
        let (sender, receiver) = sync_channel(1);
        let poller = Some(spawn_poll_thread(
            Arc::clone(&consumer),
            Arc::clone(&committer),
            control.clone(),
            settings.clone(),
            Arc::clone(&paused),
            sender,
        ));
        Ok(Self {
            consumer,
            committer,
            control,
            settings,
            receiver,
            buffer: VecDeque::new(),
            poller,
        })
    }

    fn read_next(&mut self) -> datum::StreamResult<Option<KafkaEnvelope>> {
        loop {
            if let Some(record) = self.buffer.pop_front() {
                return Ok(Some(profile::measure(ProfileBucket::SourceEmit, || {
                    self.envelope_for(record)
                })));
            }
            if self.control.state.shutdown_now.load(Ordering::SeqCst) {
                return Ok(None);
            }
            if self.control.state.draining.load(Ordering::SeqCst) {
                if self.committer.outstanding() == 0 {
                    return Ok(None);
                }
                thread::sleep(Duration::from_millis(1));
                continue;
            }
            match self.receiver.recv_timeout(self.settings.poll_timeout) {
                Ok(Ok(records)) => self.buffer.extend(records),
                Ok(Err(error)) => return Err(StreamError::Failed(error)),
                Err(RecvTimeoutError::Timeout) => continue,
                Err(RecvTimeoutError::Disconnected) => return Ok(None),
            }
        }
    }

    fn envelope_for(&self, record: ConsumerRecord) -> KafkaEnvelope {
        let position = OffsetPosition::new(record.topic.clone(), record.partition, record.offset);
        let outstanding = self.committer.observe(&position);
        self.control.metrics().emitted(outstanding, record.offset);
        let offset_committer: Arc<dyn OffsetCommitter> = self.committer.clone();
        let offset = KafkaOffset::new(position, offset_committer);
        KafkaEnvelope { record, offset }
    }

    fn close(self) -> datum::StreamResult<()> {
        let KafkaSourceResource {
            consumer,
            committer,
            control,
            settings,
            receiver,
            buffer: _,
            poller,
        } = self;
        if control.state.draining.load(Ordering::SeqCst) {
            let deadline = Instant::now() + settings.drain_timeout;
            while committer.outstanding() != 0 && Instant::now() < deadline {
                thread::sleep(Duration::from_millis(1));
            }
        }
        committer
            .flush_watermarks(true, CommitMode::Sync)
            .map_err(StreamError::from)?;
        control.state.shutdown_now.store(true, Ordering::SeqCst);
        drop(receiver);
        if let Some(poller) = poller {
            poller
                .join()
                .map_err(|_| StreamError::Failed("Kafka source poll thread panicked".to_owned()))?;
        }
        consumer.unsubscribe();
        control.notify_progress();
        Ok(())
    }
}

#[derive(Debug)]
struct PayloadBatchRaw {
    records: Vec<KafkaPayloadRecord>,
    payloads: Vec<u8>,
    watermarks: Vec<PayloadWatermark>,
}

impl PayloadBatchRaw {
    fn high_watermark(&self) -> i64 {
        self.watermarks
            .iter()
            .map(|watermark| watermark.last_offset)
            .max()
            .unwrap_or(-1)
    }

    fn observe_positions(&self) -> Vec<OffsetPosition> {
        let mut positions = Vec::with_capacity(self.watermarks.len() * 2);
        for watermark in &self.watermarks {
            positions.push(OffsetPosition::new(
                watermark.topic.clone(),
                watermark.partition,
                watermark.first_offset,
            ));
            if watermark.last_offset != watermark.first_offset {
                positions.push(OffsetPosition::new(
                    watermark.topic.clone(),
                    watermark.partition,
                    watermark.last_offset,
                ));
            }
        }
        positions
    }

    fn commit_positions(&self) -> Vec<OffsetPosition> {
        self.watermarks
            .iter()
            .map(|watermark| {
                OffsetPosition::new(
                    watermark.topic.clone(),
                    watermark.partition,
                    watermark.last_offset,
                )
            })
            .collect()
    }
}

#[derive(Debug)]
struct PayloadWatermark {
    topic: String,
    partition: i32,
    first_offset: i64,
    last_offset: i64,
}

#[derive(Debug)]
struct PayloadBatchBuilder {
    records: Vec<KafkaPayloadRecord>,
    payloads: Vec<u8>,
    watermarks: Vec<PayloadWatermark>,
}

impl PayloadBatchBuilder {
    fn with_capacity(records: usize, payload_bytes: usize) -> Self {
        Self {
            records: Vec::with_capacity(records),
            payloads: Vec::with_capacity(payload_bytes),
            watermarks: Vec::new(),
        }
    }

    fn from_raw(raw: PayloadBatchRaw) -> Self {
        Self {
            records: raw.records,
            payloads: raw.payloads,
            watermarks: raw.watermarks,
        }
    }

    fn is_empty(&self) -> bool {
        self.records.is_empty()
    }

    fn len(&self) -> usize {
        self.records.len()
    }

    fn push(&mut self, message: &rdkafka::message::BorrowedMessage<'_>) {
        let payload_start = self.payloads.len();
        if let Some(payload) = message.payload() {
            self.payloads.extend_from_slice(payload);
        }
        let payload_end = self.payloads.len();
        self.records.push(KafkaPayloadRecord {
            partition: message.partition(),
            offset: message.offset(),
            timestamp: message.timestamp().into(),
            payload_start,
            payload_end,
        });
        self.observe(message.topic(), message.partition(), message.offset());
    }

    fn take(&mut self) -> Option<PayloadBatchRaw> {
        if self.is_empty() {
            return None;
        }
        let mut outgoing =
            Self::with_capacity(self.records.capacity(), self.payloads.capacity().max(1));
        std::mem::swap(self, &mut outgoing);
        Some(PayloadBatchRaw {
            records: outgoing.records,
            payloads: outgoing.payloads,
            watermarks: outgoing.watermarks,
        })
    }

    fn observe(&mut self, topic: &str, partition: i32, offset: i64) {
        if let Some(watermark) = self
            .watermarks
            .iter_mut()
            .find(|watermark| watermark.topic == topic && watermark.partition == partition)
        {
            watermark.first_offset = watermark.first_offset.min(offset);
            watermark.last_offset = watermark.last_offset.max(offset);
        } else {
            self.watermarks.push(PayloadWatermark {
                topic: topic.to_owned(),
                partition,
                first_offset: offset,
                last_offset: offset,
            });
        }
    }
}

struct KafkaPayloadBatchSourceResource {
    consumer: Arc<BaseConsumer<DatumConsumerContext>>,
    committer: Arc<ConsumerCommitter>,
    control: KafkaControl,
    settings: KafkaConsumerSettings,
    receiver: Receiver<PayloadPollMessage>,
    poller: Option<thread::JoinHandle<()>>,
}

impl KafkaPayloadBatchSourceResource {
    fn start(
        settings: KafkaConsumerSettings,
        subscription: Subscription,
        control: KafkaControl,
    ) -> MqResult<Self> {
        subscription.validate()?;
        control.mark_initialized();
        let shared = Arc::new(ConsumerSharedState {
            metrics: control.metrics(),
            lost: Mutex::new(BTreeSet::new()),
            tracker: Mutex::new(OffsetTracker::default()),
            last_flush: Mutex::new(Instant::now()),
            commit_policy: settings.commit_policy,
            commit_batch_size: settings.commit_batch_size,
            commit_interval: settings.commit_interval,
        });
        let context = DatumConsumerContext {
            shared: Arc::clone(&shared),
        };
        let consumer: BaseConsumer<DatumConsumerContext> =
            settings.to_client_config().create_with_context(context)?;
        subscribe(&consumer, &subscription)?;
        let consumer = Arc::new(consumer);
        let committer = Arc::new(ConsumerCommitter {
            consumer: Arc::clone(&consumer),
            shared,
            control: control.clone(),
            commit_mode: if settings.commit_sync {
                CommitMode::Sync
            } else {
                CommitMode::Async
            },
        });
        let paused = Arc::new(AtomicBool::new(false));
        let (sender, receiver) = sync_channel(1);
        let poller = Some(spawn_payload_poll_thread(
            Arc::clone(&consumer),
            Arc::clone(&committer),
            control.clone(),
            settings.clone(),
            Arc::clone(&paused),
            sender,
        ));
        Ok(Self {
            consumer,
            committer,
            control,
            settings,
            receiver,
            poller,
        })
    }

    fn read_next(&mut self) -> datum::StreamResult<Option<KafkaPayloadBatch>> {
        loop {
            if self.control.state.shutdown_now.load(Ordering::SeqCst) {
                return Ok(None);
            }
            if self.control.state.draining.load(Ordering::SeqCst) {
                if self.committer.outstanding() == 0 {
                    return Ok(None);
                }
                thread::sleep(Duration::from_millis(1));
                continue;
            }
            match self.receiver.recv_timeout(self.settings.poll_timeout) {
                Ok(Ok(raw)) => {
                    return Ok(Some(profile::measure(ProfileBucket::SourceEmit, || {
                        self.batch_for(raw)
                    })));
                }
                Ok(Err(error)) => return Err(StreamError::Failed(error)),
                Err(RecvTimeoutError::Timeout) => continue,
                Err(RecvTimeoutError::Disconnected) => return Ok(None),
            }
        }
    }

    fn batch_for(&self, raw: PayloadBatchRaw) -> KafkaPayloadBatch {
        let observe_positions = raw.observe_positions();
        let outstanding = self.committer.observe_positions(&observe_positions);
        self.control.metrics().emitted_batch(
            raw.records.len() as u64,
            outstanding,
            raw.high_watermark(),
        );
        let offset_committer: Arc<dyn OffsetCommitter> = self.committer.clone();
        let offset = KafkaBatchOffset::new(raw.commit_positions(), offset_committer);
        KafkaPayloadBatch {
            records: raw.records,
            payloads: Bytes::from(raw.payloads),
            offset,
        }
    }

    fn close(self) -> datum::StreamResult<()> {
        let KafkaPayloadBatchSourceResource {
            consumer,
            committer,
            control,
            settings,
            receiver,
            poller,
        } = self;
        if control.state.draining.load(Ordering::SeqCst) {
            let deadline = Instant::now() + settings.drain_timeout;
            while committer.outstanding() != 0 && Instant::now() < deadline {
                thread::sleep(Duration::from_millis(1));
            }
        }
        committer
            .flush_watermarks(true, CommitMode::Sync)
            .map_err(StreamError::from)?;
        control.state.shutdown_now.store(true, Ordering::SeqCst);
        drop(receiver);
        if let Some(poller) = poller {
            poller
                .join()
                .map_err(|_| StreamError::Failed("Kafka source poll thread panicked".to_owned()))?;
        }
        consumer.unsubscribe();
        control.notify_progress();
        Ok(())
    }
}

fn spawn_poll_thread(
    consumer: Arc<BaseConsumer<DatumConsumerContext>>,
    committer: Arc<ConsumerCommitter>,
    control: KafkaControl,
    settings: KafkaConsumerSettings,
    paused: Arc<AtomicBool>,
    sender: SyncSender<PollMessage>,
) -> thread::JoinHandle<()> {
    thread::Builder::new()
        .name("datum-mq-consumer-poll".to_owned())
        .spawn(move || poll_loop(consumer, committer, control, settings, paused, sender))
        .expect("Kafka source poll thread must start")
}

fn spawn_payload_poll_thread(
    consumer: Arc<BaseConsumer<DatumConsumerContext>>,
    committer: Arc<ConsumerCommitter>,
    control: KafkaControl,
    settings: KafkaConsumerSettings,
    paused: Arc<AtomicBool>,
    sender: SyncSender<PayloadPollMessage>,
) -> thread::JoinHandle<()> {
    thread::Builder::new()
        .name("datum-mq-consumer-poll".to_owned())
        .spawn(move || payload_poll_loop(consumer, committer, control, settings, paused, sender))
        .expect("Kafka payload source poll thread must start")
}

fn poll_loop(
    consumer: Arc<BaseConsumer<DatumConsumerContext>>,
    committer: Arc<ConsumerCommitter>,
    control: KafkaControl,
    settings: KafkaConsumerSettings,
    paused: Arc<AtomicBool>,
    sender: SyncSender<PollMessage>,
) {
    let mut batch = Vec::with_capacity(settings.poll_batch_size);
    loop {
        if control.state.shutdown_now.load(Ordering::SeqCst) {
            break;
        }
        if control.state.draining.load(Ordering::SeqCst) && committer.outstanding() == 0 {
            break;
        }
        if let Err(error) = update_pause_resume(&consumer, &committer, &control, &settings, &paused)
        {
            let _ = sender.send(Err(error.to_string()));
            break;
        }
        let timeout = if batch.is_empty() {
            settings.poll_timeout
        } else {
            Duration::ZERO
        };
        match profile::measure(ProfileBucket::PollFetch, || {
            consumer.poll(Timeout::After(timeout))
        }) {
            Some(Ok(message)) => {
                let record = profile::measure(ProfileBucket::RecordConversion, || {
                    ConsumerRecord::from_borrowed(&message)
                });
                if settings.low_latency && batch.is_empty() {
                    match profile::measure(ProfileBucket::ChannelHandoff, || {
                        sender.try_send(Ok(vec![record]))
                    }) {
                        Ok(()) => continue,
                        Err(TrySendError::Full(Ok(mut records))) => {
                            batch.append(&mut records);
                        }
                        Err(TrySendError::Disconnected(_)) => break,
                        Err(TrySendError::Full(Err(error))) => {
                            let _ = sender.send(Err(error));
                            break;
                        }
                    }
                } else {
                    batch.push(record);
                }
                if batch.len() >= settings.poll_batch_size
                    && send_consumer_batch(&sender, &mut batch).is_err()
                {
                    break;
                }
            }
            Some(Err(KafkaError::PartitionEOF(_))) | None => {
                if !batch.is_empty() && send_consumer_batch(&sender, &mut batch).is_err() {
                    break;
                }
            }
            Some(Err(error)) => {
                let _ = send_consumer_batch(&sender, &mut batch);
                let _ = sender.send(Err(MqError::Kafka(error).to_string()));
                break;
            }
        }
    }
}

fn payload_poll_loop(
    consumer: Arc<BaseConsumer<DatumConsumerContext>>,
    committer: Arc<ConsumerCommitter>,
    control: KafkaControl,
    settings: KafkaConsumerSettings,
    paused: Arc<AtomicBool>,
    sender: SyncSender<PayloadPollMessage>,
) {
    let mut batch = PayloadBatchBuilder::with_capacity(
        settings.poll_batch_size,
        settings.poll_batch_size.saturating_mul(256),
    );
    loop {
        if control.state.shutdown_now.load(Ordering::SeqCst) {
            break;
        }
        if control.state.draining.load(Ordering::SeqCst) && committer.outstanding() == 0 {
            break;
        }
        if let Err(error) = update_pause_resume(&consumer, &committer, &control, &settings, &paused)
        {
            let _ = sender.send(Err(error.to_string()));
            break;
        }
        let timeout = if batch.is_empty() {
            settings.poll_timeout
        } else {
            Duration::ZERO
        };
        match profile::measure(ProfileBucket::PollFetch, || {
            consumer.poll(Timeout::After(timeout))
        }) {
            Some(Ok(message)) => {
                profile::measure(ProfileBucket::RecordConversion, || batch.push(&message));
                if settings.low_latency && batch.len() == 1 {
                    let raw = batch.take().expect("payload batch has one record");
                    match profile::measure(ProfileBucket::ChannelHandoff, || {
                        sender.try_send(Ok(raw))
                    }) {
                        Ok(()) => continue,
                        Err(TrySendError::Full(Ok(raw))) => {
                            batch = PayloadBatchBuilder::from_raw(raw);
                        }
                        Err(TrySendError::Disconnected(_)) => break,
                        Err(TrySendError::Full(Err(error))) => {
                            let _ = sender.send(Err(error));
                            break;
                        }
                    }
                }
                if batch.len() >= settings.poll_batch_size
                    && send_payload_batch(&sender, &mut batch).is_err()
                {
                    break;
                }
            }
            Some(Err(KafkaError::PartitionEOF(_))) | None => {
                if !batch.is_empty() && send_payload_batch(&sender, &mut batch).is_err() {
                    break;
                }
            }
            Some(Err(error)) => {
                let _ = send_payload_batch(&sender, &mut batch);
                let _ = sender.send(Err(MqError::Kafka(error).to_string()));
                break;
            }
        }
    }
}

fn send_consumer_batch(
    sender: &SyncSender<PollMessage>,
    batch: &mut Vec<ConsumerRecord>,
) -> Result<(), ()> {
    if batch.is_empty() {
        return Ok(());
    }
    let mut outgoing = Vec::with_capacity(batch.capacity());
    std::mem::swap(batch, &mut outgoing);
    profile::measure(ProfileBucket::ChannelHandoff, || sender.send(Ok(outgoing))).map_err(|_| ())
}

fn send_payload_batch(
    sender: &SyncSender<PayloadPollMessage>,
    batch: &mut PayloadBatchBuilder,
) -> Result<(), ()> {
    let Some(outgoing) = batch.take() else {
        return Ok(());
    };
    profile::measure(ProfileBucket::ChannelHandoff, || sender.send(Ok(outgoing))).map_err(|_| ())
}

fn update_pause_resume(
    consumer: &BaseConsumer<DatumConsumerContext>,
    committer: &ConsumerCommitter,
    control: &KafkaControl,
    settings: &KafkaConsumerSettings,
    paused: &AtomicBool,
) -> MqResult<()> {
    let outstanding = committer.outstanding() as usize;
    if !paused.load(Ordering::SeqCst) && outstanding >= settings.high_watermark {
        if let Ok(assignment) = consumer.assignment()
            && assignment.count() > 0
        {
            consumer.pause(&assignment)?;
            paused.store(true, Ordering::SeqCst);
            control.metrics().set_paused(true);
        }
    } else if paused.load(Ordering::SeqCst)
        && outstanding <= settings.low_watermark
        && let Ok(assignment) = consumer.assignment()
        && assignment.count() > 0
    {
        consumer.resume(&assignment)?;
        paused.store(false, Ordering::SeqCst);
        control.metrics().set_paused(false);
    }
    Ok(())
}

fn subscribe(
    consumer: &BaseConsumer<DatumConsumerContext>,
    subscription: &Subscription,
) -> MqResult<()> {
    match subscription {
        Subscription::Topics(topics) => {
            let refs = topics.iter().map(String::as_str).collect::<Vec<_>>();
            consumer.subscribe(&refs)?;
        }
        Subscription::Pattern(pattern) => {
            consumer.subscribe(&[pattern.as_str()])?;
        }
        Subscription::Assignment(partitions) => {
            let mut tpl = TopicPartitionList::with_capacity(partitions.len());
            for partition in partitions {
                tpl.add_partition_offset(&partition.topic, partition.partition, partition.offset)?;
            }
            consumer.assign(&tpl)?;
        }
    }
    Ok(())
}

#[derive(Debug)]
struct ConsumerSharedState {
    metrics: KafkaMetrics,
    lost: Mutex<BTreeSet<TopicPartition>>,
    tracker: Mutex<OffsetTracker>,
    last_flush: Mutex<Instant>,
    commit_policy: CommitPolicy,
    commit_batch_size: usize,
    commit_interval: Duration,
}

#[derive(Clone)]
struct DatumConsumerContext {
    shared: Arc<ConsumerSharedState>,
}

impl ClientContext for DatumConsumerContext {
    fn stats(&self, _statistics: Statistics) {}
}

impl ConsumerContext for DatumConsumerContext {
    fn pre_rebalance(&self, base_consumer: &BaseConsumer<Self>, rebalance: &Rebalance<'_>) {
        match rebalance {
            Rebalance::Assign(tpl) => {
                self.shared.metrics.add_assigned(tpl.count() as u64);
            }
            Rebalance::Revoke(tpl) => {
                let partitions = topic_partitions(tpl);
                if base_consumer.assignment_lost() {
                    if let Ok(mut lost) = self.shared.lost.lock() {
                        lost.extend(partitions.iter().cloned());
                    }
                    self.shared.metrics.add_lost(partitions.len() as u64);
                } else {
                    if self
                        .shared
                        .flush_watermarks(base_consumer, true, CommitMode::Sync)
                        .is_err()
                    {
                        self.shared.metrics.commit_failed();
                    }
                    self.shared.metrics.add_revoked(partitions.len() as u64);
                }
            }
            Rebalance::Error(_) => {}
        }
    }

    fn post_rebalance(&self, _base_consumer: &BaseConsumer<Self>, rebalance: &Rebalance<'_>) {
        if let Rebalance::Assign(tpl) = rebalance
            && let Ok(mut lost) = self.shared.lost.lock()
        {
            for partition in topic_partitions(tpl) {
                lost.remove(&partition);
            }
        }
    }

    fn commit_callback(
        &self,
        result: rdkafka::error::KafkaResult<()>,
        _offsets: &TopicPartitionList,
    ) {
        if result.is_err() {
            self.shared.metrics.commit_failed();
        }
    }

    fn main_queue_min_poll_interval(&self) -> Timeout {
        Timeout::After(Duration::from_millis(100))
    }
}

fn topic_partitions(tpl: &TopicPartitionList) -> Vec<TopicPartition> {
    tpl.elements()
        .into_iter()
        .map(|element| TopicPartition::new(element.topic(), element.partition()))
        .collect()
}

impl ConsumerSharedState {
    fn observe(&self, position: &OffsetPosition) -> MqResult<u64> {
        let mut tracker = self
            .tracker
            .lock()
            .map_err(|_| MqError::Failed("Kafka offset tracker poisoned".to_owned()))?;
        Ok(tracker.observe(position))
    }

    fn observe_many(&self, positions: &[OffsetPosition]) -> MqResult<u64> {
        let mut tracker = self
            .tracker
            .lock()
            .map_err(|_| MqError::Failed("Kafka offset tracker poisoned".to_owned()))?;
        for position in positions {
            tracker.observe(position);
        }
        Ok(tracker.outstanding())
    }

    fn process(&self, position: &OffsetPosition) -> MqResult<(u64, u64)> {
        let mut tracker = self
            .tracker
            .lock()
            .map_err(|_| MqError::Failed("Kafka offset tracker poisoned".to_owned()))?;
        Ok(tracker.process(position))
    }

    fn outstanding(&self) -> u64 {
        self.tracker
            .lock()
            .expect("Kafka offset tracker poisoned")
            .outstanding()
    }

    fn flush_watermarks(
        &self,
        consumer: &BaseConsumer<DatumConsumerContext>,
        force: bool,
        mode: CommitMode,
    ) -> MqResult<()> {
        let due_by_count = {
            let tracker = self
                .tracker
                .lock()
                .map_err(|_| MqError::Failed("Kafka offset tracker poisoned".to_owned()))?;
            tracker.uncommitted() >= self.commit_batch_size as u64
        };
        let due_by_time = self
            .last_flush
            .lock()
            .map_err(|_| MqError::Failed("Kafka commit flush clock poisoned".to_owned()))?
            .elapsed()
            >= self.commit_interval;
        if !force && !due_by_count && !due_by_time {
            return Ok(());
        }

        let commits = {
            let tracker = self
                .tracker
                .lock()
                .map_err(|_| MqError::Failed("Kafka offset tracker poisoned".to_owned()))?;
            tracker.due_commits()
        };
        if commits.is_empty() {
            if let Ok(mut last_flush) = self.last_flush.lock() {
                *last_flush = Instant::now();
            }
            return Ok(());
        }

        if self.commit_policy != CommitPolicy::External {
            let mut tpl = TopicPartitionList::with_capacity(commits.len());
            for commit in &commits {
                tpl.add_partition_offset(
                    &commit.topic,
                    commit.partition,
                    Offset::Offset(commit.offset),
                )?;
            }
            consumer.commit(&tpl, mode)?;
        }

        let (outstanding, committed_watermark) = {
            let mut tracker = self
                .tracker
                .lock()
                .map_err(|_| MqError::Failed("Kafka offset tracker poisoned".to_owned()))?;
            tracker.mark_committed(&commits)
        };
        if let Ok(mut last_flush) = self.last_flush.lock() {
            *last_flush = Instant::now();
        }
        self.metrics.committed(outstanding, committed_watermark);
        Ok(())
    }
}

struct ConsumerCommitter {
    consumer: Arc<BaseConsumer<DatumConsumerContext>>,
    shared: Arc<ConsumerSharedState>,
    control: KafkaControl,
    commit_mode: CommitMode,
}

impl ConsumerCommitter {
    fn observe(&self, position: &OffsetPosition) -> u64 {
        self.shared.observe(position).unwrap_or_else(|_| {
            self.control.metrics().commit_failed();
            0
        })
    }

    fn observe_positions(&self, positions: &[OffsetPosition]) -> u64 {
        self.shared.observe_many(positions).unwrap_or_else(|_| {
            self.control.metrics().commit_failed();
            0
        })
    }

    fn outstanding(&self) -> u64 {
        self.shared.outstanding()
    }

    fn flush_watermarks(&self, force: bool, mode: CommitMode) -> MqResult<()> {
        self.shared.flush_watermarks(&self.consumer, force, mode)
    }
}

impl OffsetCommitter for ConsumerCommitter {
    fn stage_offset(&self, position: OffsetPosition) -> MqResult<()> {
        profile::measure(ProfileBucket::CommitBookkeeping, || {
            let (outstanding, _) = self.shared.process(&position)?;
            self.control.metrics().set_outstanding(outstanding);
            self.control.notify_progress();
            Ok(())
        })
    }

    fn commit_offset(&self, position: OffsetPosition) -> MqResult<()> {
        profile::measure(ProfileBucket::CommitBookkeeping, || {
            let topic_partition = position.topic_partition();
            if self
                .shared
                .lost
                .lock()
                .map_err(|_| MqError::Failed("Kafka lost-assignment lock poisoned".to_owned()))?
                .contains(&topic_partition)
                || self.consumer.assignment_lost()
            {
                return Err(MqError::AssignmentLost {
                    topic: topic_partition.topic,
                    partition: topic_partition.partition,
                });
            }

            let (outstanding, _) = self.shared.process(&position)?;
            self.control.metrics().set_outstanding(outstanding);
            self.flush_watermarks(false, self.commit_mode)
                .inspect_err(|_| self.control.metrics().commit_failed())?;
            self.control.notify_progress();
            Ok(())
        })
    }
}